[RISCV] Fold binary op into select if profitable.

[mgudim]

Summary

Consider the following pattern binOp (select cond, x, c0), c1.
Where c0 and c1 are constants.
We can transform it to select cond, binOp(x, c1), binOp(c0, c1).

If binOp(c0, c1) ends up being 0 or -1 we can turn the select into
a more profitable sequence.

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5527

dyn_cast->cast if it can't fail

5527

const APInt &

[mgudim]

Addressed comments.
Fixed a bug: the constants were passed to DAG.FoldConstantArithmetic in the wrong order.
Moved call to foldBinOpIntoSelectIfProfitable after combineSelectToBinOp - this way we will not fold binop if the constant of select is already 0 or -1.
Minor typo fixes.

[mgudim]

Pre-merge checks passed.

[craig.topper]

LGTM

[mgudim]

Thanks @craig.topper , I'll merge it (once I figure out the process). What about https://reviews.llvm.org/D151750 Can it be merged after this one?

[jrtc27]

This was committed as d0189584631e587279ee5f0af5feb94d8045bb31 without the normal Phab metadata (i.e. without arc patch), but is causing assertion failures building the Linux kernel per https://github.com/llvm/llvm-project/issues/63293.

Please revert this, debug the issue and re-land with an appropriate fix and the correct metadata per the LLVM Phabricator documentation (https://llvm.org/docs/Phabricator.html#committing-a-change).

[craig.topper]

Please revert. I think the issue is that it never checks the opcode of the binary operator which might not even be a binary operator. It might not even be something that FoldConstantArithmetic can handle.

[craig.topper]

I have reverted the change f9d0bf06319203a8cbb47d89c2f39d2c782f3887

[mgudim]

added logic for the correct order of operands when creating new opepands of select.
added test for all four possible positions of constants in select and binop.

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5700

Please upload patches with full context using git diff -U99999

[mgudim]

Uploaded diff with -u99999 for more context.

[mgudim]

Can't reproduce the failed tests locally. Put a print statement in the code to make sure it is the most recent patch being tested.

[mgudim]

(1) Added a check if FoldConstantArithmetic returns null.
(2) Fixed the incorrect use of ArrayRef. This was causing a bug, which is not visible when compiled in Debug mode.

[mgudim]

Changed ArrayRef to SmallVector in two more lines, which I missed in the previous update.

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5818

Use a plain old C array?

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5803

You can use a plain C array and std::swap(NewConstOps[0], NewConstOps[1]) if SelOpNo == 1.

[mgudim]

Addressed comments: use plain array instead of SmallVector

[mgudim]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5803

Makes sence. Done.

[mgudim]

@craig.topper is it ok to recommit this?

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5784

Can ConstSelOp and ConstSelOpNode be moved below the if by using ConstSelOpNo?

5846

Is this safe for division/remainder? The select could be sanitizing inputs. RISC-V doesn't have traps for div/rem, but IR/SelectionDAG semantics consider division by 0 to be immediate UB. Same with (sdiv -1, INT_MIN)

[mgudim]

addressed comments:
(1) Moved definitions of ConstSelOp and ConstSelOpNode lower.
(2) Added check DAG.isSafeToSpeculativelyExecute(...) to check if commuting select and binop is legal.

[mgudim]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

5846

Thanks, I never thought about that.

[craig.topper]

LGTM

[mgudim]

rebased on current main.

[mgudim]

I have accidentally commited this in "Needs Review" state. What should I do?

[craig.topper]

In D152147#4502074, @mgudim wrote:

I have accidentally commited this in "Needs Review" state. What should I do?

It was only in Needs Review due to @jrtc27's. I believe those comments were addressed. The new commit has the phabricator metadata. I think you're ok unless Jessica speaks up.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

if (IsExternWeak) {
           MachineMemOperand::MOLoad | MachineMemOperand::MODereferenceable |
               MachineMemOperand::MOInvariant,
           LLT(Ty.getSimpleVT()), Align(Ty.getFixedSizeInBits() / 8));
       SDValue Load =
           DAG.getMemIntrinsicNode(RISCVISD::LGA, DL,
                                   DAG.getVTList(Ty, MVT::Other),
                                   {DAG.getEntryNode(), Addr}, Ty, MemOp);
       return Load;
     }
 
     // Generate a sequence for accessing addresses within any 2GiB range within
     // the address space. This generates the pattern (PseudoLLA sym), which
     // expands to (addi (auipc %pcrel_hi(sym)) %pcrel_lo(auipc)).
     return DAG.getNode(RISCVISD::LLA, DL, Ty, Addr);
   }
   }
 }
