[llvm] [AMDGPU] Support i8/i16 GEP indices when promoting allocas to vectors (PR #175489)

[Carl Ritson via llvm-commits]

================
@@ -551,33 +556,61 @@ computeGEPToVectorIndex(GetElementPtrInst *GEP, AllocaInst *Alloca,
   if (VarOffsets.size() > 1)
     return {};
 
-  APInt IndexQuot;
-  int64_t Rem;
-  APInt::sdivrem(ConstOffset, VecElemSize, IndexQuot, Rem);
-  if (Rem != 0)
+  // We support vector indices of the form ((A * stride) >> shift) + B.
+  // IndexQuot represents B. Check that the constant offset is a multiple
+  // of the vector element size.
+  if (ConstOffset.srem(VecElemSize) != 0)
     return {};
+  APInt IndexQuot = ConstOffset.sdiv(VecElemSize);
 
   GEPToVectorIndex Result;
 
   if (!ConstOffset.isZero())
     Result.ConstIndex = ConstantInt::get(Ctx, IndexQuot.sextOrTrunc(BW));
 
+  // If there are no variable offsets, only a constant offset, then we're done.
   if (VarOffsets.empty())
     return Result;
 
+  // Scale is the stride in the (A * stride) part. Check that there is only one
+  // variable offset and extract the scale factor.
   const auto &VarOffset = VarOffsets.front();
-  APInt OffsetQuot;
-  APInt::sdivrem(VarOffset.second, VecElemSize, OffsetQuot, Rem);
-  if (Rem != 0 || OffsetQuot.isZero())
+  uint64_t Scale = VarOffset.second.getZExtValue();
+  if (Scale == 0)
     return {};
 
   Result.VarIndex = VarOffset.first;
   auto *OffsetType = dyn_cast<IntegerType>(Result.VarIndex->getType());
   if (!OffsetType)
     return {};
 
-  if (!OffsetQuot.isOne())
-    Result.VarMul = ConstantInt::get(Ctx, OffsetQuot.sextOrTrunc(BW));
+  // The vector index for the variable part is: A * Scale / VecElemSize.
+  if (Scale >= (uint64_t)VecElemSize) {
+    // Scale is a multiple of VecElemSize, so the index is just A * (Scale /
+    // VecElemSize). Only the multiplier is needed.
+    if (Scale % VecElemSize != 0)
+      return {};
+
+    uint64_t VarMul = Scale / VecElemSize;
+    if (VarMul != 1)
+      Result.VarMul = ConstantInt::get(Ctx, APInt(BW, VarMul));
+  } else {
+    // VecElemSize is a multiple of Scale, so the index is A / (VecElemSize /
----------------
perlfu wrote:

Could you perhaps break this comment down next to the pieces of code which implement the various conditions?
Specifically at this point `VecElemSize` is not yet proven to be a multiple of scale.
And on a first read it was not immediately clear how the `computeKnownBits` was used to prove that `A` can be cleanly divided/shifted by `Divisor`.

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