diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp --- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp +++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp @@ -5819,6 +5819,72 @@ return true; } +// 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; + if (!dyn_cast(Sel->getOperand(ConstSelOpNo))) { + ConstSelOpNo = 2; + OtherSelOpNo = 1; + } + SDValue ConstSelOp = Sel->getOperand(ConstSelOpNo); + ConstantSDNode *ConstSelOpNode = dyn_cast(ConstSelOp); + if (!ConstSelOpNode || ConstSelOpNode->isOpaque()) + return SDValue(); + + SDValue ConstBinOp = BO->getOperand(SelOpNo ^ 1); + ConstantSDNode *ConstBinOpNode = dyn_cast(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(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); @@ -5899,6 +5965,18 @@ if (SDValue V = combineSelectToBinOp(Op.getNode(), DAG, Subtarget)) return V; + if (Op.hasOneUse()) { + unsigned UseOpc = Op->use_begin()->getOpcode(); + if (isBinOp(UseOpc) && DAG.isSafeToSpeculativelyExecute(UseOpc)) { + 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(TrueV); diff --git a/llvm/test/CodeGen/RISCV/fold-binop-into-select.ll b/llvm/test/CodeGen/RISCV/fold-binop-into-select.ll new file mode 100644 --- /dev/null +++ b/llvm/test/CodeGen/RISCV/fold-binop-into-select.ll @@ -0,0 +1,60 @@ +; 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 +}