[VectorCombine][TTI] Prevent extract/ins rewrite to GEP

llvm/include/llvm/Analysis/TargetTransformInfo.h

@@ -1950,6 +1950,10 @@ class TargetTransformInfo {
       const Function &F,
       SmallVectorImpl<std::pair<StringRef, int64_t>> &LB) const;
 
+  /// Returns true if GEP should not be used to index into vectors for this
+  /// target.
+  LLVM_ABI bool allowVectorElementIndexingUsingGEP() const;
+
 private:
   std::unique_ptr<const TargetTransformInfoImplBase> TTIImpl;
 };

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

@@ -1145,6 +1145,8 @@ class TargetTransformInfoImplBase {
       const Function &F,
       SmallVectorImpl<std::pair<StringRef, int64_t>> &LB) const {}
 
+  virtual bool allowVectorElementIndexingUsingGEP() const { return true; }
+
 protected:
   // Obtain the minimum required size to hold the value (without the sign)
   // In case of a vector it returns the min required size for one element.

llvm/lib/Analysis/TargetTransformInfo.cpp

@@ -1486,6 +1486,10 @@ void TargetTransformInfo::collectKernelLaunchBounds(
   return TTIImpl->collectKernelLaunchBounds(F, LB);
 }
 
+bool TargetTransformInfo::allowVectorElementIndexingUsingGEP() const {
+  return TTIImpl->allowVectorElementIndexingUsingGEP();
+}
+
 TargetTransformInfoImplBase::~TargetTransformInfoImplBase() = default;
 
 TargetIRAnalysis::TargetIRAnalysis() : TTICallback(&getDefaultTTI) {}

llvm/lib/Target/SPIRV/SPIRVTargetTransformInfo.h

@@ -59,6 +59,8 @@ class SPIRVTTIImpl final : public BasicTTIImplBase<SPIRVTTIImpl> {
                                   Intrinsic::ID IID) const override;
   Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II, Value *OldV,
                                           Value *NewV) const override;
+
+  bool allowVectorElementIndexingUsingGEP() const override { return false; }
 };
 
 } // namespace llvm

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

@@ -1664,6 +1664,8 @@ static Align computeAlignmentAfterScalarization(Align VectorAlignment,
 //   %1 = getelementptr inbounds i32, i32* %0, i64 0, i64 1
 //   store i32 %b, i32* %1
 bool VectorCombine::foldSingleElementStore(Instruction &I) {
+  if (!TTI.allowVectorElementIndexingUsingGEP())
+    return false;
   auto *SI = cast<StoreInst>(&I);
   if (!SI->isSimple() || !isa<VectorType>(SI->getValueOperand()->getType()))
     return false;
@@ -1719,6 +1721,9 @@ bool VectorCombine::foldSingleElementStore(Instruction &I) {
 
 /// Try to scalarize vector loads feeding extractelement instructions.
 bool VectorCombine::scalarizeLoadExtract(Instruction &I) {
+  if (!TTI.allowVectorElementIndexingUsingGEP())
+    return false;
+
   Value *Ptr;
   if (!match(&I, m_Load(m_Value(Ptr))))
     return false;
@@ -1827,6 +1832,8 @@ bool VectorCombine::scalarizeLoadExtract(Instruction &I) {
 }
 
 bool VectorCombine::scalarizeExtExtract(Instruction &I) {
+  if (!TTI.allowVectorElementIndexingUsingGEP())
+    return false;
   auto *Ext = dyn_cast<ZExtInst>(&I);
   if (!Ext)
     return false;

llvm/test/Transforms/VectorCombine/SPIRV/lit.local.cfg

@@ -0,0 +1,2 @@
+if not "SPIRV" in config.root.targets:
+    config.unsupported = True