[AArch64] Optimize overflow checks for [s|u]mul.with.overflow.i32.

[efriedma]

Summary

Saves one instruction for signed, uses a cheaper instruction for unsigned.

[dmgreen]

Sounds good to me. Is the fast ISel mostly updated to keep the test insync? I know that code less, but it seems OK to me.

llvm/lib/Target/AArch64/AArch64FastISel.cpp

3684–3690

Apparently we prefer "Register" now, over unsigned.

[efriedma]

In D105770#2870456, @dmgreen wrote:

Is the fast ISel mostly updated to keep the test insync?

Yes. In retrospect, maybe I should have just updated the test instead, but the change wasn't too complicated.

llvm/lib/Target/AArch64/AArch64FastISel.cpp

case Intrinsic::smul_with_overflow: {
         return false;
 
       unsigned RHSReg = getRegForValue(RHS);
       if (!RHSReg)
         return false;
 
       if (VT == MVT::i32) {
         MulReg = emitSMULL_rr(MVT::i64, LHSReg, RHSReg);
-        unsigned ShiftReg = emitLSR_ri(MVT::i64, MVT::i64, MulReg, 32);
-        MulReg = fastEmitInst_extractsubreg(VT, MulReg, AArch64::sub_32);
-        ShiftReg = fastEmitInst_extractsubreg(VT, ShiftReg, AArch64::sub_32);
-        emitSubs_rs(VT, ShiftReg, MulReg, AArch64_AM::ASR, 31,
+        unsigned MulSubReg =
+            fastEmitInst_extractsubreg(VT, MulReg, AArch64::sub_32);
+        // cmp xreg, wreg, sxtw
+        emitAddSub_rx(/*UseAdd=*/false, MVT::i64, MulReg, MulSubReg,
+                      AArch64_AM::SXTW, /*ShiftImm=*/0, /*SetFlags=*/true,
                       /*WantResult=*/false);
+        MulReg = MulSubReg;
       } else {
         assert(VT == MVT::i64 && "Unexpected value type.");
         // LHSReg and RHSReg cannot be killed by this Mul, since they are
         // reused in the next instruction.
         MulReg = emitMul_rr(VT, LHSReg, RHSReg);
         unsigned SMULHReg = fastEmit_rr(VT, VT, ISD::MULHS, LHSReg, RHSReg);
         emitSubs_rs(VT, SMULHReg, MulReg, AArch64_AM::ASR, 63,
                     /*WantResult=*/false);
       }
       break;
     }
     case Intrinsic::umul_with_overflow: {
       CC = AArch64CC::NE;
       unsigned LHSReg = getRegForValue(LHS);
       if (!LHSReg)
         return false;
 
       unsigned RHSReg = getRegForValue(RHS);
       if (!RHSReg)
         return false;
 
       if (VT == MVT::i32) {
         MulReg = emitUMULL_rr(MVT::i64, LHSReg, RHSReg);
-        emitSubs_rs(MVT::i64, AArch64::XZR, MulReg, AArch64_AM::LSR, 32,
-                    /*WantResult=*/false);
+        // tst xreg, #0xffffffff00000000
+        BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc,
+                TII.get(AArch64::ANDSXri), AArch64::XZR)
+            .addReg(MulReg)
+            .addImm(AArch64_AM::encodeLogicalImmediate(0xFFFFFFFF00000000, 64));
         MulReg = fastEmitInst_extractsubreg(VT, MulReg, AArch64::sub_32);
       } else {
         assert(VT == MVT::i64 && "Unexpected value type.");
         // LHSReg and RHSReg cannot be killed by this Mul, since they are
         // reused in the next instruction.
         MulReg = emitMul_rr(VT, LHSReg, RHSReg);
         unsigned UMULHReg = fastEmit_rr(VT, VT, ISD::MULHU, LHSReg, RHSReg);
         emitSubs_rr(VT, AArch64::XZR, UMULHReg, /*WantResult=*/false);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

case ISD::USUBO:
     CC = AArch64CC::LO;
     break;
   // Multiply needs a little bit extra work.
   case ISD::SMULO:
   case ISD::UMULO: {
     CC = AArch64CC::NE;
     bool IsSigned = Op.getOpcode() == ISD::SMULO;
     if (Op.getValueType() == MVT::i32) {
+      // Extend to 64-bits, then perform a 64-bit multiply.
       unsigned ExtendOpc = IsSigned ? ISD::SIGN_EXTEND : ISD::ZERO_EXTEND;
-      // For a 32 bit multiply with overflow check we want the instruction
-      // selector to generate a widening multiply (SMADDL/UMADDL). For that we
-      // need to generate the following pattern:
-      // (i64 add 0, (i64 mul (i64 sext|zext i32 %a), (i64 sext|zext i32 %b))
       LHS = DAG.getNode(ExtendOpc, DL, MVT::i64, LHS);
       RHS = DAG.getNode(ExtendOpc, DL, MVT::i64, RHS);
       SDValue Mul = DAG.getNode(ISD::MUL, DL, MVT::i64, LHS, RHS);
-      SDValue Add = DAG.getNode(ISD::ADD, DL, MVT::i64, Mul,
-                                DAG.getConstant(0, DL, MVT::i64));
-      // On AArch64 the upper 32 bits are always zero extended for a 32 bit
-      // operation. We need to clear out the upper 32 bits, because we used a
-      // widening multiply that wrote all 64 bits. In the end this should be a
-      // noop.
-      Value = DAG.getNode(ISD::TRUNCATE, DL, MVT::i32, Add);
+      Value = DAG.getNode(ISD::TRUNCATE, DL, MVT::i32, Mul);
+
+      // Check that the result fits into a 32-bit integer.
+      SDVTList VTs = DAG.getVTList(MVT::i64, MVT_CC);
       if (IsSigned) {
-        // The signed overflow check requires more than just a simple check for
-        // any bit set in the upper 32 bits of the result. These bits could be
-        // just the sign bits of a negative number. To perform the overflow
-        // check we have to arithmetic shift right the 32nd bit of the result by
-        // 31 bits. Then we compare the result to the upper 32 bits.
-        SDValue UpperBits = DAG.getNode(ISD::SRL, DL, MVT::i64, Add,
-                                        DAG.getConstant(32, DL, MVT::i64));
-        UpperBits = DAG.getNode(ISD::TRUNCATE, DL, MVT::i32, UpperBits);
-        SDValue LowerBits = DAG.getNode(ISD::SRA, DL, MVT::i32, Value,
-                                        DAG.getConstant(31, DL, MVT::i64));
-        // It is important that LowerBits is last, otherwise the arithmetic
-        // shift will not be folded into the compare (SUBS).
-        SDVTList VTs = DAG.getVTList(MVT::i32, MVT::i32);
-        Overflow = DAG.getNode(AArch64ISD::SUBS, DL, VTs, UpperBits, LowerBits)
-                       .getValue(1);
+        // cmp xreg, wreg, sxtw
+        SDValue SExtMul = DAG.getNode(ISD::SIGN_EXTEND, DL, MVT::i64, Value);
+        Overflow =
+            DAG.getNode(AArch64ISD::SUBS, DL, VTs, Mul, SExtMul).getValue(1);
       } else {
-        // The overflow check for unsigned multiply is easy. We only need to
-        // check if any of the upper 32 bits are set. This can be done with a
-        // CMP (shifted register). For that we need to generate the following
-        // pattern:
-        // (i64 AArch64ISD::SUBS i64 0, (i64 srl i64 %Mul, i64 32)
-        SDValue UpperBits = DAG.getNode(ISD::SRL, DL, MVT::i64, Mul,
-                                        DAG.getConstant(32, DL, MVT::i64));
-        SDVTList VTs = DAG.getVTList(MVT::i64, MVT::i32);
+        // tst xreg, #0xffffffff00000000
+        SDValue UpperBits = DAG.getConstant(0xFFFFFFFF00000000, DL, MVT::i64);
         Overflow =
-            DAG.getNode(AArch64ISD::SUBS, DL, VTs,
-                        DAG.getConstant(0, DL, MVT::i64),
-                        UpperBits).getValue(1);
+            DAG.getNode(AArch64ISD::ANDS, DL, VTs, Mul, UpperBits).getValue(1);
       }
       break;
     }
     assert(Op.getValueType() == MVT::i64 && "Expected an i64 value type");
     // For the 64 bit multiply
     Value = DAG.getNode(ISD::MUL, DL, MVT::i64, LHS, RHS);
     if (IsSigned) {
       SDValue UpperBits = DAG.getNode(ISD::MULHS, DL, MVT::i64, LHS, RHS);

llvm/test/CodeGen/AArch64/arm64-xaluo.ll

; CHECK-NEXT:   cset {{w[0-9]+}}, lo
   store i64 %val, i64* %res
   ret i1 %obit
 }
 
 define zeroext i1 @smulo.i32(i32 %v1, i32 %v2, i32* %res) {
 entry:
 ; CHECK-LABEL:  smulo.i32
 ; CHECK:        smull x[[MREG:[0-9]+]], w0, w1
-; CHECK-NEXT:   lsr x[[SREG:[0-9]+]], x[[MREG]], #32
-; CHECK-NEXT:   cmp w[[SREG]], w[[MREG]], asr #31
+; CHECK-NEXT:   cmp x[[MREG]], w[[MREG]], sxtw
 ; CHECK-NEXT:   cset {{w[0-9]+}}, ne
   %t = call {i32, i1} @llvm.smul.with.overflow.i32(i32 %v1, i32 %v2)
   %val = extractvalue {i32, i1} %t, 0
   %obit = extractvalue {i32, i1} %t, 1
   store i32 %val, i32* %res
   ret i1 %obit
 }
 
; CHECK-NEXT:   cset {{w[0-9]+}}, vs
   store i64 %val, i64* %res
   ret i1 %obit
 }
 
 define zeroext i1 @umulo.i32(i32 %v1, i32 %v2, i32* %res) {
 entry:
 ; CHECK-LABEL:  umulo.i32
 ; CHECK:        umull [[MREG:x[0-9]+]], w0, w1
-; CHECK-NEXT:   cmp xzr, [[MREG]], lsr #32
+; CHECK-NEXT:   tst [[MREG]], #0xffffffff00000000
 ; CHECK-NEXT:   cset {{w[0-9]+}}, ne
   %t = call {i32, i1} @llvm.umul.with.overflow.i32(i32 %v1, i32 %v2)
   %val = extractvalue {i32, i1} %t, 0
   %obit = extractvalue {i32, i1} %t, 1
   store i32 %val, i32* %res
   ret i1 %obit
 }
 
; CHECK-NEXT:   cset w0, hs
   %ret = xor i1 %obit, true
   ret i1 %ret
 }
 
 define i32 @smulo.select.i32(i32 %v1, i32 %v2) {
 entry:
 ; CHECK-LABEL:  smulo.select.i32
 ; CHECK:        smull   x[[MREG:[0-9]+]], w0, w1
-; CHECK-NEXT:   lsr     x[[SREG:[0-9]+]], x[[MREG]], #32
-; CHECK-NEXT:   cmp     w[[SREG]], w[[MREG]], asr #31
+; CHECK-NEXT:   cmp     x[[MREG]], w[[MREG]], sxtw
 ; CHECK-NEXT:   csel    w0, w0, w1, ne
   %t = call {i32, i1} @llvm.smul.with.overflow.i32(i32 %v1, i32 %v2)
   %obit = extractvalue {i32, i1} %t, 1
   %ret = select i1 %obit, i32 %v1, i32 %v2
   ret i32 %ret
 }
 
 define i1 @smulo.not.i32(i32 %v1, i32 %v2) {
 entry:
 ; CHECK-LABEL:  smulo.not.i32
 ; CHECK:        smull   x[[MREG:[0-9]+]], w0, w1
-; CHECK-NEXT:   lsr     x[[SREG:[0-9]+]], x[[MREG]], #32
-; CHECK-NEXT:   cmp     w[[SREG]], w[[MREG]], asr #31
+; CHECK-NEXT:   cmp     x[[MREG]], w[[MREG]], sxtw
 ; CHECK-NEXT:   cset    w0, eq
   %t = call {i32, i1} @llvm.smul.with.overflow.i32(i32 %v1, i32 %v2)
   %obit = extractvalue {i32, i1} %t, 1
   %ret = xor i1 %obit, true
   ret i1 %ret
 }
 
 define i64 @smulo.select.i64(i64 %v1, i64 %v2) {
; CHECK-NEXT:   cset    w0, eq
   %ret = xor i1 %obit, true
   ret i1 %ret
 }
 
 define i32 @umulo.select.i32(i32 %v1, i32 %v2) {
 entry:
 ; CHECK-LABEL:  umulo.select.i32
 ; CHECK:        umull   [[MREG:x[0-9]+]], w0, w1
-; CHECK-NEXT:   cmp     xzr, [[MREG]], lsr #32
+; CHECK-NEXT:   tst     [[MREG]], #0xffffffff00000000
 ; CHECK-NEXT:   csel    w0, w0, w1, ne
   %t = call {i32, i1} @llvm.umul.with.overflow.i32(i32 %v1, i32 %v2)
   %obit = extractvalue {i32, i1} %t, 1
   %ret = select i1 %obit, i32 %v1, i32 %v2
   ret i32 %ret
 }
 
 define i1 @umulo.not.i32(i32 %v1, i32 %v2) {
 entry:
 ; CHECK-LABEL:  umulo.not.i32
 ; CHECK:        umull   [[MREG:x[0-9]+]], w0, w1
-; CHECK-NEXT:   cmp     xzr, [[MREG]], lsr #32
+; CHECK-NEXT:   tst     [[MREG]], #0xffffffff00000000
 ; CHECK-NEXT:   cset    w0, eq
   %t = call {i32, i1} @llvm.umul.with.overflow.i32(i32 %v1, i32 %v2)
   %obit = extractvalue {i32, i1} %t, 1
   %ret = xor i1 %obit, true
   ret i1 %ret
 }
 
 define i64 @umulo.select.i64(i64 %v1, i64 %v2) {

 continue:
   ret i1 true
 }
 
 define zeroext i1 @smulo.br.i32(i32 %v1, i32 %v2) {
 entry:
 ; CHECK-LABEL:  smulo.br.i32
 ; CHECK:        smull   x[[MREG:[0-9]+]], w0, w1
-; CHECK-NEXT:   lsr     x[[SREG:[0-9]+]], x8, #32
-; CHECK-NEXT:   cmp     w[[SREG]], w[[MREG]], asr #31
+; CHECK-NEXT:   cmp     x[[MREG]], w[[MREG]], sxtw
 ; CHECK-NEXT:   b.eq
   %t = call {i32, i1} @llvm.smul.with.overflow.i32(i32 %v1, i32 %v2)
   %val = extractvalue {i32, i1} %t, 0
   %obit = extractvalue {i32, i1} %t, 1
   br i1 %obit, label %overflow, label %continue
 
 overflow:
   ret i1 false

 continue:
   ret i1 true
 }
 
 define zeroext i1 @umulo.br.i32(i32 %v1, i32 %v2) {
 entry:
 ; CHECK-LABEL:  umulo.br.i32
 ; CHECK:        umull   [[MREG:x[0-9]+]], w0, w1
-; CHECK-NEXT:   cmp     xzr, [[MREG]], lsr #32
+; CHECK-NEXT:   tst     [[MREG]], #0xffffffff00000000
 ; CHECK-NEXT:   b.eq
   %t = call {i32, i1} @llvm.umul.with.overflow.i32(i32 %v1, i32 %v2)
   %val = extractvalue {i32, i1} %t, 0
   %obit = extractvalue {i32, i1} %t, 1
   br i1 %obit, label %overflow, label %continue
 
 overflow:
   ret i1 false