[Global-ISel][AArch64] Port ShouldBeAdjustedToZero and allow negative cmp immediates in global-isel

To undo the mitigations, shouldBeAdjustedToZero also includes one-use checks for cmn 1 to cmp 0 because of the pl/mi favoritism. This was also brought to SelectionDAG.

Author: AZero13

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -3961,20 +3961,24 @@ static unsigned getCmpOperandFoldingProfit(SDValue Op) {
 }
 
 // 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.
+// with 0.
 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)
+  // TODO: Is this too restrictive? This is just to prevent CSE with other
+  // comparisons.
+  if (LHS.getOpcode() != ISD::AND && LHS.getOpcode() != AArch64ISD::ANDS &&
+      !LHS.hasOneUse())
     return false;
 
-  if (C.isOne() && (CC == ISD::SETLT || CC == ISD::SETGE)) {
-    CC = (CC == ISD::SETLT) ? ISD::SETLE : ISD::SETGT;
+  if (C.isAllOnes() && (CC == ISD::SETLE || CC == ISD::SETGT)) {
+    CC = (CC == ISD::SETLE) ? ISD::SETLT : ISD::SETGE;
     return true;
   }
 
-  if (C.isAllOnes() && (CC == ISD::SETLE || CC == ISD::SETGT)) {
-    CC = (CC == ISD::SETLE) ? ISD::SETLT : ISD::SETGE;
+  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;
   }
 
@@ -4035,13 +4039,12 @@ static SDValue getAArch64Cmp(SDValue LHS, SDValue RHS, ISD::CondCode CC,
         break;
       case ISD::SETULE:
       case ISD::SETUGT: {
-        if (!C.isAllOnes()) {
-          APInt CPlusOne = C + 1;
-          if (isLegalCmpImmed(CPlusOne) ||
-              (NumImmForC > numberOfInstrToLoadImm(CPlusOne))) {
-            CC = (CC == ISD::SETULE) ? ISD::SETULT : ISD::SETUGE;
-            RHS = DAG.getConstant(CPlusOne, DL, VT);
-          }
+        assert(!C.isAllOnes() && "C should not be -1 here");
+        APInt CPlusOne = C + 1;
+        if (isLegalCmpImmed(CPlusOne) ||
+            (NumImmForC > numberOfInstrToLoadImm(CPlusOne))) {
+          CC = (CC == ISD::SETULE) ? ISD::SETULT : ISD::SETUGE;
+          RHS = DAG.getConstant(CPlusOne, DL, VT);
         }
         break;
       }

llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp

@@ -561,6 +561,46 @@ void applyVAshrLshrImm(MachineInstr &MI, MachineRegisterInfo &MRI,
   MI.eraseFromParent();
 }
 
+bool isLegalCmpImmed(APInt C) {
+  // Works for negative immediates too, as it can be written as an ADDS
+  // instruction with a negated immediate.
+  return isLegalArithImmed(C.abs().getZExtValue());
+}
+
+/// Check if a comparison with 1 or -1 should be adjusted to compare with 0.
+/// This only works for signed comparisons because of how ANDS works.
+///
+/// \param LHS - The left-hand side register of the comparison
+/// \param C - The constant value (1 or -1)
+/// \param P - The predicate to potentially adjust
+/// \param MRI - Machine register info for looking up definitions
+/// \returns true if the comparison should be adjusted to compare with 0
+static bool shouldBeAdjustedToZero(Register LHS, APInt C, CmpInst::Predicate &P,
+                                   const MachineRegisterInfo &MRI) {
+  // Only works for AND operations
+  MachineInstr *LHSDef = getDefIgnoringCopies(LHS, MRI);
+
+  // TODO: Too restrictive?
+  if (!LHSDef ||
+      (LHSDef->getOpcode() != TargetOpcode::G_AND && !MRI.hasOneUse(LHS)))
+    return false;
+
+  if (C.isAllOnes() && (P == CmpInst::ICMP_SLE || P == CmpInst::ICMP_SGT)) {
+    P = (P == CmpInst::ICMP_SLE) ? CmpInst::ICMP_SLT : CmpInst::ICMP_SGE;
+    return true;
+  }
+
+  if (LHSDef->getOpcode() != TargetOpcode::G_AND)
+    return false;
+
+  if (C.isOne() && (P == CmpInst::ICMP_SLT || P == CmpInst::ICMP_SGE)) {
+    P = (P == CmpInst::ICMP_SLT) ? CmpInst::ICMP_SLE : CmpInst::ICMP_SGT;
+    return true;
+  }
+
+  return false;
+}
+
 /// Determine if it is possible to modify the \p RHS and predicate \p P of a
 /// G_ICMP instruction such that the right-hand side is an arithmetic immediate.
 ///
@@ -569,7 +609,7 @@ void applyVAshrLshrImm(MachineInstr &MI, MachineRegisterInfo &MRI,
 ///
 /// \note This assumes that the comparison has been legalized.
 std::optional<std::pair<uint64_t, CmpInst::Predicate>>
-tryAdjustICmpImmAndPred(Register RHS, CmpInst::Predicate P,
+tryAdjustICmpImmAndPred(Register LHS, Register RHS, CmpInst::Predicate P,
                         const MachineRegisterInfo &MRI) {
   const auto &Ty = MRI.getType(RHS);
   if (Ty.isVector())
@@ -582,11 +622,18 @@ tryAdjustICmpImmAndPred(Register RHS, CmpInst::Predicate P,
   auto ValAndVReg = getIConstantVRegValWithLookThrough(RHS, MRI);
   if (!ValAndVReg)
     return std::nullopt;
-  uint64_t OriginalC = ValAndVReg->Value.getZExtValue();
-  uint64_t C = OriginalC;
-  if (isLegalArithImmed(C))
+
+  APInt C = ValAndVReg->Value;
+
+  // Check if this is a comparison with 1 or -1 that should be adjusted to 0
+  if (shouldBeAdjustedToZero(LHS, C, P, MRI))
+    return {{0, P}};
+
+  if (isLegalCmpImmed(C))
     return std::nullopt;
 
+  uint64_t OriginalC = C.getZExtValue();
+
   // We have a non-arithmetic immediate. Check if adjusting the immediate and
   // adjusting the predicate will result in a legal arithmetic immediate.
   switch (P) {
@@ -599,9 +646,7 @@ tryAdjustICmpImmAndPred(Register RHS, CmpInst::Predicate P,
     // x slt c => x sle c - 1
     // x sge c => x sgt c - 1
     //
-    // When c is not the smallest possible negative number.
-    if ((Size == 64 && static_cast<int64_t>(C) == INT64_MIN) ||
-        (Size == 32 && static_cast<int32_t>(C) == INT32_MIN))
+    if (C.isMinSignedValue())
       return std::nullopt;
     P = (P == CmpInst::ICMP_SLT) ? CmpInst::ICMP_SLE : CmpInst::ICMP_SGT;
     C -= 1;
@@ -614,9 +659,9 @@ tryAdjustICmpImmAndPred(Register RHS, CmpInst::Predicate P,
     // x uge c => x ugt c - 1
     //
     // When c is not zero.
-    assert(C != 0 && "C should not be zero here!");
+    assert(!C.isZero() && "C should not be zero here!");
     P = (P == CmpInst::ICMP_ULT) ? CmpInst::ICMP_ULE : CmpInst::ICMP_UGT;
-    C -= 1;
+    C = C - 1;
     break;
   case CmpInst::ICMP_SLE:
   case CmpInst::ICMP_SGT:
@@ -626,11 +671,10 @@ tryAdjustICmpImmAndPred(Register RHS, CmpInst::Predicate P,
     // x sgt c => s sge c + 1
     //
     // When c is not the largest possible signed integer.
-    if ((Size == 32 && static_cast<int32_t>(C) == INT32_MAX) ||
-        (Size == 64 && static_cast<int64_t>(C) == INT64_MAX))
+    if (C.isMaxSignedValue())
       return std::nullopt;
     P = (P == CmpInst::ICMP_SLE) ? CmpInst::ICMP_SLT : CmpInst::ICMP_SGE;
-    C += 1;
+    C = C + 1;
     break;
   case CmpInst::ICMP_ULE:
   case CmpInst::ICMP_UGT:
@@ -640,29 +684,27 @@ tryAdjustICmpImmAndPred(Register RHS, CmpInst::Predicate P,
     // x ugt c => s uge c + 1
     //
     // When c is not the largest possible unsigned integer.
-    if ((Size == 32 && static_cast<uint32_t>(C) == UINT32_MAX) ||
-        (Size == 64 && C == UINT64_MAX))
-      return std::nullopt;
+    assert(!C.isMaxValue() &&
+           "C should not be -1 here, as it is a valid legal immediate!");
     P = (P == CmpInst::ICMP_ULE) ? CmpInst::ICMP_ULT : CmpInst::ICMP_UGE;
-    C += 1;
+    C = C + 1;
     break;
   }
 
   // Check if the new constant is valid, and return the updated constant and
   // predicate if it is.
-  if (Size == 32)
-    C = static_cast<uint32_t>(C);
-  if (isLegalArithImmed(C))
-    return {{C, P}};
+  uint64_t NewC = C.getZExtValue();
+  if (isLegalCmpImmed(C))
+    return {{NewC, P}};
 
   auto NumberOfInstrToLoadImm = [=](uint64_t Imm) {
     SmallVector<AArch64_IMM::ImmInsnModel> Insn;
     AArch64_IMM::expandMOVImm(Imm, 32, Insn);
     return Insn.size();
   };
 
-  if (NumberOfInstrToLoadImm(OriginalC) > NumberOfInstrToLoadImm(C))
-    return {{C, P}};
+  if (NumberOfInstrToLoadImm(OriginalC) > NumberOfInstrToLoadImm(NewC))
+    return {{NewC, P}};
 
   return std::nullopt;
 }
@@ -679,9 +721,10 @@ bool matchAdjustICmpImmAndPred(
     MachineInstr &MI, const MachineRegisterInfo &MRI,
     std::pair<uint64_t, CmpInst::Predicate> &MatchInfo) {
   assert(MI.getOpcode() == TargetOpcode::G_ICMP);
+  Register LHS = MI.getOperand(2).getReg();
   Register RHS = MI.getOperand(3).getReg();
   auto Pred = static_cast<CmpInst::Predicate>(MI.getOperand(1).getPredicate());
-  if (auto MaybeNewImmAndPred = tryAdjustICmpImmAndPred(RHS, Pred, MRI)) {
+  if (auto MaybeNewImmAndPred = tryAdjustICmpImmAndPred(LHS, RHS, Pred, MRI)) {
     MatchInfo = *MaybeNewImmAndPred;
     return true;
   }

llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-lowering-adjust-icmp-imm.mir

@@ -649,8 +649,8 @@ body:             |
     ; LOWER-NEXT: {{  $}}
     ; LOWER-NEXT: %reg0:_(s32) = COPY $w0
     ; LOWER-NEXT: %reg1:_(s32) = COPY $w1
-    ; LOWER-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
-    ; LOWER-NEXT: %cmp:_(s32) = G_ICMP intpred(slt), %reg0(s32), [[C]]
+    ; LOWER-NEXT: %cst:_(s32) = G_CONSTANT i32 -1
+    ; LOWER-NEXT: %cmp:_(s32) = G_ICMP intpred(sle), %reg0(s32), %cst
     ; LOWER-NEXT: %select:_(s32) = G_SELECT %cmp(s32), %reg0, %reg1
     ; LOWER-NEXT: $w0 = COPY %select(s32)
     ; LOWER-NEXT: RET_ReallyLR implicit $w0
@@ -660,8 +660,8 @@ body:             |
     ; SELECT-NEXT: {{  $}}
     ; SELECT-NEXT: %reg0:gpr32common = COPY $w0
     ; SELECT-NEXT: %reg1:gpr32 = COPY $w1
-    ; SELECT-NEXT: [[SUBSWri:%[0-9]+]]:gpr32 = SUBSWri %reg0, 0, 0, implicit-def $nzcv
-    ; SELECT-NEXT: %select:gpr32 = CSELWr %reg0, %reg1, 4, implicit $nzcv
+    ; SELECT-NEXT: [[ADDSWri:%[0-9]+]]:gpr32 = ADDSWri %reg0, 1, 0, implicit-def $nzcv
+    ; SELECT-NEXT: %select:gpr32 = CSELWr %reg0, %reg1, 13, implicit $nzcv
     ; SELECT-NEXT: $w0 = COPY %select
     ; SELECT-NEXT: RET_ReallyLR implicit $w0
     %reg0:_(s32) = COPY $w0

llvm/test/CodeGen/AArch64/check-sign-bit-before-extension.ll

@@ -14,8 +14,8 @@ define i32 @f_i8_sign_extend_inreg(i8 %in, i32 %a, i32 %b) nounwind {
 ; CHECK-LABEL: f_i8_sign_extend_inreg:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    sxtb w8, w0
-; CHECK-NEXT:    cmn w8, #1
-; CHECK-NEXT:    csel w8, w1, w2, gt
+; CHECK-NEXT:    cmp w8, #0
+; CHECK-NEXT:    csel w8, w1, w2, pl
 ; CHECK-NEXT:    add w0, w8, w0, uxtb
 ; CHECK-NEXT:    ret
 entry:
@@ -36,8 +36,8 @@ define i32 @f_i16_sign_extend_inreg(i16 %in, i32 %a, i32 %b) nounwind {
 ; CHECK-LABEL: f_i16_sign_extend_inreg:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    sxth w8, w0
-; CHECK-NEXT:    cmn w8, #1
-; CHECK-NEXT:    csel w8, w1, w2, gt
+; CHECK-NEXT:    cmp w8, #0
+; CHECK-NEXT:    csel w8, w1, w2, pl
 ; CHECK-NEXT:    add w0, w8, w0, uxth
 ; CHECK-NEXT:    ret
 entry:
@@ -145,8 +145,8 @@ define i64 @f_i32_sign_extend_i64(i32 %in, i64 %a, i64 %b) nounwind {
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    // kill: def $w0 killed $w0 def $x0
 ; CHECK-NEXT:    sxtw x8, w0
-; CHECK-NEXT:    cmn x8, #1
-; CHECK-NEXT:    csel x8, x1, x2, gt
+; CHECK-NEXT:    cmp x8, #0
+; CHECK-NEXT:    csel x8, x1, x2, pl
 ; CHECK-NEXT:    add x0, x8, w0, uxtw
 ; CHECK-NEXT:    ret
 entry: