[llvm] [DAGCombiner] Preserve nsw when folding (mul x, 2^c) to (shl x, c) (PR #192366)

[Jim Lin via llvm-commits]

https://github.com/tclin914 created https://github.com/llvm/llvm-project/pull/192366

When the multiplier is a power of two strictly less than 2^(BitWidth-1),
the signed multiplication cannot overflow any more than the equivalent
shift, so nsw transfers from the mul to the shl. At 2^(BitWidth-1) the
multiplier is the signed minimum and the equivalence breaks, so keep
the existing behaviour of dropping nsw in that case.

>From a55cc25b34e2e9c34b2ea70300db0db3b110f237 Mon Sep 17 00:00:00 2001
From: Jim Lin <jim at andestech.com>
Date: Wed, 15 Apr 2026 05:51:48 +0000
Subject: [PATCH 1/2] [RISCV] Pre-commit test for (mul x, 2^c) -> (shl x, c)
 flag preservation

Baseline test showing that DAGCombiner currently drops the nsw flag when
folding (mul nsw x, 2^c) to (shl x, c). Preparatory for a follow-up that
preserves nsw when the shift amount is less than BitWidth - 1.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply at anthropic.com>
---
 llvm/test/CodeGen/RISCV/mul-to-shl-nsw.ll | 40 +++++++++++++++++++++++
 1 file changed, 40 insertions(+)
 create mode 100644 llvm/test/CodeGen/RISCV/mul-to-shl-nsw.ll

diff --git a/llvm/test/CodeGen/RISCV/mul-to-shl-nsw.ll b/llvm/test/CodeGen/RISCV/mul-to-shl-nsw.ll
new file mode 100644
index 0000000000000..1a0a792c32628
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/mul-to-shl-nsw.ll
@@ -0,0 +1,40 @@
+; REQUIRES: asserts
+; RUN: llc -mtriple=riscv64 -debug-only=isel -o /dev/null %s 2>&1 | FileCheck %s
+
+; Verify that DAGCombiner's (mul x, 2^c) -> (shl x, c) fold preserves the nsw
+; flag when the shift amount is strictly less than BitWidth - 1, and does not
+; preserve it when the multiplier is the signed minimum (2^(BitWidth - 1)).
+; RISC-V has no unified IMUL-with-flags instruction, so this inspects the
+; SelectionDAG flags directly via the isel debug dump.
+
+define i32 @mul_nsw_pow2(i32 %x) {
+; CHECK-LABEL: === mul_nsw_pow2
+; CHECK: Optimized lowered selection DAG
+; CHECK: i32 = shl {{[^n].*}}, Constant:i64<2>
+  %mul = mul nsw i32 %x, 4
+  ret i32 %mul
+}
+
+define i32 @mul_nsw_pow2_signmin(i32 %x) {
+; CHECK-LABEL: === mul_nsw_pow2_signmin
+; CHECK: Optimized lowered selection DAG
+; CHECK: i32 = shl {{[^n].*}}, Constant:i64<31>
+  %mul = mul nsw i32 %x, -2147483648
+  ret i32 %mul
+}
+
+define i32 @mul_no_flags_pow2(i32 %x) {
+; CHECK-LABEL: === mul_no_flags_pow2
+; CHECK: Optimized lowered selection DAG
+; CHECK: i32 = shl {{[^n].*}}, Constant:i64<2>
+  %mul = mul i32 %x, 4
+  ret i32 %mul
+}
+
+define i32 @mul_nuw_pow2(i32 %x) {
+; CHECK-LABEL: === mul_nuw_pow2
+; CHECK: Optimized lowered selection DAG
+; CHECK: i32 = shl nuw {{.*}}, Constant:i64<2>
+  %mul = mul nuw i32 %x, 4
+  ret i32 %mul
+}

>From 0e261a96335287aab2f7d0537e4269dc572a7487 Mon Sep 17 00:00:00 2001
From: Jim Lin <jim at andestech.com>
Date: Wed, 15 Apr 2026 05:52:31 +0000
Subject: [PATCH 2/2] [DAGCombiner] Preserve nsw when folding (mul x, 2^c) to
 (shl x, c)

When the multiplier is a power of two strictly less than 2^(BitWidth-1),
the signed multiplication cannot overflow any more than the equivalent
shift, so nsw transfers from the mul to the shl. At 2^(BitWidth-1) the
multiplier is the signed minimum and the equivalence breaks, so keep
the existing behaviour of dropping nsw in that case.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply at anthropic.com>
---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 6 +++++-
 llvm/test/CodeGen/RISCV/mul-to-shl-nsw.ll     | 2 +-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index aee5aabdc9e02..0625aa3f324f8 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -4832,7 +4832,11 @@ template <class MatchContextClass> SDValue DAGCombiner::visitMUL(SDNode *N) {
       SDValue Trunc = DAG.getZExtOrTrunc(LogBase2, DL, ShiftVT);
       SDNodeFlags Flags;
       Flags.setNoUnsignedWrap(N->getFlags().hasNoUnsignedWrap());
-      // TODO: Preserve setNoSignedWrap if LogBase2 isn't BitWidth - 1.
+      // Preserve nsw when the shift amount is strictly less than BitWidth - 1,
+      // i.e. the multiplier is not the signed minimum value.
+      if (N->getFlags().hasNoSignedWrap() && N1IsConst &&
+          ConstValue1.logBase2() < BitWidth - 1)
+        Flags.setNoSignedWrap(true);
       return Matcher.getNode(ISD::SHL, DL, VT, N0, Trunc, Flags);
     }
   }
diff --git a/llvm/test/CodeGen/RISCV/mul-to-shl-nsw.ll b/llvm/test/CodeGen/RISCV/mul-to-shl-nsw.ll
index 1a0a792c32628..8f926443b2a5c 100644
--- a/llvm/test/CodeGen/RISCV/mul-to-shl-nsw.ll
+++ b/llvm/test/CodeGen/RISCV/mul-to-shl-nsw.ll
@@ -10,7 +10,7 @@
 define i32 @mul_nsw_pow2(i32 %x) {
 ; CHECK-LABEL: === mul_nsw_pow2
 ; CHECK: Optimized lowered selection DAG
-; CHECK: i32 = shl {{[^n].*}}, Constant:i64<2>
+; CHECK: i32 = shl nsw {{.*}}, Constant:i64<2>
   %mul = mul nsw i32 %x, 4
   ret i32 %mul
 }