[llvm] Add MUL/SHL handling in decomposeLinearExpression (PR #222574)

[Arne Stenkrona via llvm-commits]

https://github.com/ArneStenkrona2 updated https://github.com/llvm/llvm-project/pull/222574

>From d20b308d4db6e9e99c72f587adb1df5cfcd7309b Mon Sep 17 00:00:00 2001
From: Arne Stenkrona <arne.stenkrona at arm.com>
Date: Mon, 7 Sep 2026 17:03:41 +0200
Subject: [PATCH] Add MUL/SHL handling in decomposeLinearExpression

Implements a FIXME in decomposeLinearExpression, allowing it to look
through multiplications and left-shifts.

We peel the variable index and keep track of the cumulated scale from
the multiplications and shifts and apply this scale when we return the
final expression.

This allows decomposeLinearExpression to extract more statically known
information into the `Scale` of the linear expression.

assisted-by: codex
---
 llvm/lib/Analysis/Loads.cpp                  | 34 ++++++++++++++++++--
 llvm/test/Transforms/InstCombine/load-cmp.ll | 14 ++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Analysis/Loads.cpp b/llvm/lib/Analysis/Loads.cpp
index de9022c540d42..125c23bf5fda7 100644
--- a/llvm/lib/Analysis/Loads.cpp
+++ b/llvm/lib/Analysis/Loads.cpp
@@ -24,8 +24,10 @@
 #include "llvm/IR/GetElementPtrTypeIterator.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Operator.h"
+#include "llvm/IR/PatternMatch.h"
 
 using namespace llvm;
+using namespace llvm::PatternMatch;
 
 static bool isAligned(const Value *Base, Align Alignment,
                       const DataLayout &DL) {
@@ -939,6 +941,32 @@ LinearExpression llvm::decomposeLinearExpression(const DataLayout &DL,
     if (VarIndex && !VarIndex->getType()->isIntegerTy(BitWidth))
       return Expr;
 
+    APInt IndexScale(BitWidth, 1);
+    while (auto *BO = dyn_cast_or_null<BinaryOperator>(VarIndex)) {
+      Value *UnscaledIndex = nullptr;
+      APInt Factor(BitWidth, 0);
+      const APInt *Constant;
+      if (match(BO, m_Mul(m_Value(UnscaledIndex), m_APInt(Constant)))) {
+        Factor = *Constant;
+      } else if (match(BO, m_Shl(m_Value(UnscaledIndex),
+                                 m_APInt(Constant)))) {
+        if (Constant->uge(BitWidth))
+          break;
+        Factor.setBit(Constant->getZExtValue());
+      } else {
+        break;
+      }
+
+      // If we overflow, we cannot fold the factor into the current scale
+      bool Overflow = false;
+      APInt NewIndexScale = IndexScale.umul_ov(Factor, Overflow);
+      if (Overflow)
+        break;
+
+      IndexScale = NewIndexScale;
+      VarIndex = UnscaledIndex;
+    }
+
     // We have verified that we can fully handle this GEP, so we can update Expr
     // members past this point.
     Expr.BasePtr = GEP->getPointerOperand();
@@ -963,12 +991,12 @@ LinearExpression llvm::decomposeLinearExpression(const DataLayout &DL,
         continue;
       }
 
-      // FIXME: Also look through a mul/shl in the index.
       assert(Expr.Index == nullptr && "Shouldn't have index yet");
-      Expr.Index = Index;
+      Expr.Index = VarIndex;
       // Truncate if type size exceeds index space.
       Expr.Scale = APInt(BitWidth, GTI.getSequentialElementStride(DL),
-                         /*isSigned=*/false, /*implicitTrunc=*/true);
+                         /*isSigned=*/false, /*implicitTrunc=*/true) *
+                   IndexScale;
     }
   }
 
diff --git a/llvm/test/Transforms/InstCombine/load-cmp.ll b/llvm/test/Transforms/InstCombine/load-cmp.ll
index e5b90107bb27d..8bdd88cea6959 100644
--- a/llvm/test/Transforms/InstCombine/load-cmp.ll
+++ b/llvm/test/Transforms/InstCombine/load-cmp.ll
@@ -608,3 +608,17 @@ define i1 @cmp_load_multiple_indices2(i32 %idx, i32 %idx2) {
   %cmp = icmp eq i16 %load, 0
   ret i1 %cmp
 }
+
+define i1 @cmp_load_multiplied_and_shifted_index(i32 %X) {
+; CHECK-LABEL: @cmp_load_multiplied_and_shifted_index(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X:%.*]], 1
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %mul = mul nuw i32 %X, 2
+  %shl = shl nuw i32 %mul, 1
+  %mul2 = mul nuw i32 %shl, 2
+  %gep = getelementptr inbounds [10 x i16], ptr @G16, i32 0, i32 %mul2
+  %load = load i16, ptr %gep
+  %cmp = icmp eq i16 %load, 68
+  ret i1 %cmp
+}