[llvm] [InferAlignment] Add variable offset alignment analysis for GEPs in InferAlignment (PR #185980)

[via llvm-commits]

================
@@ -133,32 +183,134 @@ bool inferAlignment(Function &F, AssumptionCache &AC, DominatorTree &DT) {
     return Align(1ull << std::min(Known.getBitWidth() - 1, TrailZ));
   };
 
+  // Helper function to compute variable offset alignment and base pointer.
+  // If ConstOffset > 0, the effective offset alignment is limited by the
+  // constant offset.
+  auto computeVariableOffsetAlignment =
+      [&](Value *Ptr, uint64_t ConstOffset = 0) -> std::pair<Value *, Align> {
+    Align VarOffsetAlign = Align(1);
+    Value *VarBasePtr = Ptr;
+
+    if (auto *GEP = dyn_cast<GEPOperator>(VarBasePtr)) {
+      // We can only handle GEPs with a single index
+      if (GEP->getNumIndices() == 1) {
+        Value *Idx = GEP->idx_begin()->get();
+        Align IndexAlign(1);
+
+        if (auto PatternAlign = getIndexAlignmentFromPattern(Idx))
+          IndexAlign = *PatternAlign;
+
+        Type *EltTy = GEP->getSourceElementType();
+        TypeSize EltSizeType = DL.getTypeAllocSize(EltTy);
+
+        // If we encounter a scalable type, we can't compute alignment.
+        if (!EltSizeType.isScalable()) {
+          uint64_t EltSize = EltSizeType.getFixedValue();
+
+          // Compute offset alignment: multiply index alignment by element size,
+          // then take the greatest power of 2 that divides the product.
+          uint64_t Product = IndexAlign.value() * EltSize;
+          uint64_t ProductAlignValue =
+              Product > 0 ? (Product & (~Product + 1)) : 1;
+          VarOffsetAlign = Align(ProductAlignValue);
+        }
+
+        VarBasePtr = GEP->getPointerOperand();
+      }
+    }
+    VarBasePtr = VarBasePtr->stripPointerCasts();
+
+    // If we have a constant offset, the effective alignment is the GCD of both.
+    if (ConstOffset > 0)
+      VarOffsetAlign = commonAlignment(VarOffsetAlign, ConstOffset);
+
+    return {VarBasePtr, VarOffsetAlign};
+  };
+
   // Propagate alignment between loads and stores that originate from the
   // same base pointer.
   ScopedHT BestBasePointerAligns;
+
+  auto updateBestBaseAlign = [&](Value *BasePtr, Align NewAlign) {
+    if (Align Old = BestBasePointerAligns.lookup(BasePtr); Old != Align()) {
+      if (NewAlign <= Old)
+        return;
+    }
+    BestBasePointerAligns.insert(BasePtr, NewAlign);
----------------
yasmincs wrote:

I have tested this with patterns such as: 
entry:
  br i1 %cond, label %then, label %else

then:
  call void @llvm.assume(i1 true) ["align"(ptr %a, i32 64)]
  br label %merge

else:
  %load = load i32, ptr %a, align 4
  br label %merge

merge:
  ret void

The code works as intended. After the for loop going through all the instructions in a BasicBlock, the loop pushes the next child node on the stack or removes the AlignmentScope if all of the children have been visited.

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