[llvm] 4a96803 - [AArch64] Avoid overflow when using shl lower mul (#97148)

[via llvm-commits]

Author: DianQK
Date: 2024-06-30T06:23:27+08:00
New Revision: 4a96803abda2ad74f326de0af6b16552067bda65

URL: https://github.com/llvm/llvm-project/commit/4a96803abda2ad74f326de0af6b16552067bda65
DIFF: https://github.com/llvm/llvm-project/commit/4a96803abda2ad74f326de0af6b16552067bda65.diff

LOG: [AArch64] Avoid overflow when using shl lower mul (#97148)

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.

Added: 
    

Modified: 
    llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
    llvm/test/CodeGen/AArch64/mul_pow2.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 0d53f71a4def8..acce9515e832c 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -18059,16 +18059,27 @@ static SDValue performMulCombine(SDNode *N, SelectionDAG &DAG,
   unsigned ShiftAmt;
 
   auto Shl = [&](SDValue N0, unsigned N1) {
+    if (!N0.getNode())
+      return SDValue();
+    // If shift causes overflow, ignore this combine.
+    if (N1 >= N0.getValueSizeInBits())
+      return SDValue();
     SDValue RHS = DAG.getConstant(N1, DL, MVT::i64);
     return DAG.getNode(ISD::SHL, DL, VT, N0, RHS);
   };
   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..7e26b877a4228 100644
--- a/llvm/test/CodeGen/AArch64/mul_pow2.ll
+++ b/llvm/test/CodeGen/AArch64/mul_pow2.ll
@@ -992,3 +992,41 @@ 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
+}
+
+; Transforming `(mul x, -(2^(N-M) - 1) * 2^M)` to `(sub (shl x, M), (shl x, N))`
+; will not cause overflow when N is 31 and M is 30.
+define i32 @shift_no_overflow(i32 %x) {
+; CHECK-LABEL: shift_no_overflow:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    lsl w8, w0, #31
+; CHECK-NEXT:    sub w0, w8, w0, lsl #30
+; CHECK-NEXT:    ret
+;
+; GISEL-LABEL: shift_no_overflow:
+; GISEL:       // %bb.0:
+; GISEL-NEXT:    mov w8, #1073741824 // =0x40000000
+; GISEL-NEXT:    mul w0, w0, w8
+; GISEL-NEXT:    ret
+  %const = bitcast i32 1073741824 to i32
+  %r = mul i32 %x, %const
+  ret i32 %r
+}