[llvm] [AArch64] Avoid overflow when using shl lower mul (PR #97148)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-aarch64

Author: DianQK (DianQK)

<details>
<summary>Changes</summary>

Fixes #<!-- -->97147.

Transforming `(mul x, -(2^(N-M) - 1) * 2^M)` to `(sub (shl x, M), (shl x, N))` will cause overflow when N is 32 and M is 31. I still added checks for all scenarios, even other scenarios don't seem to cause overflow.

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


2 Files Affected:

- (modified) llvm/lib/Target/AArch64/AArch64ISelLowering.cpp (+13-1) 
- (modified) llvm/test/CodeGen/AArch64/mul_pow2.ll (+19) 


``````````diff
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 0d53f71a4def8..e022646969b62 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -18059,16 +18059,28 @@ static SDValue performMulCombine(SDNode *N, SelectionDAG &DAG,
   unsigned ShiftAmt;
 
   auto Shl = [&](SDValue N0, unsigned N1) {
+    if (!N0.getNode())
+      return SDValue();
     SDValue RHS = DAG.getConstant(N1, DL, MVT::i64);
-    return DAG.getNode(ISD::SHL, DL, VT, N0, RHS);
+    SDValue SHL = DAG.getNode(ISD::SHL, DL, VT, N0, RHS);
+    // If shift causes overflow, ignore this combine.
+    if (SHL->isUndef())
+      return SDValue();
+    return SHL;
   };
   auto Add = [&](SDValue N0, SDValue N1) {
+    if (!N0.getNode() || !N1.getNode())
+      return SDValue();
     return DAG.getNode(ISD::ADD, DL, VT, N0, N1);
   };
   auto Sub = [&](SDValue N0, SDValue N1) {
+    if (!N0.getNode() || !N1.getNode())
+      return SDValue();
     return DAG.getNode(ISD::SUB, DL, VT, N0, N1);
   };
   auto Negate = [&](SDValue N) {
+    if (!N0.getNode())
+      return SDValue();
     SDValue Zero = DAG.getConstant(0, DL, VT);
     return DAG.getNode(ISD::SUB, DL, VT, Zero, N);
   };
diff --git a/llvm/test/CodeGen/AArch64/mul_pow2.ll b/llvm/test/CodeGen/AArch64/mul_pow2.ll
index c4839175ded5a..16a47c9a49a05 100644
--- a/llvm/test/CodeGen/AArch64/mul_pow2.ll
+++ b/llvm/test/CodeGen/AArch64/mul_pow2.ll
@@ -992,3 +992,22 @@ define <4 x i32> @muladd_demand_commute(<4 x i32> %x, <4 x i32> %y) {
   %r = and <4 x i32> %a, <i32 131071, i32 131071, i32 131071, i32 131071>
   ret <4 x i32> %r
 }
+
+; Transforming `(mul x, -(2^(N-M) - 1) * 2^M)` to `(sub (shl x, M), (shl x, N))`
+; will cause overflow when N is 32 and M is 31.
+define i32 @shift_overflow(i32 %x) {
+; CHECK-LABEL: shift_overflow:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #-2147483648 // =0x80000000
+; CHECK-NEXT:    mul w0, w0, w8
+; CHECK-NEXT:    ret
+;
+; GISEL-LABEL: shift_overflow:
+; GISEL:       // %bb.0:
+; GISEL-NEXT:    mov w8, #-2147483648 // =0x80000000
+; GISEL-NEXT:    mul w0, w0, w8
+; GISEL-NEXT:    ret
+  %const = bitcast i32 2147483648 to i32
+  %r = mul i32 %x, %const
+  ret i32 %r
+}

``````````

</details>


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