[llvm] [SPIR-V] Fix some GEP legalization (PR #150943)

[Nathan Gauër via llvm-commits]

================
@@ -593,54 +604,64 @@ void SPIRVEmitIntrinsics::maybeAssignPtrType(Type *&Ty, Value *Op, Type *RefTy,
 
 bool SPIRVEmitIntrinsics::walkLogicalAccessChain(
     GetElementPtrInst &GEP,
-    const std::function<void(Type *, uint64_t)> &OnStaticIndexing,
+    const std::function<void(Type *, uint64_t)> &OnLiteralIndexing,
     const std::function<void(Type *, Value *)> &OnDynamicIndexing) {
+  // We only rewrite i8* GEP. Other should be left as-is.
+  // Observation so-far is i8* GEP always have a single index. Making sure
+  // that's the case.
+  assert(GEP.getSourceElementType() ==
+         IntegerType::getInt8Ty(CurrF->getContext()));
+  assert(GEP.getNumIndices() == 1);
+
   auto &DL = CurrF->getDataLayout();
   Value *Src = getPointerRoot(GEP.getPointerOperand());
   Type *CurType = deduceElementType(Src, true);
 
-  for (Value *V : GEP.indices()) {
-    if (ConstantInt *CI = dyn_cast<ConstantInt>(V)) {
-      uint64_t Offset = CI->getZExtValue();
-
-      do {
-        if (ArrayType *AT = dyn_cast<ArrayType>(CurType)) {
-          uint32_t EltTypeSize = DL.getTypeSizeInBits(AT->getElementType()) / 8;
-          assert(Offset < AT->getNumElements() * EltTypeSize);
-          uint64_t Index = Offset / EltTypeSize;
-          Offset = Offset - (Index * EltTypeSize);
-          CurType = AT->getElementType();
-          OnStaticIndexing(CurType, Index);
-        } else if (StructType *ST = dyn_cast<StructType>(CurType)) {
-          uint32_t StructSize = DL.getTypeSizeInBits(ST) / 8;
-          assert(Offset < StructSize);
-          const auto &STL = DL.getStructLayout(ST);
-          unsigned Element = STL->getElementContainingOffset(Offset);
-          Offset -= STL->getElementOffset(Element);
-          CurType = ST->getElementType(Element);
-          OnStaticIndexing(CurType, Element);
-        } else {
-          // Vector type indexing should not use GEP.
-          // So if we have an index left, something is wrong. Giving up.
-          return true;
-        }
-      } while (Offset > 0);
-
-    } else if (ArrayType *AT = dyn_cast<ArrayType>(CurType)) {
-      // Index is not constant. Either we have an array and accept it, or we
-      // give up.
+  Value *Operand = *GEP.idx_begin();
+  ConstantInt *CI = dyn_cast<ConstantInt>(Operand);
+  if (!CI) {
+    ArrayType *AT = dyn_cast<ArrayType>(CurType);
+    // Operand is not constant. Either we have an array and accept it, or we
+    // give up.
+    if (AT)
+      OnDynamicIndexing(AT->getElementType(), Operand);
+    return AT == nullptr;
+  }
+
+  assert(CI);
----------------
Keenuts wrote:

This assert is more a "comment" for the reader to remind of the precondition since this is an implicit else.

https://github.com/llvm/llvm-project/pull/150943