[InstCombine] Lower flag check pattern to use a bitmask-shift

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -16736,6 +16736,52 @@ static SDValue combineOrAndToBitfieldInsert(SDNode *N, SelectionDAG &DAG,
   return DAG.getNode(RISCVISD::QC_INSB, DL, MVT::i32, Ops);
 }
 
+// or (icmp eq x, imm0), (icmp eq x, imm1) -> czero.eqz (sltui x, 64), (bext x,
+// 1 << imm0 | 1 << imm1) If [imm0, imm1] < 64
+static SDValue combineOrOfImmCmpToBitExtract(SDNode *N, SelectionDAG &DAG,
+                                             const RISCVSubtarget &Subtarget) {
+  using namespace SDPatternMatch;
+
+  auto CollectSetEqImmTree = [](auto &&Self, SmallVector<APInt, 4> &FlagVals,
+                                SDNode *N, SDValue &X) -> bool {
+    APInt Imm;
+    if (X ? sd_match(N, m_OneUse(m_SetCC(m_Specific(X), m_ConstInt(Imm),
+                                         m_SpecificCondCode(ISD::SETEQ))))
+          : sd_match(N, m_OneUse(m_SetCC(m_Value(X), m_ConstInt(Imm),
+                                         m_SpecificCondCode(ISD::SETEQ))))) {
+      FlagVals.push_back(Imm);
+      return true;
+    }
+    SDValue LHS, RHS;
+    if (sd_match(N, m_OneUse(m_Or(m_Value(LHS), m_Value(RHS))))) {
+      return Self(Self, FlagVals, LHS.getNode(), X) &&
+             Self(Self, FlagVals, RHS.getNode(), X);
+    }
+    return false;
+  };
+
+  SmallVector<APInt, 4> FlagVals;
+  SDValue X;
+  if (!CollectSetEqImmTree(CollectSetEqImmTree, FlagVals, N, X))
+    return SDValue();
+
+  unsigned XLen = Subtarget.getXLen();
+  uint64_t BitMask = 0;
+  for (auto &Imm : FlagVals) {
+    if (Imm.uge(XLen))
+      return SDValue();
+    BitMask |= ((uint64_t)1 << Imm.getZExtValue());
+  }
+
+  SDLoc DL(N);
+  EVT VT = N->getValueType(0);
+  SDValue BitExtract =
+      DAG.getNode(ISD::SRL, DL, VT, DAG.getConstant(BitMask, DL, VT), X);
+  SDValue Lt64Check =
+      DAG.getSetCC(DL, VT, X, DAG.getConstant(XLen, DL, VT), ISD::SETULT);
+  return DAG.getNode(ISD::AND, DL, VT, Lt64Check, BitExtract);
+}
+
 static SDValue performORCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI,
                                 const RISCVSubtarget &Subtarget) {
   SelectionDAG &DAG = DCI.DAG;
@@ -16748,6 +16794,9 @@ static SDValue performORCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI,
     return V;
   if (SDValue V = combineBinOpOfExtractToReduceTree(N, DAG, Subtarget))
     return V;
+  if (DCI.isAfterLegalizeDAG())
+    if (SDValue V = combineOrOfImmCmpToBitExtract(N, DAG, Subtarget))
+      return V;
 
   if (DCI.isAfterLegalizeDAG())
     if (SDValue V = combineDeMorganOfBoolean(N, DAG))

llvm/test/CodeGen/RISCV/flag_check.ll

@@ -0,0 +1,241 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc -mtriple=riscv32 -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,RV32
+; RUN: llc -mtriple=riscv64 -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,RV64
+
+define i1 @or_icmp_2(i32 signext %type) {
+; RV32-LABEL: or_icmp_2:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    li a1, 65
+; RV32-NEXT:    srl a1, a1, a0
+; RV32-NEXT:    sltiu a0, a0, 32
+; RV32-NEXT:    and a0, a0, a1
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: or_icmp_2:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    li a1, 65
+; RV64-NEXT:    srl a1, a1, a0
+; RV64-NEXT:    sltiu a0, a0, 64
+; RV64-NEXT:    and a0, a0, a1
+; RV64-NEXT:    ret
+entry:
+  %cmp = icmp eq i32 %type, 6
+  %cmp1 = icmp eq i32 %type, 0
+  %or.cond = or i1 %cmp, %cmp1
+  ret i1 %or.cond
+}
+
+define i1 @or_icmp_3(i32 signext %type) {
+; RV32-LABEL: or_icmp_3:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    lui a1, 8
+; RV32-NEXT:    addi a1, a1, 65
+; RV32-NEXT:    srl a1, a1, a0
+; RV32-NEXT:    sltiu a0, a0, 32
+; RV32-NEXT:    and a0, a0, a1
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: or_icmp_3:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    lui a1, 8
+; RV64-NEXT:    addi a1, a1, 65
+; RV64-NEXT:    srl a1, a1, a0
+; RV64-NEXT:    sltiu a0, a0, 64
+; RV64-NEXT:    and a0, a0, a1
+; RV64-NEXT:    ret
+entry:
+  %cmp = icmp eq i32 %type, 6
+  %cmp1 = icmp eq i32 %type, 0
+  %or.cond = or i1 %cmp, %cmp1
+  %cmp3 = icmp eq i32 %type, 15
+  %or.cond1 = or i1 %cmp3, %or.cond
+  ret i1 %or.cond1
+}
+
+define i1 @or_icmp_4_tree(i32 signext %type) {
+; RV32-LABEL: or_icmp_4_tree:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    lui a1, 1032
+; RV32-NEXT:    addi a1, a1, 65
+; RV32-NEXT:    srl a1, a1, a0
+; RV32-NEXT:    sltiu a0, a0, 32
+; RV32-NEXT:    and a0, a0, a1
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: or_icmp_4_tree:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    lui a1, 1032
+; RV64-NEXT:    addi a1, a1, 65
+; RV64-NEXT:    srl a1, a1, a0
+; RV64-NEXT:    sltiu a0, a0, 64
+; RV64-NEXT:    and a0, a0, a1
+; RV64-NEXT:    ret
+entry:
+  %cmp = icmp eq i32 %type, 6
+  %cmp1 = icmp eq i32 %type, 0
+  %or.cond = or i1 %cmp, %cmp1
+  %cmp2 = icmp eq i32 %type, 15
+  %cmp3 = icmp eq i32 %type, 22
+  %or.cond1 = or i1 %cmp2, %cmp3
+  %or.cond2 = or i1 %or.cond1, %or.cond
+  ret i1 %or.cond2
+}
+
+define i1 @or_icmp_7(i32 signext %type) {
+; RV32-LABEL: or_icmp_7:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    lui a1, 589860
+; RV32-NEXT:    addi a1, a1, 73
+; RV32-NEXT:    srl a1, a1, a0
+; RV32-NEXT:    sltiu a0, a0, 32
+; RV32-NEXT:    and a0, a0, a1
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: or_icmp_7:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    lui a1, 147465
+; RV64-NEXT:    slli a1, a1, 2
+; RV64-NEXT:    addi a1, a1, 73
+; RV64-NEXT:    srl a1, a1, a0
+; RV64-NEXT:    sltiu a0, a0, 64
+; RV64-NEXT:    and a0, a0, a1
+; RV64-NEXT:    ret
+entry:
+  %cmp = icmp eq i32 %type, 6
+  %cmp1 = icmp eq i32 %type, 0
+  %or.cond = or i1 %cmp, %cmp1
+  %cmp2 = icmp eq i32 %type, 17
+  %or.cond1 = or i1 %cmp2, %or.cond
+  %cmp3 = icmp eq i32 %type, 3
+  %or.cond2 = or i1 %cmp3, %or.cond1
+  %cmp4 = icmp eq i32 %type, 31
+  %or.cond3 = or i1 %cmp4, %or.cond2
+  %cmp5 = icmp eq i32 %type, 14
+  %or.cond4 = or i1 %cmp5, %or.cond3
+  %cmp6 = icmp eq i32 %type, 28
+  %or.cond5 = or i1 %cmp6, %or.cond4
+  ret i1 %or.cond5
+}
+
+define i1 @or_icmp_gte_64(i32 signext %type) {
+; CHECK-LABEL: or_icmp_gte_64:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    addi a1, a0, -6
+; CHECK-NEXT:    addi a0, a0, -64
+; CHECK-NEXT:    seqz a1, a1
+; CHECK-NEXT:    seqz a0, a0
+; CHECK-NEXT:    or a0, a1, a0
+; CHECK-NEXT:    ret
+entry:
+  %cmp = icmp eq i32 %type, 6
+  %cmp1 = icmp eq i32 %type, 64
+  %or.cond = or i1 %cmp, %cmp1
+  ret i1 %or.cond
+}
+
+define i1 @or_icmp_multiple_uses(i32 signext %type) {
+; CHECK-LABEL: or_icmp_multiple_uses:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    addi a1, a0, -6
+; CHECK-NEXT:    seqz a1, a1
+; CHECK-NEXT:    seqz a0, a0
+; CHECK-NEXT:    or a0, a1, a0
+; CHECK-NEXT:    xor a0, a1, a0
+; CHECK-NEXT:    ret
+entry:
+  %cmp = icmp eq i32 %type, 6
+  %cmp1 = icmp eq i32 %type, 0
+  %or.cond = or i1 %cmp, %cmp1
+  %or.cond1 = xor i1 %cmp, %or.cond
+  ret i1 %or.cond1
+}
+
+
+define i1 @or_icmp_not_eq(i32 signext %type) {
+; CHECK-LABEL: or_icmp_not_eq:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    addi a1, a0, -6
+; CHECK-NEXT:    seqz a1, a1
+; CHECK-NEXT:    snez a0, a0
+; CHECK-NEXT:    or a0, a1, a0
+; CHECK-NEXT:    ret
+entry:
+  %cmp = icmp eq i32 %type, 6
+  %cmp1 = icmp ugt i32 %type, 0
+  %or.cond = or i1 %cmp, %cmp1
+  ret i1 %or.cond
+}
+
+define i1 @or_icmp_xlen(i32 signext %type) {
+; RV32-LABEL: or_icmp_xlen:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    addi a1, a0, -6
+; RV32-NEXT:    addi a0, a0, -32
+; RV32-NEXT:    seqz a1, a1
+; RV32-NEXT:    seqz a0, a0
+; RV32-NEXT:    or a0, a1, a0
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: or_icmp_xlen:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    li a1, 1
+; RV64-NEXT:    slli a1, a1, 32
+; RV64-NEXT:    addi a1, a1, 64
+; RV64-NEXT:    srl a1, a1, a0
+; RV64-NEXT:    sltiu a0, a0, 64
+; RV64-NEXT:    and a0, a0, a1
+; RV64-NEXT:    ret
+entry:
+  %cmp = icmp eq i32 %type, 6
+  %cmp1 = icmp eq i32 %type, 32
+  %or.cond = or i1 %cmp, %cmp1
+  ret i1 %or.cond
+}
+
+define i1 @or_icmp_i64(i64 signext %type) {
+; RV32-LABEL: or_icmp_i64:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    xori a2, a0, 6
+; RV32-NEXT:    or a3, a0, a1
+; RV32-NEXT:    xori a0, a0, 15
+; RV32-NEXT:    or a2, a2, a1
+; RV32-NEXT:    seqz a3, a3
+; RV32-NEXT:    or a0, a0, a1
+; RV32-NEXT:    seqz a1, a2
+; RV32-NEXT:    or a1, a1, a3
+; RV32-NEXT:    seqz a0, a0
+; RV32-NEXT:    or a0, a0, a1
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: or_icmp_i64:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    lui a1, 8
+; RV64-NEXT:    addi a1, a1, 65
+; RV64-NEXT:    srl a1, a1, a0
+; RV64-NEXT:    sltiu a0, a0, 64
+; RV64-NEXT:    and a0, a0, a1
+; RV64-NEXT:    ret
+entry:
+  %cmp = icmp eq i64 %type, 6
+  %cmp1 = icmp eq i64 %type, 0
+  %or.cond = or i1 %cmp, %cmp1
+  %cmp3 = icmp eq i64 %type, 15
+  %or.cond1 = or i1 %cmp3, %or.cond
+  ret i1 %or.cond1
+}
+
+define i1 @or_icmp_specific(i32 signext %type, i32 signext %type1) {
+; CHECK-LABEL: or_icmp_specific:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    addi a0, a0, -6
+; CHECK-NEXT:    addi a1, a1, -32
+; CHECK-NEXT:    seqz a0, a0
+; CHECK-NEXT:    seqz a1, a1
+; CHECK-NEXT:    or a0, a0, a1
+; CHECK-NEXT:    ret
+entry:
+  %cmp = icmp eq i32 %type, 6
+  %cmp1 = icmp eq i32 %type1, 32
+  %or.cond = or i1 %cmp, %cmp1
+  ret i1 %or.cond
+}