[RISCV] Fold (add (select lhs, rhs, cc, 0, y), x) -> (select lhs, rhs, cc, x, (add x, y))

[craig.topper]

Summary

Similar for sub except sub isn't commutative.

Modify the existing and/or/xor folds to also work on ISD::SELECT
and not just RISCVISD::SELECT_CC. This is needed to make sure
we do this transform before type legalization turns i32 add/sub
into add/sub+sign_extend_inreg on RV64. If we don't do this before
that, the sign_extend_inreg will still be after the select.

[frasercrmck]

A general observation is that expanding this to ISD::SELECT now opens this up to vector types. I think the only thing that's now preventing it is isZeroOrAllOnes which "silently" (for want of a better word) only accepts scalar constants.

My first thoughts are that we could either be clearer about not accepting vector types, or - without wanting to expand the scope of this patch - we do open it up to vectors. From a quick play around on top of your patch it doesn't look too useful in that case, though. A vector add_select_all_zeros_v4i32 just introduces an extra vsetvli.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5818

"the" condition, perhaps

5846

I know you didn't introduce this but /*AllOnes*/ might be a good change to make while you're here.

[craig.topper]

Address review comments

[craig.topper]

In D107603#2930673, @frasercrmck wrote:

A general observation is that expanding this to ISD::SELECT now opens this up to vector types. I think the only thing that's now preventing it is isZeroOrAllOnes which "silently" (for want of a better word) only accepts scalar constants.

My first thoughts are that we could either be clearer about not accepting vector types, or - without wanting to expand the scope of this patch - we do open it up to vectors. From a quick play around on top of your patch it doesn't look too useful in that case, though. A vector add_select_all_zeros_v4i32 just introduces an extra vsetvli.

I disabled it for vectors.

[frasercrmck]

LGTM, bearing in mind there are others with more knowledge about the "scalar" side of this target

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

