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

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-analysis

Author: Arne Stenkrona (ArneStenkrona2)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/222574.diff


2 Files Affected:

- (modified) llvm/lib/Analysis/Loads.cpp (+25-3) 
- (modified) llvm/unittests/Analysis/LoadsTest.cpp (+44) 


``````````diff
diff --git a/llvm/lib/Analysis/Loads.cpp b/llvm/lib/Analysis/Loads.cpp
index de9022c540d42..cd5aede82385b 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) {
@@ -935,6 +937,26 @@ LinearExpression llvm::decomposeLinearExpression(const DataLayout &DL,
       VarIndex = Index;
     }
 
+    APInt IndexScale(BitWidth, 1);
+    while (auto *BO = dyn_cast_or_null<BinaryOperator>(VarIndex)) {
+      Value *UnscaledIndex = nullptr;
+      APInt Factor(BitWidth, 0);
+      ConstantInt *Constant;
+      if (match(BO, m_c_Mul(m_Value(UnscaledIndex), m_ConstantInt(Constant)))) {
+        Factor = Constant->getValue().zextOrTrunc(BitWidth);
+      } else if (match(BO, m_Shl(m_Value(UnscaledIndex),
+                                 m_ConstantInt(Constant)))) {
+        if (Constant->getValue().uge(BitWidth))
+          break;
+        Factor.setBit(Constant->getZExtValue());
+      } else {
+        break;
+      }
+
+      IndexScale *= Factor;
+      VarIndex = UnscaledIndex;
+    }
+
     // Don't return non-canonical indexes.
     if (VarIndex && !VarIndex->getType()->isIntegerTy(BitWidth))
       return Expr;
@@ -963,12 +985,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/unittests/Analysis/LoadsTest.cpp b/llvm/unittests/Analysis/LoadsTest.cpp
index 8b15bda08485e..a5aa27715e4a2 100644
--- a/llvm/unittests/Analysis/LoadsTest.cpp
+++ b/llvm/unittests/Analysis/LoadsTest.cpp
@@ -17,6 +17,7 @@
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/ValueSymbolTable.h"
 #include "llvm/Support/SourceMgr.h"
 #include "gtest/gtest.h"
 
@@ -243,3 +244,46 @@ loop.end:
   ASSERT_TRUE((NonDerefLoads.size() == 1) &&
               (NonDerefLoads[0]->getName() == "ld1"));
 }
+
+TEST(LoadsTest, DecomposeLinearExpressionScaling) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"IR(
+target datalayout = "e-p:64:64"
+
+define void @f(ptr %base, i64 %index) {
+  %mul = mul i64 %index, 3
+  %shl = shl i64 %index, 2
+  %scaled = shl i64 %mul, 2
+  %gep.mul = getelementptr i32, ptr %base, i64 %mul
+  %gep.shl = getelementptr i32, ptr %base, i64 %shl
+  %gep.scaled = getelementptr i32, ptr %base, i64 %scaled
+  ret void
+}
+)IR");
+  ASSERT_TRUE(M);
+
+  Function *F = M->getFunction("f");
+  ASSERT_TRUE(F);
+  Value *Base = F->getArg(0);
+  Value *Index = F->getArg(1);
+  Value *GEPMul = F->getValueSymbolTable()->lookup("gep.mul");
+  Value *GEPShl = F->getValueSymbolTable()->lookup("gep.shl");
+  Value *GEPScaled = F->getValueSymbolTable()->lookup("gep.scaled");
+  ASSERT_TRUE(GEPMul && GEPShl && GEPScaled);
+
+  const DataLayout &DL = M->getDataLayout();
+  LinearExpression Expr = decomposeLinearExpression(DL, GEPMul);
+  EXPECT_EQ(Base, Expr.BasePtr);
+  EXPECT_EQ(Index, Expr.Index);
+  EXPECT_EQ(APInt(64, 12), Expr.Scale);
+
+  Expr = decomposeLinearExpression(DL, GEPShl);
+  EXPECT_EQ(Base, Expr.BasePtr);
+  EXPECT_EQ(Index, Expr.Index);
+  EXPECT_EQ(APInt(64, 16), Expr.Scale);
+
+  Expr = decomposeLinearExpression(DL, GEPScaled);
+  EXPECT_EQ(Base, Expr.BasePtr);
+  EXPECT_EQ(Index, Expr.Index);
+  EXPECT_EQ(APInt(64, 48), Expr.Scale);
+}

``````````

</details>


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