[RISCV] Efficiently lower select %cond, %constant, %variable using zicond extension

Provides the following lowerings, which reduces instruction count by 1 for rv64:
(select c, c1, t) -> (add (czero_nez t - c1, c), c1)
(select c, t, c1) -> (add (czero_eqz t - c1, c), c1)

Author: bababuck

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -9043,15 +9043,20 @@ SDValue RISCVTargetLowering::lowerSELECT(SDValue Op, SelectionDAG &DAG) const {
 
     // (select c, c1, c2) -> (add (czero_nez c2 - c1, c), c1)
     // (select c, c1, c2) -> (add (czero_eqz c1 - c2, c), c2)
-    if (isa<ConstantSDNode>(TrueV) && isa<ConstantSDNode>(FalseV)) {
-      const APInt &TrueVal = TrueV->getAsAPIntVal();
-      const APInt &FalseVal = FalseV->getAsAPIntVal();
+    // (select c, c1, t) -> (add (czero_nez t - c1, c), c1)
+    // (select c, t, c1) -> (add (czero_eqz t - c1, c), c1)
+    if (isa<ConstantSDNode>(TrueV) || isa<ConstantSDNode>(FalseV)) {
+      bool BothAreConstant =
+          isa<ConstantSDNode>(TrueV) && isa<ConstantSDNode>(FalseV);
 
       // Prefer these over Zicond to avoid materializing an immediate:
       //   (select (x < 0), y, z)  -> x >> (XLEN - 1) & (y - z) + z
       //   (select (x > -1), z, y) -> x >> (XLEN - 1) & (y - z) + z
       if (CondV.getOpcode() == ISD::SETCC &&
-          CondV.getOperand(0).getValueType() == VT && CondV.hasOneUse()) {
+          CondV.getOperand(0).getValueType() == VT && CondV.hasOneUse() &&
+          BothAreConstant) {
+        const APInt &TrueVal = TrueV->getAsAPIntVal();
+        const APInt &FalseVal = FalseV->getAsAPIntVal();
         ISD::CondCode CCVal = cast<CondCodeSDNode>(CondV.getOperand(2))->get();
         if ((CCVal == ISD::SETLT && isNullConstant(CondV.getOperand(1))) ||
             (CCVal == ISD::SETGT && isAllOnesConstant(CondV.getOperand(1)))) {
@@ -9073,19 +9078,36 @@ SDValue RISCVTargetLowering::lowerSELECT(SDValue Op, SelectionDAG &DAG) const {
         }
       }
 
-      const int TrueValCost = RISCVMatInt::getIntMatCost(
-          TrueVal, Subtarget.getXLen(), Subtarget, /*CompressionCost=*/true);
-      const int FalseValCost = RISCVMatInt::getIntMatCost(
-          FalseVal, Subtarget.getXLen(), Subtarget, /*CompressionCost=*/true);
-      bool IsCZERO_NEZ = TrueValCost <= FalseValCost;
-      SDValue LHSVal = DAG.getConstant(
-          IsCZERO_NEZ ? FalseVal - TrueVal : TrueVal - FalseVal, DL, VT);
-      SDValue RHSVal =
-          DAG.getConstant(IsCZERO_NEZ ? TrueVal : FalseVal, DL, VT);
-      SDValue CMOV =
-          DAG.getNode(IsCZERO_NEZ ? RISCVISD::CZERO_NEZ : RISCVISD::CZERO_EQZ,
-                      DL, VT, LHSVal, CondV);
-      return DAG.getNode(ISD::ADD, DL, VT, CMOV, RHSVal);
+      bool IsCZERO_NEZ;
+      SDValue ConstVal, SubOp;
+      if (BothAreConstant) {
+        const APInt &TrueVal = TrueV->getAsAPIntVal();
+        const APInt &FalseVal = FalseV->getAsAPIntVal();
+        const int TrueValCost = RISCVMatInt::getIntMatCost(
+            TrueVal, Subtarget.getXLen(), Subtarget, /*CompressionCost=*/true);
+        const int FalseValCost = RISCVMatInt::getIntMatCost(
+            FalseVal, Subtarget.getXLen(), Subtarget, /*CompressionCost=*/true);
+        IsCZERO_NEZ = TrueValCost <= FalseValCost;
+
+        SubOp = DAG.getConstant(
+            IsCZERO_NEZ ? FalseVal - TrueVal : TrueVal - FalseVal, DL, VT);
+        ConstVal = DAG.getConstant(IsCZERO_NEZ ? TrueVal : FalseVal, DL, VT);
+      } else {
+        IsCZERO_NEZ = isa<ConstantSDNode>(TrueV);
+        ConstVal = IsCZERO_NEZ ? TrueV : FalseV;
+
+        SDValue &RegV = IsCZERO_NEZ ? FalseV : TrueV;
+        SubOp = DAG.getNode(ISD::SUB, DL, VT, RegV, ConstVal);
+      }
+      // The one constant case is only efficient is the constant fits into
+      // `ADDI`
+      if (BothAreConstant ||
+          isInt<12>(dyn_cast<ConstantSDNode>(ConstVal)->getSExtValue())) {
+        SDValue CMOV =
+            DAG.getNode(IsCZERO_NEZ ? RISCVISD::CZERO_NEZ : RISCVISD::CZERO_EQZ,
+                        DL, VT, SubOp, CondV);
+        return DAG.getNode(ISD::ADD, DL, VT, CMOV, ConstVal);
+      }
     }
 
     // (select c, t, f) -> (or (czero_eqz t, c), (czero_nez f, c))

llvm/test/CodeGen/RISCV/short-forward-branch-opt.ll

@@ -173,14 +173,21 @@ define signext i32 @test6(i32 signext %x, i32 signext %z) {
 ; NOSFB-NEXT:    or a0, a0, a1
 ; NOSFB-NEXT:    ret
 ;
-; SFB-LABEL: test6:
-; SFB:       # %bb.0:
-; SFB-NEXT:    li a2, -1
-; SFB-NEXT:    beqz a1, .LBB5_2
-; SFB-NEXT:  # %bb.1:
-; SFB-NEXT:    mv a0, a2
-; SFB-NEXT:  .LBB5_2:
-; SFB-NEXT:    ret
+; NOZICOND-LABEL: test6:
+; NOZICOND:       # %bb.0:
+; NOZICOND-NEXT:    li a2, -1
+; NOZICOND-NEXT:    beqz a1, .LBB5_2
+; NOZICOND-NEXT:  # %bb.1:
+; NOZICOND-NEXT:    mv a0, a2
+; NOZICOND-NEXT:  .LBB5_2:
+; NOZICOND-NEXT:    ret
+;
+; ZICOND-LABEL: test6:
+; ZICOND:       # %bb.0:
+; ZICOND-NEXT:    addi a0, a0, 1
+; ZICOND-NEXT:    czero.nez a0, a0, a1
+; ZICOND-NEXT:    addi a0, a0, -1
+; ZICOND-NEXT:    ret
   %c = icmp eq i32 %z, 0
   %b = select i1 %c, i32 %x, i32 -1
   ret i32 %b
@@ -195,14 +202,21 @@ define signext i32 @test7(i32 signext %x, i32 signext %z) {
 ; NOSFB-NEXT:    or a0, a0, a1
 ; NOSFB-NEXT:    ret
 ;
-; SFB-LABEL: test7:
-; SFB:       # %bb.0:
-; SFB-NEXT:    li a2, -1
-; SFB-NEXT:    bnez a1, .LBB6_2
-; SFB-NEXT:  # %bb.1:
-; SFB-NEXT:    mv a0, a2
-; SFB-NEXT:  .LBB6_2:
-; SFB-NEXT:    ret
+; NOZICOND-LABEL: test7:
+; NOZICOND:       # %bb.0:
+; NOZICOND-NEXT:    li a2, -1
+; NOZICOND-NEXT:    bnez a1, .LBB6_2
+; NOZICOND-NEXT:  # %bb.1:
+; NOZICOND-NEXT:    mv a0, a2
+; NOZICOND-NEXT:  .LBB6_2:
+; NOZICOND-NEXT:    ret
+;
+; ZICOND-LABEL: test7:
+; ZICOND:       # %bb.0:
+; ZICOND-NEXT:    addi a0, a0, 1
+; ZICOND-NEXT:    czero.eqz a0, a0, a1
+; ZICOND-NEXT:    addi a0, a0, -1
+; ZICOND-NEXT:    ret
   %c = icmp eq i32 %z, 0
   %b = select i1 %c, i32 -1, i32 %x
   ret i32 %b

llvm/test/CodeGen/RISCV/zicond-opts.ll

@@ -100,21 +100,19 @@ define dso_local noundef i64 @select_imm_reg(i64 noundef %0, i64 noundef %1) loc
 ; RV32ZICOND-LABEL: select_imm_reg:
 ; RV32ZICOND:       # %bb.0:
 ; RV32ZICOND-NEXT:    xori a0, a0, 2
-; RV32ZICOND-NEXT:    or a1, a0, a1
-; RV32ZICOND-NEXT:    li a0, 3
-; RV32ZICOND-NEXT:    czero.eqz a2, a2, a1
-; RV32ZICOND-NEXT:    czero.nez a0, a0, a1
-; RV32ZICOND-NEXT:    or a0, a0, a2
-; RV32ZICOND-NEXT:    czero.eqz a1, a3, a1
+; RV32ZICOND-NEXT:    addi a2, a2, -3
+; RV32ZICOND-NEXT:    or a0, a0, a1
+; RV32ZICOND-NEXT:    czero.eqz a1, a3, a0
+; RV32ZICOND-NEXT:    czero.eqz a0, a2, a0
+; RV32ZICOND-NEXT:    addi a0, a0, 3
 ; RV32ZICOND-NEXT:    ret
 ;
 ; RV64ZICOND-LABEL: select_imm_reg:
 ; RV64ZICOND:       # %bb.0:
 ; RV64ZICOND-NEXT:    addi a0, a0, -2
-; RV64ZICOND-NEXT:    li a2, 3
-; RV64ZICOND-NEXT:    czero.eqz a1, a1, a0
-; RV64ZICOND-NEXT:    czero.nez a0, a2, a0
-; RV64ZICOND-NEXT:    or a0, a0, a1
+; RV64ZICOND-NEXT:    addi a1, a1, -3
+; RV64ZICOND-NEXT:    czero.eqz a0, a1, a0
+; RV64ZICOND-NEXT:    addi a0, a0, 3
 ; RV64ZICOND-NEXT:    ret
   %3 = icmp eq i64 %0, 2
   %4 = select i1 %3, i64 3, i64 %1