[llvm] [RISCV][GISel] Select G_SELECT (G_ICMP, A, B) (PR #68247)

[Craig Topper via llvm-commits]

================
@@ -376,6 +392,121 @@ bool RISCVInstructionSelector::selectSExtInreg(MachineInstr &MI,
   return true;
 }
 
+/// Returns the RISCVCC::CondCode that corresponds to the CmpInst::Predicate CC.
+/// CC Must be an ICMP Predicate.
+static RISCVCC::CondCode getRISCVCCFromICMP(CmpInst::Predicate CC) {
+  switch (CC) {
+  default:
+    llvm_unreachable("Expected ICMP CmpInst::Predicate.");
+  case CmpInst::Predicate::ICMP_EQ:
+    return RISCVCC::COND_EQ;
+  case CmpInst::Predicate::ICMP_NE:
+    return RISCVCC::COND_NE;
+  case CmpInst::Predicate::ICMP_ULT:
+    return RISCVCC::COND_LTU;
+  case CmpInst::Predicate::ICMP_SLT:
+    return RISCVCC::COND_LT;
+  case CmpInst::Predicate::ICMP_UGE:
+    return RISCVCC::COND_GEU;
+  case CmpInst::Predicate::ICMP_SGE:
+    return RISCVCC::COND_GE;
+  }
+}
+
+void RISCVInstructionSelector::setICMPOperandsForBranch(
+    MachineInstr &MI, MachineIRBuilder &MIB, MachineRegisterInfo &MRI,
+    CmpInst::Predicate &CC, MachineOperand &LHS, MachineOperand &RHS) const {
+  assert(MI.getOpcode() == TargetOpcode::G_ICMP);
+  CC = static_cast<CmpInst::Predicate>(MI.getOperand(1).getPredicate());
+  LHS = MI.getOperand(2);
+  RHS = MI.getOperand(3);
+
+  // Adjust comparisons to use comparison with 0 if possible.
+  MachineInstr *MaybeConstant = MRI.getVRegDef(RHS.getReg());
+  if (MaybeConstant && MaybeConstant->getOpcode() == TargetOpcode::G_CONSTANT) {
+    switch (CC) {
+    case CmpInst::Predicate::ICMP_SGT:
+      // Convert X > -1 to X >= 0
+      if (MaybeConstant->getOperand(1).getCImm()->getSExtValue() == -1) {
+        MachineInstr *Zero = MIB.buildConstant(
+            MRI.getType(MaybeConstant->getOperand(0).getReg()), 0);
+        selectConstant(*Zero, MIB, MRI);
+        CC = CmpInst::Predicate::ICMP_SGE;
+        RHS = MachineOperand::CreateReg(Zero->getOperand(0).getReg(), false);
+      }
+      return;
+    case CmpInst::Predicate::ICMP_SLT:
+      // Convert X < 1 to 0 >= X
+      if (MaybeConstant->getOperand(1).getCImm()->getSExtValue() == 1) {
+        MachineInstr *Zero = MIB.buildConstant(
+            MRI.getType(MaybeConstant->getOperand(0).getReg()), 0);
+        selectConstant(*Zero, MIB, MRI);
+        CC =  CmpInst::Predicate::ICMP_SGE;
+        RHS = LHS;
+        LHS = MachineOperand::CreateReg(Zero->getOperand(0).getReg(), false);
+      }
+      return;
+    default:
+      break;
+    }
+  }
+
+  switch (CC) {
+  default:
+    llvm_unreachable("Expected ICMP CmpInst::Predicate.");
+  case CmpInst::Predicate::ICMP_EQ:
+  case CmpInst::Predicate::ICMP_NE:
+  case CmpInst::Predicate::ICMP_ULT:
+  case CmpInst::Predicate::ICMP_SLT:
+  case CmpInst::Predicate::ICMP_UGE:
+  case CmpInst::Predicate::ICMP_SGE:
+    // These CCs are supported directly by RISC-V branches.
+    return;
+  case CmpInst::Predicate::ICMP_SGT:
+  case CmpInst::Predicate::ICMP_SLE:
+  case CmpInst::Predicate::ICMP_UGT:
+  case CmpInst::Predicate::ICMP_ULE:
+    // These CCs are not supported directly by RISC-V branches, but changing the
+    // direction of the CC and swapping LHS and RHS are.
+    CC = CmpInst::getSwappedPredicate(CC);
+    std::swap(LHS, RHS);
+    return;
+  }
+}
+
+bool RISCVInstructionSelector::selectSelect(MachineInstr &MI,
+                                            MachineIRBuilder &MIB,
+                                            MachineRegisterInfo &MRI) const {
+  assert(MI.getOpcode() == TargetOpcode::G_SELECT);
+  MachineInstr *Result;
+  MachineInstr *MaybeICMP = MRI.getVRegDef(MI.getOperand(1).getReg());
+  if (MaybeICMP && MaybeICMP->getOpcode() == TargetOpcode::G_ICMP) {
+    // If MI is a G_SELECT(G_ICMP(tst, A, B), C, D) then we can use (A, B, tst)
+    // as the (LHS, RHS, CC) of the Select_GPR_Using_CC_GPR.
+    CmpInst::Predicate CC;
+    MachineOperand LHS = MaybeICMP->getOperand(2);
+    MachineOperand RHS = MaybeICMP->getOperand(3);
+    setICMPOperandsForBranch(*MaybeICMP, MIB, MRI, CC, LHS, RHS);
+    Result = MIB.buildInstr(RISCV::Select_GPR_Using_CC_GPR)
+                 .addDef(MI.getOperand(0).getReg());
+    Result->addOperand(LHS);
+    Result->addOperand(RHS);
+    Result->addOperand(MachineOperand::CreateImm(getRISCVCCFromICMP(CC)));
----------------
topperc wrote:

Should we call `getRISCVCCFromICMP` from `setICMPOperandsForBranch` and return the translated value insead of the CmpInst::Predicate? We're going to need to do this again for G_BRCOND.

https://github.com/llvm/llvm-project/pull/68247