[llvm] [SCEV] Fix scMulExpr cost for multiply by -1 and power-of-2 (PR #191033)

[Adel Ejjeh via llvm-commits]

https://github.com/adelejjeh created https://github.com/llvm/llvm-project/pull/191033

## Summary

The SCEV expansion cost model in `costAndCollectOperands` always asks
`TTI.getArithmeticInstrCost(Mul, ...)` for `scMulExpr`, but
`SCEVExpander::visitMulExpr` — the code that actually generates IR
from the SCEV expression — never emits a multiply for two common
cases:

- `-1 * X` → generates `sub 0, X`
- `2^k * X` → generates `shl X, k`

This is target-independent: the expander unconditionally checks for
all-ones and power-of-2 constants before falling back to `mul`.

The mismatch matters on targets where integer multiply is expensive.
On AMDGPU, `i32 mul` has TTI cost 4. A trivial trip count like
`(-1 * start) + end` from `for (i = start; i < end; i++)` gets priced
at 4 (mul) + 1 (add) = 5, exceeding the default
`SCEVCheapExpansionBudget` of 4, and runtime loop unrolling is
rejected. The actual expansion cost is 1 (sub) + 1 (add) = 2.

This patch checks the constant operand in two-operand `scMulExpr` and
costs it as the operation the expander will generate. This is the same
approach already used for `scUDivExpr`, which costs power-of-2
division as `LShr` (lines directly above the changed code).

The general case (non-constant multipliers, >2 operands, binpow) is
left as-is; the existing TODO is preserved.

## Changes

- **`llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp`** — In
  `costAndCollectOperands`, for two-operand `scMulExpr` with a
  `SCEVConstant` LHS: cost all-ones (-1) as `Instruction::Sub` and
  power-of-2 as `Instruction::Shl` instead of `Instruction::Mul`.

- **`llvm/test/Transforms/LoopUnroll/scev-mul-expansion-cost.ll`**
  *(new)* — Verifies an AMDGPU offset-start loop
  (`for (i = start; i < end; i++)`) is runtime-unrolled 8x with the
  default expansion budget. Uses AMDGPU triple because its i32 mul
  cost (4) makes the over-counting observable against the default
  budget (4).

Assisted-by: Cursor (Claude)

Made with [Cursor](https://cursor.com)

>From 2edc5e05b26b4fe69ea11d4ab200e824b8ffc4e2 Mon Sep 17 00:00:00 2001
From: Adel Ejjeh <adel.ejjeh at amd.com>
Date: Wed, 8 Apr 2026 14:06:11 -0500
Subject: [PATCH] [SCEV] Fix scMulExpr cost for multiply by -1 and power-of-2
 Assisted-by: Cursor (Claude)

---
 .../Utils/ScalarEvolutionExpander.cpp         | 24 ++++++++---
 .../LoopUnroll/scev-mul-expansion-cost.ll     | 40 +++++++++++++++++++
 2 files changed, 59 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/Transforms/LoopUnroll/scev-mul-expansion-cost.ll

diff --git a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
index a560c324b4f1e..5f2468f54e2fd 100644
--- a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
+++ b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
@@ -2023,12 +2023,26 @@ template<typename T> static InstructionCost costAndCollectOperands(
   case scAddExpr:
     Cost = ArithCost(Instruction::Add, S->getNumOperands() - 1);
     break;
-  case scMulExpr:
-    // TODO: this is a very pessimistic cost modelling for Mul,
-    // because of Bin Pow algorithm actually used by the expander,
-    // see SCEVExpander::visitMulExpr(), ExpandOpBinPowN().
-    Cost = ArithCost(Instruction::Mul, S->getNumOperands() - 1);
+  case scMulExpr: {
+    // Match the actual expansion in visitMulExpr: multiply by -1 is
+    // expanded as a negate (sub 0, x), and multiply by a power of 2 is
+    // expanded as a shift.  Only handle the common two-operand case with a
+    // constant LHS; for everything else fall back to the pessimistic
+    // all-multiplies estimate.
+    // TODO: this is still pessimistic for the general case because of the
+    // Bin Pow algorithm actually used by the expander, see
+    // SCEVExpander::visitMulExpr(), ExpandOpBinPowN().
+    unsigned MulOpc = Instruction::Mul;
+    if (S->getNumOperands() == 2)
+      if (auto *SC = dyn_cast<SCEVConstant>(S->getOperand(0))) {
+        if (SC->getAPInt().isAllOnes()) // -1
+          MulOpc = Instruction::Sub;
+        else if (SC->getAPInt().isPowerOf2())
+          MulOpc = Instruction::Shl;
+      }
+    Cost = ArithCost(MulOpc, S->getNumOperands() - 1);
     break;
+  }
   case scSMaxExpr:
   case scUMaxExpr:
   case scSMinExpr:
diff --git a/llvm/test/Transforms/LoopUnroll/scev-mul-expansion-cost.ll b/llvm/test/Transforms/LoopUnroll/scev-mul-expansion-cost.ll
new file mode 100644
index 0000000000000..8e188f280fdf3
--- /dev/null
+++ b/llvm/test/Transforms/LoopUnroll/scev-mul-expansion-cost.ll
@@ -0,0 +1,40 @@
+; RUN: opt -S -passes=loop-unroll -unroll-runtime < %s | FileCheck %s
+
+; The trip count SCEV for "for (i = start; i < end; i++)" is
+; (-1 * start) + end.  The SCEV expansion cost model should recognize that
+; (-1 * X) is expanded as a negate (sub 0, X), not a real multiply.  On
+; targets where i32 mul is expensive (e.g. AMDGPU quarter-rate), the old
+; cost model would over-count this as a mul (cost 4) exceeding the default
+; budget of 4, and reject runtime unrolling.  With the fix, it's costed as
+; a sub (cost 1), well within budget.
+
+; AMDGPU is used because its i32 mul has a cost of 4,
+; making the over-counting observable against the default budget of 4.
+target datalayout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128-p9:192:256:256:32-ni:7:8:9"
+target triple = "amdgcn-amd-amdhsa"
+
+; CHECK-LABEL: @offset_start_loop
+; The loop should be runtime-unrolled 8x (prologue + main unrolled body).
+; CHECK: %xtraiter = and i32 %{{.*}}, 7
+; CHECK: loop.prol:
+; CHECK: loop:
+; CHECK: %iv.next.7 = add nsw i32 %iv, 8
+define void @offset_start_loop(ptr addrspace(1) %out, ptr addrspace(1) %in, i32 %start, i32 %end) {
+entry:
+  %cmp = icmp slt i32 %start, %end
+  br i1 %cmp, label %loop, label %exit
+
+loop:
+  %iv = phi i32 [ %start, %entry ], [ %iv.next, %loop ]
+  %idx = sext i32 %iv to i64
+  %gep.in = getelementptr inbounds float, ptr addrspace(1) %in, i64 %idx
+  %val = load float, ptr addrspace(1) %gep.in, align 4
+  %gep.out = getelementptr inbounds float, ptr addrspace(1) %out, i64 %idx
+  store float %val, ptr addrspace(1) %gep.out, align 4
+  %iv.next = add nsw i32 %iv, 1
+  %cond = icmp slt i32 %iv.next, %end
+  br i1 %cond, label %loop, label %exit
+
+exit:
+  ret void
+}