Use single ISel opcode for register-register and register-immediate CB

dtellenbach · dtellenbach · commit df6e8314e22c · 2025-02-18T20:50:06.000-08:00
diff --git a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -507,6 +507,9 @@ class AArch64DAGToDAGISel : public SelectionDAGISel {
 
   bool SelectAllActivePredicate(SDValue N);
   bool SelectAnyPredicate(SDValue N);
+
+  template <int Bits>
+  bool SelectCmpBranchUImm6Operand(SDNode *P, SDValue N, SDValue &Imm);
 };
 
 class AArch64DAGToDAGISelLegacy : public SelectionDAGISelLegacy {
@@ -7489,3 +7492,57 @@ bool AArch64DAGToDAGISel::SelectSMETileSlice(SDValue N, unsigned MaxSize,
   Offset = CurDAG->getTargetConstant(0, SDLoc(N), MVT::i64);
   return true;
 }
+
+template <int Bits>
+bool AArch64DAGToDAGISel::SelectCmpBranchUImm6Operand(SDNode *P, SDValue N,
+                                                      SDValue &Imm) {
+  ConstantSDNode *C = dyn_cast<ConstantSDNode>(P->getOperand(1));
+  if (!C)
+    return false;
+
+  AArch64CC::CondCode CC = static_cast<AArch64CC::CondCode>(C->getZExtValue());
+  if (auto *CN = dyn_cast<ConstantSDNode>(N)) {
+    // Check conservatively if the immediate fits the valid range [0, 64).
+    // Immediate variants for GE and HS definitely need to be decremented
+    // when lowering the pseudos later, so an immediate of 1 would become 0.
+    // For the inverse conditions LT and LO we don't know for sure if they
+    // will need a decrement but should the decision be made to reverse the
+    // branch condition, we again end up with the need to decrement.
+    // The same argument holds for LE, LS, GT and HI and possibly
+    // incremented immediates. This can lead to slightly less optimal
+    // codegen, e.g. we never codegen the legal case
+    //    cblt w0, #63, A
+    // because we could end up with the illegal case
+    //    cbge w0, #64, B
+    // should the decision to reverse the branch direction be made. For the
+    // lower bound cases this is no problem since we can express comparisons
+    // against 0 with either tbz/tnbz or using wzr/xzr.
+    uint64_t LowerBound = 0, UpperBound = 64;
+    switch (CC) {
+    case AArch64CC::GE:
+    case AArch64CC::HS:
+    case AArch64CC::LT:
+    case AArch64CC::LO:
+      LowerBound = 1;
+      break;
+    case AArch64CC::LE:
+    case AArch64CC::LS:
+    case AArch64CC::GT:
+    case AArch64CC::HI:
+      UpperBound = 63;
+      break;
+    default:
+      break;
+    }
+
+    if (CN->getAPIntValue().uge(LowerBound) &&
+        CN->getAPIntValue().ult(UpperBound)) {
+      SDLoc DL(N);
+      Imm = CurDAG->getTargetConstant(CN->getZExtValue(), DL,
+                                      Bits == 32 ? MVT::i32 : MVT::i64);
+      return true;
+    }
+  }
+
+  return false;
+}
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -2983,8 +2983,7 @@ const char *AArch64TargetLowering::getTargetNodeName(unsigned Opcode) const {
     MAKE_CASE(AArch64ISD::CTTZ_ELTS)
     MAKE_CASE(AArch64ISD::CALL_ARM64EC_TO_X64)
     MAKE_CASE(AArch64ISD::URSHR_I_PRED)
-    MAKE_CASE(AArch64ISD::CBRR)
-    MAKE_CASE(AArch64ISD::CBRI)
+    MAKE_CASE(AArch64ISD::CB)
   }
 #undef MAKE_CASE
   return nullptr;
@@ -10600,49 +10599,10 @@ SDValue AArch64TargetLowering::LowerBR_CC(SDValue Op, SelectionDAG &DAG) const {
     if (Subtarget->hasCMPBR() &&
         AArch64CC::isValidCBCond(changeIntCCToAArch64CC(CC)) &&
         ProduceNonFlagSettingCondBr) {
-      AArch64CC::CondCode ACC = changeIntCCToAArch64CC(CC);
-      unsigned Opc = AArch64ISD::CBRR;
-      if (auto *Imm = dyn_cast<ConstantSDNode>(RHS)) {
-        // Check conservatively if the immediate fits the valid range [0, 64).
-        // Immediate variants for GE and HS definitely need to be decremented
-        // when lowering the pseudos later, so an immediate of 1 would become 0.
-        // For the inverse conditions LT and LO we don't know for sure if they
-        // will need a decrement but should the decision be made to reverse the
-        // branch condition, we again end up with the need to decrement.
-        // The same argument holds for LE, LS, GT and HI and possibly
-        // incremented immediates. This can lead to slightly less optimal
-        // codegen, e.g. we never codegen the legal case
-        //    cblt w0, #63, A
-        // because we could end up with the illegal case
-        //    cbge w0, #64, B
-        // should the decision to reverse the branch direction be made. For the
-        // lower bound cases this is no problem since we can express comparisons
-        // against 0 with either tbz/tnbz or using wzr/xzr.
-        uint64_t LowerBound = 0, UpperBound = 64;
-        switch (ACC) {
-        case AArch64CC::GE:
-        case AArch64CC::HS:
-        case AArch64CC::LT:
-        case AArch64CC::LO:
-          LowerBound = 1;
-          break;
-        case AArch64CC::LE:
-        case AArch64CC::LS:
-        case AArch64CC::GT:
-        case AArch64CC::HI:
-          UpperBound = 63;
-          break;
-        default:
-          break;
-        }
-
-        if (Imm->getAPIntValue().uge(LowerBound) &&
-            Imm->getAPIntValue().ult(UpperBound))
-          Opc = AArch64ISD::CBRI;
-      }
-
-      SDValue Cond = DAG.getTargetConstant(ACC, dl, MVT::i32);
-      return DAG.getNode(Opc, dl, MVT::Other, Chain, Cond, LHS, RHS, Dest);
+      SDValue Cond =
+          DAG.getTargetConstant(changeIntCCToAArch64CC(CC), dl, MVT::i32);
+      return DAG.getNode(AArch64ISD::CB, dl, MVT::Other, Chain, Cond, LHS, RHS,
+                         Dest);
     }
 
     SDValue CCVal;
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.h b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
@@ -531,8 +531,7 @@ enum NodeType : unsigned {
   SME_ZA_STR,
 
   // Compare-and-branch
-  CBRR,
-  CBRI,
+  CB,
 };
 
 } // end namespace AArch64ISD
diff --git a/llvm/lib/Target/AArch64/AArch64InstrFormats.td b/llvm/lib/Target/AArch64/AArch64InstrFormats.td
@@ -400,6 +400,20 @@ def uimm6_32b : Operand<i32>, ImmLeaf<i32, [{ return Imm >= 0 && Imm < 64; }]> {
    let ParserMatchClass = UImm6Operand;
 }
 
+def CmpBranchUImm6Operand_32b
+    : Operand<i32>,
+      ComplexPattern<i32, 1, "SelectCmpBranchUImm6Operand<32>"> {
+  let ParserMatchClass = UImm6Operand;
+  let WantsParent = true;
+}
+
+def CmpBranchUImm6Operand_64b
+    : Operand<i64>,
+      ComplexPattern<i64, 1, "SelectCmpBranchUImm6Operand<64>"> {
+  let ParserMatchClass = UImm6Operand;
+  let WantsParent = true;
+}
+
 def UImm6Plus1Operand : AsmOperandClass {
   let Name = "UImm6P1";
   let DiagnosticType = "InvalidImm1_64";
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -518,8 +518,9 @@ def SDT_AArch64TBL : SDTypeProfile<1, 2, [
   SDTCisVec<0>, SDTCisSameAs<0, 1>, SDTCisInt<2>
 ]>;
 
-def SDT_AArch64cbrr : SDTypeProfile<0, 4, [SDTCisVT<0, i32>, SDTCisInt<1>, SDTCisSameAs<1, 2>, SDTCisVT<3, OtherVT>]>;
-def SDT_AArch64cbri : SDTypeProfile<0, 4, [SDTCisVT<0, i32>, SDTCisInt<1>, SDTCisInt<2>, SDTCisVT<3, OtherVT>]>;
+def SDT_AArch64cb : SDTypeProfile<0, 4,
+                                  [SDTCisVT<0, i32>, SDTCisInt<1>, SDTCisInt<2>,
+                                   SDTCisVT<3, OtherVT>]>;
 
 // non-extending masked load fragment.
 def nonext_masked_load :
@@ -697,8 +698,7 @@ def topbitsallzero64: PatLeaf<(i64 GPR64:$src), [{
   }]>;
 
 // Node definitions.
-def AArch64CBrr : SDNode<"AArch64ISD::CBRR", SDT_AArch64cbrr, [SDNPHasChain]>;
-def AArch64CBri : SDNode<"AArch64ISD::CBRI", SDT_AArch64cbri, [SDNPHasChain]>;
+def AArch64CB : SDNode<"AArch64ISD::CB", SDT_AArch64cb, [SDNPHasChain]>;
 def AArch64adrp          : SDNode<"AArch64ISD::ADRP", SDTIntUnaryOp, []>;
 def AArch64adr           : SDNode<"AArch64ISD::ADR", SDTIntUnaryOp, []>;
 def AArch64addlow        : SDNode<"AArch64ISD::ADDlow", SDTIntBinOp, []>;
@@ -10552,15 +10552,18 @@ let Predicates = [HasCMPBR] in {
   def CBWPri : CmpBranchImmediatePseudo<GPR32, uimm6_32b>;
   def CBXPri : CmpBranchImmediatePseudo<GPR64, uimm6_64b>;
 
-def : Pat<(AArch64CBrr i32:$Cond, GPR32:$Rn, GPR32:$Rt, bb:$Target),
-          (CBWPrr ccode:$Cond, GPR32:$Rn, GPR32:$Rt, am_brcmpcond:$Target)>;
-def : Pat<(AArch64CBrr i32:$Cond, GPR64:$Rn, GPR64:$Rt, bb:$Target),
-          (CBXPrr ccode:$Cond, GPR64:$Rn, GPR64:$Rt, am_brcmpcond:$Target)>;
-def : Pat<(AArch64CBri i32:$Cond, GPR32:$Rn, i32:$Imm, bb:$Target),
-          (CBWPri ccode:$Cond, GPR32:$Rn, uimm6_32b:$Imm, am_brcmpcond:$Target)>;
-def : Pat<(AArch64CBri i32:$Cond, GPR64:$Rn, i64:$Imm, bb:$Target),
-          (CBXPri ccode:$Cond, GPR64:$Rn, uimm6_64b:$Imm, am_brcmpcond:$Target)>;
-
+  def : Pat<(AArch64CB i32:$Cond, GPR32:$Rn, CmpBranchUImm6Operand_32b:$Imm,
+                bb:$Target),
+            (CBWPri i32:$Cond, GPR32:$Rn, uimm6_32b:$Imm,
+                am_brcmpcond:$Target)>;
+  def : Pat<(AArch64CB i32:$Cond, GPR64:$Rn, CmpBranchUImm6Operand_64b:$Imm,
+                bb:$Target),
+            (CBXPri i32:$Cond, GPR64:$Rn, uimm6_64b:$Imm,
+                am_brcmpcond:$Target)>;
+  def : Pat<(AArch64CB i32:$Cond, GPR32:$Rn, GPR32:$Rt, bb:$Target),
+            (CBWPrr ccode:$Cond, GPR32:$Rn, GPR32:$Rt, am_brcmpcond:$Target)>;
+  def : Pat<(AArch64CB i32:$Cond, GPR64:$Rn, GPR64:$Rt, bb:$Target),
+            (CBXPrr ccode:$Cond, GPR64:$Rn, GPR64:$Rt, am_brcmpcond:$Target)>;
 } // HasCMPBR
 
 
