[AArch64] Convert comparisons with 1 and -1 to 0 if it is profitable

AZero13 · AZero13 · commit 27e6a0e13909 · 2025-05-29T15:04:58.000-04:00
This relies on the fact that we can use tst and ands for comparisons as by emitComparison. Relies on #140999. Fixes: #141137.
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -3770,13 +3770,41 @@ static unsigned getCmpOperandFoldingProfit(SDValue Op) {
   return 0;
 }
 
+// emitComparison() converts comparison with one or negative one to comparison
+// with 0.
+
+// Note that this only works for signed comparisons because of how ANDS works.
+static bool shouldBeAdjustedToZero(SDValue LHS, APInt C, ISD::CondCode &CC) {
+  // Only works for ANDS and AND.
+  if (LHS.getOpcode() != ISD::AND && LHS.getOpcode() != AArch64ISD::ANDS)
+    return false;
+
+  if (C.isOne() && (CC == ISD::SETLT || CC == ISD::SETGE)) {
+    CC = (CC == ISD::SETLT) ? ISD::SETLE : ISD::SETGT;
+    return true;
+  }
+
+  if (C.isAllOnes() && (CC == ISD::SETLE || CC == ISD::SETGT)) {
+    CC = (CC == ISD::SETLE) ? ISD::SETLT : ISD::SETGE;
+    return true;
+  }
+
+  return false;
+}
+
 static SDValue getAArch64Cmp(SDValue LHS, SDValue RHS, ISD::CondCode CC,
                              SDValue &AArch64cc, SelectionDAG &DAG,
                              const SDLoc &dl) {
   if (ConstantSDNode *RHSC = dyn_cast<ConstantSDNode>(RHS.getNode())) {
     EVT VT = RHS.getValueType();
     APInt C = RHSC->getAPIntValue();
-    if (!isLegalCmpImmed(C)) {
+    // shouldBeAdjustedToZero is a special case to better fold with
+    // emitComparison().
+    if (shouldBeAdjustedToZero(LHS, C, CC)) {
+      // Adjust the constant to zero.
+      // CC has already been adjusted.
+      RHS = DAG.getConstant(0, dl, VT);
+    } else if (!isLegalCmpImmed(C)) {
       // Constant does not fit, try adjusting it by one?
       switch (CC) {
       default:
diff --git a/llvm/test/CodeGen/AArch64/fptosi-sat-scalar.ll b/llvm/test/CodeGen/AArch64/fptosi-sat-scalar.ll
@@ -23,9 +23,8 @@ define i1 @test_signed_i1_f32(float %f) nounwind {
 ; CHECK-SD-LABEL: test_signed_i1_f32:
 ; CHECK-SD:       // %bb.0:
 ; CHECK-SD-NEXT:    fcvtzs w8, s0
-; CHECK-SD-NEXT:    and w8, w8, w8, asr #31
-; CHECK-SD-NEXT:    cmn w8, #1
-; CHECK-SD-NEXT:    csinv w8, w8, wzr, gt
+; CHECK-SD-NEXT:    ands w8, w8, w8, asr #31
+; CHECK-SD-NEXT:    csinv w8, w8, wzr, ge
 ; CHECK-SD-NEXT:    and w0, w8, #0x1
 ; CHECK-SD-NEXT:    ret
 ;
@@ -269,9 +268,8 @@ define i1 @test_signed_i1_f64(double %f) nounwind {
 ; CHECK-SD-LABEL: test_signed_i1_f64:
 ; CHECK-SD:       // %bb.0:
 ; CHECK-SD-NEXT:    fcvtzs w8, d0
-; CHECK-SD-NEXT:    and w8, w8, w8, asr #31
-; CHECK-SD-NEXT:    cmn w8, #1
-; CHECK-SD-NEXT:    csinv w8, w8, wzr, gt
+; CHECK-SD-NEXT:    ands w8, w8, w8, asr #31
+; CHECK-SD-NEXT:    csinv w8, w8, wzr, ge
 ; CHECK-SD-NEXT:    and w0, w8, #0x1
 ; CHECK-SD-NEXT:    ret
 ;
@@ -520,18 +518,16 @@ define i1 @test_signed_i1_f16(half %f) nounwind {
 ; CHECK-SD-CVT:       // %bb.0:
 ; CHECK-SD-CVT-NEXT:    fcvt s0, h0
 ; CHECK-SD-CVT-NEXT:    fcvtzs w8, s0
-; CHECK-SD-CVT-NEXT:    and w8, w8, w8, asr #31
-; CHECK-SD-CVT-NEXT:    cmn w8, #1
-; CHECK-SD-CVT-NEXT:    csinv w8, w8, wzr, gt
+; CHECK-SD-CVT-NEXT:    ands w8, w8, w8, asr #31
+; CHECK-SD-CVT-NEXT:    csinv w8, w8, wzr, ge
 ; CHECK-SD-CVT-NEXT:    and w0, w8, #0x1
 ; CHECK-SD-CVT-NEXT:    ret
 ;
 ; CHECK-SD-FP16-LABEL: test_signed_i1_f16:
 ; CHECK-SD-FP16:       // %bb.0:
 ; CHECK-SD-FP16-NEXT:    fcvtzs w8, h0
-; CHECK-SD-FP16-NEXT:    and w8, w8, w8, asr #31
-; CHECK-SD-FP16-NEXT:    cmn w8, #1
-; CHECK-SD-FP16-NEXT:    csinv w8, w8, wzr, gt
+; CHECK-SD-FP16-NEXT:    ands w8, w8, w8, asr #31
+; CHECK-SD-FP16-NEXT:    csinv w8, w8, wzr, ge
 ; CHECK-SD-FP16-NEXT:    and w0, w8, #0x1
 ; CHECK-SD-FP16-NEXT:    ret
 ;
diff --git a/llvm/test/CodeGen/AArch64/fptosi-sat-vector.ll b/llvm/test/CodeGen/AArch64/fptosi-sat-vector.ll
@@ -2371,12 +2371,10 @@ define <2 x i1> @test_signed_v2f64_v2i1(<2 x double> %f) {
 ; CHECK-SD-NEXT:    mov d1, v0.d[1]
 ; CHECK-SD-NEXT:    fcvtzs w9, d0
 ; CHECK-SD-NEXT:    fcvtzs w8, d1
-; CHECK-SD-NEXT:    and w9, w9, w9, asr #31
-; CHECK-SD-NEXT:    and w8, w8, w8, asr #31
-; CHECK-SD-NEXT:    cmn w8, #1
-; CHECK-SD-NEXT:    csinv w8, w8, wzr, gt
-; CHECK-SD-NEXT:    cmn w9, #1
-; CHECK-SD-NEXT:    csinv w9, w9, wzr, gt
+; CHECK-SD-NEXT:    ands w8, w8, w8, asr #31
+; CHECK-SD-NEXT:    csinv w8, w8, wzr, ge
+; CHECK-SD-NEXT:    ands w9, w9, w9, asr #31
+; CHECK-SD-NEXT:    csinv w9, w9, wzr, ge
 ; CHECK-SD-NEXT:    fmov s0, w9
 ; CHECK-SD-NEXT:    mov v0.s[1], w8
 ; CHECK-SD-NEXT:    // kill: def $d0 killed $d0 killed $q0
diff --git a/llvm/test/CodeGen/AArch64/logical_shifted_reg.ll b/llvm/test/CodeGen/AArch64/logical_shifted_reg.ll
@@ -257,9 +257,8 @@ define void @flag_setting() {
 ; CHECK-NEXT:    tst x9, x10, lsl #63
 ; CHECK-NEXT:    b.lt .LBB2_4
 ; CHECK-NEXT:  // %bb.2: // %test3
-; CHECK-NEXT:    and x10, x9, x10, asr #12
-; CHECK-NEXT:    cmp x10, #1
-; CHECK-NEXT:    b.ge .LBB2_4
+; CHECK-NEXT:    tst x9, x10, asr #12
+; CHECK-NEXT:    b.gt .LBB2_4
 ; CHECK-NEXT:  // %bb.3: // %other_exit
 ; CHECK-NEXT:    str x9, [x8]
 ; CHECK-NEXT:  .LBB2_4: // %common.ret
diff --git a/llvm/test/CodeGen/AArch64/tbz-tbnz.ll b/llvm/test/CodeGen/AArch64/tbz-tbnz.ll
@@ -199,16 +199,14 @@ if.end:
 define void @test8(i64 %val1, i64 %val2, i64 %val3) {
 ; CHECK-SD-LABEL: test8:
 ; CHECK-SD:       // %bb.0:
-; CHECK-SD-NEXT:    and x8, x0, x1
-; CHECK-SD-NEXT:    cmn x8, #1
-; CHECK-SD-NEXT:    b.gt .LBB7_3
+; CHECK-SD-NEXT:    tst x0, x1
+; CHECK-SD-NEXT:    b.ge .LBB7_3
 ; CHECK-SD-NEXT:  // %bb.1:
 ; CHECK-SD-NEXT:    and x8, x1, x2
 ; CHECK-SD-NEXT:    tbnz x8, #63, .LBB7_3
 ; CHECK-SD-NEXT:  // %bb.2: // %if.then2
-; CHECK-SD-NEXT:    and x8, x0, x1, lsl #63
-; CHECK-SD-NEXT:    cmn x8, #1
-; CHECK-SD-NEXT:    b.le .LBB7_4
+; CHECK-SD-NEXT:    tst x0, x1, lsl #63
+; CHECK-SD-NEXT:    b.lt .LBB7_4
 ; CHECK-SD-NEXT:  .LBB7_3: // %if.end
 ; CHECK-SD-NEXT:    ret
 ; CHECK-SD-NEXT:  .LBB7_4: // %if.then3