RISCVTargetLowering::RISCVTargetLowering(const TargetMachine &TM,
   setMinimumJumpTableEntries(5);
 
   // Jumps are expensive, compared to logic
   setJumpIsExpensive();
 
   // We can use any register for comparisons
   setHasMultipleConditionRegisters();
 
+  setTargetDAGCombine(ISD::ADD);
+  setTargetDAGCombine(ISD::SUB);
   setTargetDAGCombine(ISD::AND);
   setTargetDAGCombine(ISD::OR);
   setTargetDAGCombine(ISD::XOR);
   setTargetDAGCombine(ISD::ANY_EXTEND);
   setTargetDAGCombine(ISD::ZERO_EXTEND);
   if (Subtarget.hasStdExtV()) {
     setTargetDAGCombine(ISD::FCOPYSIGN);
     setTargetDAGCombine(ISD::MGATHER);
static SDValue combineGREVI_GORCI(SDNode *N, SelectionDAG &DAG) {
   SDLoc DL(N);
   return DAG.getNode(
       N->getOpcode(), DL, N->getValueType(0), Src,
       DAG.getConstant(CombinedShAmt, DL, N->getOperand(1).getValueType()));
 }
 
 // Combine a constant select operand into its use:
 //
-// (and (select_cc lhs, rhs, cc, -1, c), x)
-//   -> (select_cc lhs, rhs, cc, x, (and, x, c))  [AllOnes=1]
-// (or  (select_cc lhs, rhs, cc, 0, c), x)
-//   -> (select_cc lhs, rhs, cc, x, (or, x, c))  [AllOnes=0]
-// (xor (select_cc lhs, rhs, cc, 0, c), x)
-//   -> (select_cc lhs, rhs, cc, x, (xor, x, c))  [AllOnes=0]
-static SDValue combineSelectCCAndUse(SDNode *N, SDValue Slct, SDValue OtherOp,
-                                     SelectionDAG &DAG, bool AllOnes) {
+// (and (select cond, -1, c), x)
+//   -> (select cond, x, (and x, c))  [AllOnes=1]
+// (or  (select cond, 0, c), x)
+//   -> (select cond, x, (or x, c))  [AllOnes=0]
+// (xor (select cond, 0, c), x)
+//   -> (select cond, x, (xor x, c))  [AllOnes=0]
+// (add (select cond, 0, c), x)
+//   -> (select cond, x, (add x, c))  [AllOnes=0]
+// (sub x, (select cond, 0, c))
+//   -> (select cond, x, (sub x, c))  [AllOnes=0]
+static SDValue combineSelectAndUse(SDNode *N, SDValue Slct, SDValue OtherOp,
   EVT VT = N->getValueType(0);
 
-  if (Slct.getOpcode() != RISCVISD::SELECT_CC || !Slct.hasOneUse())
+  // Skip vectors.
+  if (VT.isVector())
+    return SDValue();
+
+  if ((Slct.getOpcode() != ISD::SELECT &&
+       Slct.getOpcode() != RISCVISD::SELECT_CC) ||
+      !Slct.hasOneUse())
     return SDValue();
 
   auto isZeroOrAllOnes = [](SDValue N, bool AllOnes) {
     return AllOnes ? isAllOnesConstant(N) : isNullConstant(N);
   };
 
   bool SwapSelectOps;
-  SDValue TrueVal = Slct.getOperand(3);
-  SDValue FalseVal = Slct.getOperand(4);
+  unsigned OpOffset = Slct.getOpcode() == RISCVISD::SELECT_CC ? 2 : 0;
+  SDValue TrueVal = Slct.getOperand(1 + OpOffset);
+  SDValue FalseVal = Slct.getOperand(2 + OpOffset);
   SDValue NonConstantVal;
   if (isZeroOrAllOnes(TrueVal, AllOnes)) {
     SwapSelectOps = false;
     NonConstantVal = FalseVal;
   } else if (isZeroOrAllOnes(FalseVal, AllOnes)) {
     SwapSelectOps = true;
     NonConstantVal = TrueVal;
   } else
     return SDValue();
 
   // Slct is now know to be the desired identity constant when CC is true.
   TrueVal = OtherOp;
   FalseVal = DAG.getNode(N->getOpcode(), SDLoc(N), VT, OtherOp, NonConstantVal);
-  // Unless SwapSelectOps says CC should be false.
+  // Unless SwapSelectOps says the condition should be false.
   if (SwapSelectOps)
     std::swap(TrueVal, FalseVal);
 
+  if (Slct.getOpcode() == RISCVISD::SELECT_CC)
     return DAG.getNode(RISCVISD::SELECT_CC, SDLoc(N), VT,
                        {Slct.getOperand(0), Slct.getOperand(1),
                         Slct.getOperand(2), TrueVal, FalseVal});
+
+  return DAG.getNode(ISD::SELECT, SDLoc(N), VT,
+                     {Slct.getOperand(0), TrueVal, FalseVal});
 }
 
 // Attempt combineSelectAndUse on each operand of a commutative operator N.
-static SDValue combineSelectCCAndUseCommutative(SDNode *N, SelectionDAG &DAG,
-                                                bool AllOnes) {
+static SDValue combineSelectAndUseCommutative(SDNode *N, SelectionDAG &DAG,
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
-  if (SDValue Result = combineSelectCCAndUse(N, N0, N1, DAG, AllOnes))
+  if (SDValue Result = combineSelectAndUse(N, N0, N1, DAG, AllOnes))
     return Result;
-  if (SDValue Result = combineSelectCCAndUse(N, N1, N0, DAG, AllOnes))
+  if (SDValue Result = combineSelectAndUse(N, N1, N0, DAG, AllOnes))
     return Result;
   return SDValue();
 }
 
-static SDValue performANDCombine(SDNode *N,
-                                 TargetLowering::DAGCombinerInfo &DCI,
-                                 const RISCVSubtarget &Subtarget) {
-  SelectionDAG &DAG = DCI.DAG;
+static SDValue performADDCombine(SDNode *N, SelectionDAG &DAG) {
+  // fold (add (select lhs, rhs, cc, 0, y), x) ->
+  //      (select lhs, rhs, cc, x, (add x, y))
+  return combineSelectAndUseCommutative(N, DAG, /*AllOnes*/ false);
+}
 
-  // fold (and (select_cc lhs, rhs, cc, -1, y), x) ->
+static SDValue performSUBCombine(SDNode *N, SelectionDAG &DAG) {
+  // fold (sub x, (select lhs, rhs, cc, 0, y)) ->
+  //      (select lhs, rhs, cc, x, (sub x, y))
+  SDValue N0 = N->getOperand(0);
+  SDValue N1 = N->getOperand(1);
+  return combineSelectAndUse(N, N1, N0, DAG, /*AllOnes*/ false);
+}
+
+static SDValue performANDCombine(SDNode *N, SelectionDAG &DAG) {
+  // fold (and (select lhs, rhs, cc, -1, y), x) ->
   //      (select lhs, rhs, cc, x, (and x, y))
-  return combineSelectCCAndUseCommutative(N, DAG, true);
+  return combineSelectAndUseCommutative(N, DAG, /*AllOnes*/ true);
 }
 
-static SDValue performORCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI,
+static SDValue performORCombine(SDNode *N, SelectionDAG &DAG,
                                 const RISCVSubtarget &Subtarget) {
-  SelectionDAG &DAG = DCI.DAG;
   if (Subtarget.hasStdExtZbp()) {
     if (auto GREV = combineORToGREV(SDValue(N, 0), DAG, Subtarget))
       return GREV;
     if (auto GORC = combineORToGORC(SDValue(N, 0), DAG, Subtarget))
       return GORC;
     if (auto SHFL = combineORToSHFL(SDValue(N, 0), DAG, Subtarget))
       return SHFL;
   }
 
-  // fold (or (select_cc lhs, rhs, cc, 0, y), x) ->
-  //      (select lhs, rhs, cc, x, (or x, y))
-  return combineSelectCCAndUseCommutative(N, DAG, false);
+  // fold (or (select cond, 0, y), x) ->
+  //      (select cond, x, (or x, y))
+  return combineSelectAndUseCommutative(N, DAG, /*AllOnes*/ false);
 }
 
-static SDValue performXORCombine(SDNode *N,
-                                 TargetLowering::DAGCombinerInfo &DCI,
-                                 const RISCVSubtarget &Subtarget) {
-  SelectionDAG &DAG = DCI.DAG;
-
-  // fold (xor (select_cc lhs, rhs, cc, 0, y), x) ->
-  //      (select lhs, rhs, cc, x, (xor x, y))
-  return combineSelectCCAndUseCommutative(N, DAG, false);
+static SDValue performXORCombine(SDNode *N, SelectionDAG &DAG) {
+  // fold (xor (select cond, 0, y), x) ->
+  //      (select cond, x, (xor x, y))
+  return combineSelectAndUseCommutative(N, DAG, /*AllOnes*/ false);
 }
 
 // Attempt to turn ANY_EXTEND into SIGN_EXTEND if the input to the ANY_EXTEND
 // has users that require SIGN_EXTEND and the SIGN_EXTEND can be done for free
 // by an instruction like ADDW/SUBW/MULW. Without this the ANY_EXTEND would be
 // removed during type legalization leaving an ADD/SUB/MUL use that won't use
 // ADDW/SUBW/MULW.
 static SDValue performANY_EXTENDCombine(SDNode *N,
case RISCVISD::FMV_X_ANYEXTW_RV64: {
     if (Op0.getOpcode() == ISD::FNEG)
       return DAG.getNode(ISD::XOR, DL, VT, NewFMV,
                          DAG.getConstant(SignBit, DL, VT));
 
     assert(Op0.getOpcode() == ISD::FABS);
     return DAG.getNode(ISD::AND, DL, VT, NewFMV,
                        DAG.getConstant(~SignBit, DL, VT));
   }
+  case ISD::ADD:
+    return performADDCombine(N, DAG);
+  case ISD::SUB:
+    return performSUBCombine(N, DAG);
   case ISD::AND:
-    return performANDCombine(N, DCI, Subtarget);
+    return performANDCombine(N, DAG);
   case ISD::OR:
-    return performORCombine(N, DCI, Subtarget);
+    return performORCombine(N, DAG, Subtarget);
   case ISD::XOR:
-    return performXORCombine(N, DCI, Subtarget);
+    return performXORCombine(N, DAG);
   case ISD::ANY_EXTEND:
     return performANY_EXTENDCombine(N, DCI, Subtarget);
   case ISD::ZERO_EXTEND:
     // Fold (zero_extend (fp_to_uint X)) to prevent forming fcvt+zexti32 during
     // type legalization. This is safe because fp_to_uint produces poison if
     // it overflows.
     if (N->getValueType(0) == MVT::i64 && Subtarget.is64Bit() &&
         N->getOperand(0).getOpcode() == ISD::FP_TO_UINT &&

llvm/test/CodeGen/RISCV/rv32zbs.ll

 ; RV32I-NEXT:  .LBB2_2:
 ; RV32I-NEXT:    sll a2, a3, a2
 ; RV32I-NEXT:    not a2, a2
 ; RV32I-NEXT:    and a0, a0, a2
 ; RV32I-NEXT:    ret
 ;
 ; RV32IB-LABEL: sbclr_i64:
 ; RV32IB:       # %bb.0:
-; RV32IB-NEXT:    andi a3, a2, 63
-; RV32IB-NEXT:    addi a3, a3, -32
-; RV32IB-NEXT:    bset a4, zero, a3
-; RV32IB-NEXT:    slti a5, a3, 0
-; RV32IB-NEXT:    cmov a4, a5, zero, a4
-; RV32IB-NEXT:    bset a2, zero, a2
-; RV32IB-NEXT:    srai a3, a3, 31
-; RV32IB-NEXT:    and a2, a3, a2
-; RV32IB-NEXT:    andn a1, a1, a4
-; RV32IB-NEXT:    andn a0, a0, a2
+; RV32IB-NEXT:    bset a3, zero, a2
+; RV32IB-NEXT:    andi a2, a2, 63
+; RV32IB-NEXT:    addi a2, a2, -32
+; RV32IB-NEXT:    srai a4, a2, 31
+; RV32IB-NEXT:    and a3, a4, a3
+; RV32IB-NEXT:    slti a4, a2, 0
+; RV32IB-NEXT:    bclr a2, a1, a2
+; RV32IB-NEXT:    cmov a1, a4, a1, a2
+; RV32IB-NEXT:    andn a0, a0, a3
 ; RV32IB-NEXT:    ret
 ;
 ; RV32IBS-LABEL: sbclr_i64:
 ; RV32IBS:       # %bb.0:
 ; RV32IBS-NEXT:    andi a3, a2, 63
 ; RV32IBS-NEXT:    addi a3, a3, -32
 ; RV32IBS-NEXT:    bltz a3, .LBB2_2
 ; RV32IBS-NEXT:  # %bb.1:

llvm/test/CodeGen/RISCV/select-binop-identity.ll

; RV64I-NEXT:    ret
   %a = select i1 %c, i64 %x, i64 0
   %b = xor i64 %a, %y
   ret i64 %b
 }
 
 define signext i32 @add_select_all_zeros_i32(i1 zeroext %c, i32 signext %x, i32 signext %y) {
 ; RV32I-LABEL: add_select_all_zeros_i32:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    mv a3, zero
 ; RV32I-NEXT:    bnez a0, .LBB6_2
 ; RV32I-NEXT:  # %bb.1:
-; RV32I-NEXT:    mv a3, a1
+; RV32I-NEXT:    add a2, a2, a1
 ; RV32I-NEXT:  .LBB6_2:
-; RV32I-NEXT:    add a0, a2, a3
+; RV32I-NEXT:    mv a0, a2
 ; RV32I-NEXT:    ret
 ;
 ; RV64I-LABEL: add_select_all_zeros_i32:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    mv a3, zero
 ; RV64I-NEXT:    bnez a0, .LBB6_2
 ; RV64I-NEXT:  # %bb.1:
-; RV64I-NEXT:    mv a3, a1
+; RV64I-NEXT:    addw a2, a2, a1
 ; RV64I-NEXT:  .LBB6_2:
-; RV64I-NEXT:    addw a0, a2, a3
+; RV64I-NEXT:    mv a0, a2
 ; RV64I-NEXT:    ret
   %a = select i1 %c, i32 0, i32 %x
   %b = add i32 %y, %a
   ret i32 %b
 }
 
 define i64 @add_select_all_zeros_i64(i1 zeroext %c, i64 %x, i64 %y) {
 ; RV32I-LABEL: add_select_all_zeros_i64:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    bnez a0, .LBB7_2
+; RV32I-NEXT:    beqz a0, .LBB7_2
 ; RV32I-NEXT:  # %bb.1:
-; RV32I-NEXT:    mv a2, zero
-; RV32I-NEXT:    mv a1, zero
+; RV32I-NEXT:    add a0, a4, a2
+; RV32I-NEXT:    add a1, a3, a1
+; RV32I-NEXT:    sltu a2, a1, a3
+; RV32I-NEXT:    add a4, a0, a2
+; RV32I-NEXT:    mv a3, a1
 ; RV32I-NEXT:  .LBB7_2:
-; RV32I-NEXT:    add a0, a1, a3
-; RV32I-NEXT:    sltu a1, a0, a1
-; RV32I-NEXT:    add a2, a2, a4
-; RV32I-NEXT:    add a1, a2, a1
+; RV32I-NEXT:    mv a0, a3
+; RV32I-NEXT:    mv a1, a4
 ; RV32I-NEXT:    ret
 ;
 ; RV64I-LABEL: add_select_all_zeros_i64:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    bnez a0, .LBB7_2
+; RV64I-NEXT:    beqz a0, .LBB7_2
 ; RV64I-NEXT:  # %bb.1:
-; RV64I-NEXT:    mv a1, zero
+; RV64I-NEXT:    add a2, a2, a1
 ; RV64I-NEXT:  .LBB7_2:
-; RV64I-NEXT:    add a0, a1, a2
+; RV64I-NEXT:    mv a0, a2
 ; RV64I-NEXT:    ret
   %a = select i1 %c, i64 %x, i64 0
   %b = add i64 %a, %y
   ret i64 %b
 }
 
 define signext i32 @sub_select_all_zeros_i32(i1 zeroext %c, i32 signext %x, i32 signext %y) {
 ; RV32I-LABEL: sub_select_all_zeros_i32:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    mv a3, zero
 ; RV32I-NEXT:    bnez a0, .LBB8_2
 ; RV32I-NEXT:  # %bb.1:
-; RV32I-NEXT:    mv a3, a1
+; RV32I-NEXT:    sub a2, a2, a1
 ; RV32I-NEXT:  .LBB8_2:
-; RV32I-NEXT:    sub a0, a2, a3
+; RV32I-NEXT:    mv a0, a2
 ; RV32I-NEXT:    ret
 ;
 ; RV64I-LABEL: sub_select_all_zeros_i32:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    mv a3, zero
 ; RV64I-NEXT:    bnez a0, .LBB8_2
 ; RV64I-NEXT:  # %bb.1:
-; RV64I-NEXT:    mv a3, a1
+; RV64I-NEXT:    subw a2, a2, a1
 ; RV64I-NEXT:  .LBB8_2:
-; RV64I-NEXT:    subw a0, a2, a3
+; RV64I-NEXT:    mv a0, a2
 ; RV64I-NEXT:    ret
   %a = select i1 %c, i32 0, i32 %x
   %b = sub i32 %y, %a
   ret i32 %b
 }
 
 define i64 @sub_select_all_zeros_i64(i1 zeroext %c, i64 %x, i64 %y) {
 ; RV32I-LABEL: sub_select_all_zeros_i64:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    bnez a0, .LBB9_2
+; RV32I-NEXT:    beqz a0, .LBB9_2
 ; RV32I-NEXT:  # %bb.1:
-; RV32I-NEXT:    mv a2, zero
-; RV32I-NEXT:    mv a1, zero
-; RV32I-NEXT:  .LBB9_2:
 ; RV32I-NEXT:    sltu a0, a3, a1
 ; RV32I-NEXT:    sub a2, a4, a2
-; RV32I-NEXT:    sub a2, a2, a0
-; RV32I-NEXT:    sub a0, a3, a1
-; RV32I-NEXT:    mv a1, a2
+; RV32I-NEXT:    sub a4, a2, a0
+; RV32I-NEXT:    sub a3, a3, a1
+; RV32I-NEXT:  .LBB9_2:
+; RV32I-NEXT:    mv a0, a3
+; RV32I-NEXT:    mv a1, a4
 ; RV32I-NEXT:    ret
 ;
 ; RV64I-LABEL: sub_select_all_zeros_i64:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    bnez a0, .LBB9_2
+; RV64I-NEXT:    beqz a0, .LBB9_2
 ; RV64I-NEXT:  # %bb.1:
-; RV64I-NEXT:    mv a1, zero
+; RV64I-NEXT:    sub a2, a2, a1
 ; RV64I-NEXT:  .LBB9_2:
-; RV64I-NEXT:    sub a0, a2, a1
+; RV64I-NEXT:    mv a0, a2
 ; RV64I-NEXT:    ret
   %a = select i1 %c, i64 %x, i64 0
   %b = sub i64 %y, %a
   ret i64 %b
 }