if (LHS == LHS2 && RHS == RHS2) {
     if (CC == CC2)
       return true;
     if (CC == ISD::getSetCCInverse(CC2, LHS2.getValueType()))
       return false;
   }
 
   return std::nullopt;
 }
 
 static SDValue combineSelectToBinOp(SDNode *N, SelectionDAG &DAG,
                                     const RISCVSubtarget &Subtarget) {
   SDValue CondV = N->getOperand(0);
   SDValue TrueV = N->getOperand(1);
   SDValue FalseV = N->getOperand(2);
   MVT VT = N->getSimpleValueType(0);
   SDLoc DL(N);
 
if (std::optional<bool> MatchResult = matchSetCC(LHS, RHS, CC, FalseV)) {
       return DAG.getNode(*MatchResult ? ISD::AND : ISD::OR, DL, VT, TrueV,
                          FalseV);
     }
   }
 
   return SDValue();
 }
 
+// Transform `binOp (select cond, x, c0), c1` where `c0` and `c1` are constants
+// into `select cond, binOp(x, c1), binOp(c0, c1)` if profitable.
+// For now we only consider transformation profitable if `binOp(c0, c1)` ends up
+// being `0` or `-1`. In such cases we can replace `select` with `and`.
+// TODO: Should we also do this if `binOp(c0, c1)` is cheaper to materialize
+// than `c0`?
+static SDValue
+foldBinOpIntoSelectIfProfitable(SDNode *BO, SelectionDAG &DAG,
+                                const RISCVSubtarget &Subtarget) {
+  if (Subtarget.hasShortForwardBranchOpt())
+    return SDValue();
+
+  unsigned SelOpNo = 0;
+  SDValue Sel = BO->getOperand(0);
+  if (Sel.getOpcode() != ISD::SELECT || !Sel.hasOneUse()) {
+    SelOpNo = 1;
+    Sel = BO->getOperand(1);
+  }
+
+  if (Sel.getOpcode() != ISD::SELECT || !Sel.hasOneUse())
+    return SDValue();
+
+  unsigned ConstSelOpNo = 1;
+  unsigned OtherSelOpNo = 2;
+  SDValue ConstSelOp = Sel->getOperand(1);
+  ConstantSDNode *ConstSelOpNode = dyn_cast<ConstantSDNode>(ConstSelOp);
+  if (!ConstSelOpNode) {
+    ConstSelOpNo = 2;
+    OtherSelOpNo = 1;
+    ConstSelOp = Sel->getOperand(2);
+    ConstSelOpNode = dyn_cast<ConstantSDNode>(ConstSelOp);
+  }
+  if (!ConstSelOpNode || ConstSelOpNode->isOpaque())
+    return SDValue();
+
+  SDValue ConstBinOp = BO->getOperand(SelOpNo ^ 1);
+  ConstantSDNode *ConstBinOpNode = dyn_cast<ConstantSDNode>(ConstBinOp);
+  if (!ConstBinOpNode || ConstBinOpNode->isOpaque())
+    return SDValue();
+
+  SDLoc DL(Sel);
+  EVT VT = BO->getValueType(0);
+
+  SDValue NewConstOps[2] = {ConstSelOp, ConstBinOp};
+  if (SelOpNo == 1)
+    std::swap(NewConstOps[0], NewConstOps[1]);
+
+  SDValue NewConstOp =
+      DAG.FoldConstantArithmetic(BO->getOpcode(), DL, VT, NewConstOps);
+  if (!NewConstOp)
+    return SDValue();
+
+  const APInt &NewConstAPInt =
+      cast<ConstantSDNode>(NewConstOp)->getAPIntValue();
+  if (!NewConstAPInt.isZero() && !NewConstAPInt.isAllOnes())
+    return SDValue();
+
+  SDValue OtherSelOp = Sel->getOperand(OtherSelOpNo);
+  SDValue NewNonConstOps[2] = {OtherSelOp, ConstBinOp};
+  if (SelOpNo == 1)
+    std::swap(NewNonConstOps[0], NewNonConstOps[1]);
+  SDValue NewNonConstOp = DAG.getNode(BO->getOpcode(), DL, VT, NewNonConstOps);
+
+  SDValue NewT = (ConstSelOpNo == 1) ? NewConstOp : NewNonConstOp;
+  SDValue NewF = (ConstSelOpNo == 1) ? NewNonConstOp : NewConstOp;
+  return DAG.getSelect(DL, VT, Sel.getOperand(0), NewT, NewF);
+}
+
 SDValue RISCVTargetLowering::lowerSELECT(SDValue Op, SelectionDAG &DAG) const {
   SDValue CondV = Op.getOperand(0);
   SDValue TrueV = Op.getOperand(1);
   SDValue FalseV = Op.getOperand(2);
   SDLoc DL(Op);
   MVT VT = Op.getSimpleValueType();
   MVT XLenVT = Subtarget.getXLenVT();
 
   // Lower vector SELECTs to VSELECTs by splatting the condition.
   if (VT.isVector()) {
     MVT SplatCondVT = VT.changeVectorElementType(MVT::i1);
     SDValue CondSplat = DAG.getSplat(SplatCondVT, DL, CondV);
     return DAG.getNode(ISD::VSELECT, DL, VT, CondSplat, TrueV, FalseV);
   }
 
   if (SDValue V = combineSelectToBinOp(Op.getNode(), DAG, Subtarget))
     return V;
 
+  if (Op.hasOneUse() && isBinOp(Op->use_begin()->getOpcode())) {
+    SDNode *BinOp = *Op->use_begin();
+    if (SDValue NewSel =
+            foldBinOpIntoSelectIfProfitable(*Op->use_begin(), DAG, Subtarget)) {
+      DAG.ReplaceAllUsesWith(BinOp, &NewSel);
+      return lowerSELECT(NewSel, DAG);
+    }
+  }
+
   // (select cc, 1.0, 0.0) -> (sint_to_fp (zext cc))
   // (select cc, 0.0, 1.0) -> (sint_to_fp (zext (xor cc, 1)))
   const ConstantFPSDNode *FPTV = dyn_cast<ConstantFPSDNode>(TrueV);
   const ConstantFPSDNode *FPFV = dyn_cast<ConstantFPSDNode>(FalseV);
   if (FPTV && FPFV) {
     if (FPTV->isExactlyValue(1.0) && FPFV->isExactlyValue(0.0))
       return DAG.getNode(ISD::SINT_TO_FP, DL, VT, CondV);
     if (FPTV->isExactlyValue(0.0) && FPFV->isExactlyValue(1.0)) {

llvm/test/CodeGen/RISCV/fold-binop-into-select.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 2
+; RUN: llc -mtriple=riscv64 < %s  | FileCheck  %s
+
+define i64 @fold_binop_into_select_0(i1 %c, i64 %x) {
+; CHECK-LABEL: fold_binop_into_select_0:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    addi a1, a1, -2
+; CHECK-NEXT:    slli a0, a0, 63
+; CHECK-NEXT:    srai a0, a0, 63
+; CHECK-NEXT:    and a0, a0, a1
+; CHECK-NEXT:    ret
+entry:
+  %select_ = select i1 %c, i64 %x, i64 2
+  %res = sub i64 %select_, 2
+  ret i64 %res
+}
+
+define i64 @fold_binop_into_select_1(i1 %c, i64 %x) {
+; CHECK-LABEL: fold_binop_into_select_1:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    andi a0, a0, 1
+; CHECK-NEXT:    addi a1, a1, -2
+; CHECK-NEXT:    addi a0, a0, -1
+; CHECK-NEXT:    and a0, a0, a1
+; CHECK-NEXT:    ret
+entry:
+  %select_ = select i1 %c, i64 2, i64 %x
+  %res = sub i64 %select_, 2
+  ret i64 %res
+}
+
+define i64 @fold_binop_into_select_2(i1 %c, i64 %x) {
+; CHECK-LABEL: fold_binop_into_select_2:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    li a2, 2
+; CHECK-NEXT:    sub a2, a2, a1
+; CHECK-NEXT:    slli a0, a0, 63
+; CHECK-NEXT:    srai a0, a0, 63
+; CHECK-NEXT:    and a0, a0, a2
+; CHECK-NEXT:    ret
+entry:
+  %select_ = select i1 %c, i64 %x, i64 2
+  %res = sub i64 2, %select_
+  ret i64 %res
+}
+
+define i64 @fold_binop_into_select_3(i1 %c, i64 %x) {
+; CHECK-LABEL: fold_binop_into_select_3:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    andi a0, a0, 1
+; CHECK-NEXT:    li a2, 2
+; CHECK-NEXT:    sub a2, a2, a1
+; CHECK-NEXT:    addi a0, a0, -1
+; CHECK-NEXT:    and a0, a0, a2
+; CHECK-NEXT:    ret
+entry:
+  %select_ = select i1 %c, i64 2, i64 %x
+  %res = sub i64 2, %select_
+  ret i64 %res
+}