[llvm] [BasicAA] Preserve constant offset correlations through select indices (PR #228778)

[Dávid Bolvanský via llvm-commits]

https://github.com/davidbolvansky created https://github.com/llvm/llvm-project/pull/228778

BasicAA currently returns `MayAlias` for accesses indexed by `select(c, x + 1, x + 2)` and `x + 3`, despite their disjoint offsets.

Check constant differences between matching linear expressions in each select alternative, with a budget of eight comparisons. This resolves the existing `select_in_gep2` TODO to `NoAlias`; compile-time tracker results are neutral.


>From 9fcc3f729efe6dbe6183bf24d2d6c9c1c3d9eb48 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?D=C3=A1vid=20Bolvansk=C3=BD?= <david.bolvansky at gmail.com>
Date: Sat, 3 Oct 2026 23:06:24 +0200
Subject: [PATCH] [BasicAA] Check constant offsets in select index alternatives

---
 llvm/lib/Analysis/BasicAliasAnalysis.cpp      | 32 +++++++++++++++++++
 .../test/Analysis/BasicAA/non-equal-select.ll |  3 +-
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Analysis/BasicAliasAnalysis.cpp b/llvm/lib/Analysis/BasicAliasAnalysis.cpp
index 3aece3ca7d769..d6eaa2aa4dda7 100644
--- a/llvm/lib/Analysis/BasicAliasAnalysis.cpp
+++ b/llvm/lib/Analysis/BasicAliasAnalysis.cpp
@@ -2107,6 +2107,38 @@ bool BasicAAResult::computeConstantOffsetHeuristic(const DecomposedGEP &GEP,
       Var0.Val.V->getType() != Var1.Val.V->getType())
     return false;
 
+  // Without extensions, matching linear expressions have an exact modular
+  // difference. Check each select alternative before losing that correlation.
+  if (!AAQI.MayBeCrossIteration && !Var0.Val.ZExtBits && !Var0.Val.SExtBits &&
+      Var0.Val.getBitWidth() == GEP.Offset.getBitWidth() &&
+      (isa<SelectInst>(Var0.Val.V) || isa<SelectInst>(Var1.Val.V))) {
+    unsigned Budget = 8;
+    auto Check = [&](auto &&Check, const Value *V0, const Value *V1) -> bool {
+      if (!Budget)
+        return false;
+      --Budget;
+      LinearExpression E0 = GetLinearExpression(CastedValue(V0), DL, 0, AC, DT);
+      LinearExpression E1 = GetLinearExpression(CastedValue(V1), DL, 0, AC, DT);
+      if (E0.Scale == E1.Scale && E0.Val.hasSameCastsAs(E1.Val) &&
+          isValueEqualInPotentialCycles(E0.Val.V, E1.Val.V, AAQI)) {
+        APInt Offset = (E0.Offset - E1.Offset) * Var0.Scale;
+        if (Var0.IsNegated)
+          Offset = -Offset;
+        Offset += GEP.Offset;
+        return Offset.uge(V2Size) && (-Offset).uge(V1Size);
+      }
+      if (const auto *SI = dyn_cast<SelectInst>(V0))
+        return Check(Check, SI->getTrueValue(), V1) &&
+               Check(Check, SI->getFalseValue(), V1);
+      if (const auto *SI = dyn_cast<SelectInst>(V1))
+        return Check(Check, V0, SI->getTrueValue()) &&
+               Check(Check, V0, SI->getFalseValue());
+      return false;
+    };
+    if (Check(Check, Var0.Val.V, Var1.Val.V))
+      return true;
+  }
+
   // We'll strip off the Extensions of Var0 and Var1 and do another round
   // of GetLinearExpression decomposition. In the example above, if Var0
   // is zext(%x + 1) we should get V1 == %x and V1Offset == 1.
diff --git a/llvm/test/Analysis/BasicAA/non-equal-select.ll b/llvm/test/Analysis/BasicAA/non-equal-select.ll
index fe38a57eeb1b4..76986fbf4e0e1 100644
--- a/llvm/test/Analysis/BasicAA/non-equal-select.ll
+++ b/llvm/test/Analysis/BasicAA/non-equal-select.ll
@@ -17,9 +17,8 @@ entry:
 
 define void @select_in_gep2(i1 %c, i64 %x) {
 entry:
-  ; TODO: should be "NoAlias" here as well.
 ; CHECK-LABEL: Function: select_in_gep2
-; CHECK: MayAlias:     i32* %arrayidx1, i32* %arrayidx2
+; CHECK: NoAlias:      i32* %arrayidx1, i32* %arrayidx2
   %add1_ = add nsw i64 %x, 1
   %add2_ = add nsw i64 %x, 2
   %add3_ = add nsw i64 %x, 3