[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 e01b3644daae1a2e4d8046add4d1947ac14b6f67 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                  | 32 ++++++++++++++++++--
 llvm/test/Transforms/InstCombine/load-cmp.ll | 24 +++++++++++++++
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Analysis/Loads.cpp b/llvm/lib/Analysis/Loads.cpp
index de9022c540d42..dd19e38e87a1f 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) {
@@ -943,6 +945,30 @@ LinearExpression llvm::decomposeLinearExpression(const DataLayout &DL,
     // members past this point.
     Expr.BasePtr = GEP->getPointerOperand();
     Expr.Flags = Expr.Flags.intersectForOffsetAdd(GEP->getNoWrapFlags());
+
+    // Try to look through mul and shl
+    APInt IndexScale(BitWidth, 1);
+    if (auto *BO = dyn_cast_or_null<BinaryOperator>(VarIndex)) {
+      Value *UnscaledIndex = nullptr;
+      const APInt *Constant;
+      if (match(BO, m_Mul(m_Value(UnscaledIndex), m_APInt(Constant)))) {
+        IndexScale *= *Constant;
+      } else if (match(BO, m_Shl(m_Value(UnscaledIndex), m_APInt(Constant))) &&
+                 Constant->ult(BitWidth)) {
+        IndexScale <<= *Constant;
+      } else {
+        break;
+      }
+
+      // If we do not have nuw or nsw on the mul/shl we need to update the flags
+      if (!BO->hasNoSignedWrap())
+        Expr.Flags = Expr.Flags.withoutNoUnsignedSignedWrap();
+      if (!BO->hasNoSignedWrap())
+        Expr.Flags = Expr.Flags.withoutNoUnsignedWrap();
+
+      VarIndex = UnscaledIndex;
+    }
+
     for (gep_type_iterator GTI = gep_type_begin(GEP), GTE = gep_type_end(GEP);
          GTI != GTE; ++GTI) {
       Value *Index = GTI.getOperand();
@@ -963,12 +989,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..668ba298db44f 100644
--- a/llvm/test/Transforms/InstCombine/load-cmp.ll
+++ b/llvm/test/Transforms/InstCombine/load-cmp.ll
@@ -608,3 +608,27 @@ define i1 @cmp_load_multiple_indices2(i32 %idx, i32 %idx2) {
   %cmp = icmp eq i16 %load, 0
   ret i1 %cmp
 }
+
+define i1 @cmp_load_multiplied_index(i32 %X) {
+; CHECK-LABEL: @cmp_load_multiplied_index(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X:%.*]], 1
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %mul = mul nuw nsw i32 %X, 3
+  %gep = getelementptr inbounds [10 x i16], ptr @G16, i32 0, i32 %mul
+  %load = load i16, ptr %gep
+  %cmp = icmp eq i16 %load, 81
+  ret i1 %cmp
+}
+
+define i1 @cmp_load_shifted_index(i32 %X) {
+; CHECK-LABEL: @cmp_load_shifted_index(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[X:%.*]], 1
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %shl = shl nuw nsw i32 %X, 3
+  %gep = getelementptr inbounds [10 x i16], ptr @G16, i32 0, i32 %shl
+  %load = load i16, ptr %gep
+  %cmp = icmp eq i16 %load, 68
+  ret i1 %cmp
+}