[RISCV][GISel] Add isel patterns for ADDIW/SRLIW/SRAIW/SLLIW and remove custom selection. (#68470)

topperc · web-flow · commit 9d1a3fdd6278 · 2023-10-13T14:44:17.000-07:00
I had trouble getting patterns working previously because GISel was
using an i32 immediate, but the instructions expected an i64 immediate
because SelectionDAG doesn't have i32 as a legal type yet.

After looking at other targets like AMDGPU, I discovered that I could
use a SDNodeXForm and a cast to get the type checking in tablegen to
allow me to do it.
diff --git a/llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp b/llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp
@@ -65,9 +65,6 @@ class RISCVInstructionSelector : public InstructionSelector {
   bool selectSelect(MachineInstr &MI, MachineIRBuilder &MIB,
                     MachineRegisterInfo &MRI) const;
 
-  bool earlySelectShift(unsigned Opc, MachineInstr &I, MachineIRBuilder &MIB,
-                        const MachineRegisterInfo &MRI);
-
   ComplexRendererFns selectShiftMask(MachineOperand &Root) const;
   ComplexRendererFns selectAddrRegImm(MachineOperand &Root) const;
 
@@ -76,6 +73,8 @@ class RISCVInstructionSelector : public InstructionSelector {
                     int OpIdx) const;
   void renderImmPlus1(MachineInstrBuilder &MIB, const MachineInstr &MI,
                       int OpIdx) const;
+  void renderImm(MachineInstrBuilder &MIB, const MachineInstr &MI,
+                 int OpIdx) const;
 
   const RISCVSubtarget &STI;
   const RISCVInstrInfo &TII;
@@ -131,30 +130,6 @@ RISCVInstructionSelector::selectAddrRegImm(MachineOperand &Root) const {
            [=](MachineInstrBuilder &MIB) { MIB.addImm(0); }}};
 }
 
-// Tablegen doesn't allow us to write SRLIW/SRAIW/SLLIW patterns because the
-// immediate Operand has type XLenVT. GlobalISel wants it to be i32.
-bool RISCVInstructionSelector::earlySelectShift(
-    unsigned Opc, MachineInstr &I, MachineIRBuilder &MIB,
-    const MachineRegisterInfo &MRI) {
-  if (!Subtarget->is64Bit())
-    return false;
-
-  LLT Ty = MRI.getType(I.getOperand(0).getReg());
-  if (!Ty.isScalar() || Ty.getSizeInBits() != 32)
-    return false;
-
-  std::optional<int64_t> CstVal =
-      getIConstantVRegSExtVal(I.getOperand(2).getReg(), MRI);
-  if (!CstVal || !isUInt<5>(*CstVal))
-    return false;
-
-  auto NewI = MIB.buildInstr(Opc, {I.getOperand(0).getReg()},
-                             {I.getOperand(1).getReg()})
-                  .addImm(*CstVal);
-  I.eraseFromParent();
-  return constrainSelectedInstRegOperands(*NewI, TII, TRI, RBI);
-}
-
 bool RISCVInstructionSelector::select(MachineInstr &MI) {
   MachineBasicBlock &MBB = *MI.getParent();
   MachineFunction &MF = *MBB.getParent();
@@ -199,55 +174,6 @@ bool RISCVInstructionSelector::select(MachineInstr &MI) {
     return true;
   }
 
-  switch (Opc) {
-  case TargetOpcode::G_ADD: {
-    // Tablegen doesn't pick up the ADDIW pattern because i32 isn't a legal
-    // type for RV64 in SelectionDAG. Manually select it here.
-    LLT Ty = MRI.getType(MI.getOperand(0).getReg());
-    if (Subtarget->is64Bit() && Ty.isScalar() && Ty.getSizeInBits() == 32) {
-      std::optional<int64_t> CstVal =
-          getIConstantVRegSExtVal(MI.getOperand(2).getReg(), MRI);
-      if (CstVal && isInt<12>(*CstVal)) {
-        auto NewI = MIB.buildInstr(RISCV::ADDIW, {MI.getOperand(0).getReg()},
-                                   {MI.getOperand(1).getReg()})
-                        .addImm(*CstVal);
-        MI.eraseFromParent();
-        return constrainSelectedInstRegOperands(*NewI, TII, TRI, RBI);
-      }
-    }
-    break;
-  }
-  case TargetOpcode::G_SUB: {
-    // Tablegen doesn't pick up the ADDIW pattern because i32 isn't a legal
-    // type for RV64 in SelectionDAG. Manually select it here.
-    LLT Ty = MRI.getType(MI.getOperand(0).getReg());
-    if (Subtarget->is64Bit() && Ty.isScalar() && Ty.getSizeInBits() == 32) {
-      std::optional<int64_t> CstVal =
-          getIConstantVRegSExtVal(MI.getOperand(2).getReg(), MRI);
-      if (CstVal && ((isInt<12>(*CstVal) && *CstVal != -2048) || *CstVal == 2048)) {
-        auto NewI = MIB.buildInstr(RISCV::ADDIW, {MI.getOperand(0).getReg()},
-                                   {MI.getOperand(1).getReg()})
-                        .addImm(-*CstVal);
-        MI.eraseFromParent();
-        return constrainSelectedInstRegOperands(*NewI, TII, TRI, RBI);
-      }
-    }
-    break;
-  }
-  case TargetOpcode::G_ASHR:
-    if (earlySelectShift(RISCV::SRAIW, MI, MIB, MRI))
-      return true;
-    break;
-  case TargetOpcode::G_LSHR:
-    if (earlySelectShift(RISCV::SRLIW, MI, MIB, MRI))
-      return true;
-    break;
-  case TargetOpcode::G_SHL:
-    if (earlySelectShift(RISCV::SLLIW, MI, MIB, MRI))
-      return true;
-    break;
-  }
-
   if (selectImpl(MI, *CoverageInfo))
     return true;
 
@@ -323,6 +249,15 @@ void RISCVInstructionSelector::renderImmPlus1(MachineInstrBuilder &MIB,
   MIB.addImm(CstVal + 1);
 }
 
+void RISCVInstructionSelector::renderImm(MachineInstrBuilder &MIB,
+                                         const MachineInstr &MI,
+                                         int OpIdx) const {
+  assert(MI.getOpcode() == TargetOpcode::G_CONSTANT && OpIdx == -1 &&
+         "Expected G_CONSTANT");
+  int64_t CstVal = MI.getOperand(1).getCImm()->getSExtValue();
+  MIB.addImm(CstVal);
+}
+
 const TargetRegisterClass *RISCVInstructionSelector::getRegClassForTypeOnBank(
     LLT Ty, const RegisterBank &RB) const {
   if (RB.getID() == RISCV::GPRRegBankID) {
diff --git a/llvm/lib/Target/RISCV/RISCVGISel.td b/llvm/lib/Target/RISCV/RISCVGISel.td
@@ -18,6 +18,12 @@ include "RISCVCombine.td"
 
 def simm12Plus1 : ImmLeaf<XLenVT, [{
     return (isInt<12>(Imm) && Imm != -2048) || Imm == 2048;}]>;
+def simm12Plus1i32 : ImmLeaf<i32, [{
+    return (isInt<12>(Imm) && Imm != -2048) || Imm == 2048;}]>;
+
+def simm12i32 : ImmLeaf<i32, [{return isInt<12>(Imm);}]>;
+
+def uimm5i32 : ImmLeaf<i32, [{return isUInt<5>(Imm);}]>;
 
 // FIXME: This doesn't check that the G_CONSTANT we're deriving the immediate
 // from is only used once
@@ -43,6 +49,14 @@ def GIAddrRegImm :
   GIComplexOperandMatcher<s32, "selectAddrRegImm">,
   GIComplexPatternEquiv<AddrRegImm>;
 
+// Convert from i32 immediate to i64 target immediate to make SelectionDAG type
+// checking happy so we can use ADDIW which expects an XLen immediate.
+def as_i64imm : SDNodeXForm<imm, [{
+  return CurDAG->getTargetConstant(N->getSExtValue(), SDLoc(N), MVT::i64);
+}]>;
+def gi_as_i64imm : GICustomOperandRenderer<"renderImm">,
+  GISDNodeXFormEquiv<as_i64imm>;
+
 // FIXME: This is labelled as handling 's32', however the ComplexPattern it
 // refers to handles both i32 and i64 based on the HwMode. Currently this LLT
 // parameter appears to be ignored so this pattern works for both, however we
@@ -60,11 +74,23 @@ let Predicates = [IsRV64] in {
 def : Pat<(i32 (add GPR:$rs1, GPR:$rs2)), (ADDW GPR:$rs1, GPR:$rs2)>;
 def : Pat<(i32 (sub GPR:$rs1, GPR:$rs2)), (SUBW GPR:$rs1, GPR:$rs2)>;
 
+def : Pat<(i32 (add GPR:$rs1, simm12i32:$imm)),
+          (ADDIW GPR:$rs1, (i64 (as_i64imm $imm)))>;
+def : Pat<(i32 (sub GPR:$rs1, simm12Plus1i32:$imm)),
+          (ADDIW GPR:$rs1, (i64 (NegImm $imm)))>;
+
 def : Pat<(i32 (shl GPR:$rs1, (i32 GPR:$rs2))), (SLLW GPR:$rs1, GPR:$rs2)>;
 def : Pat<(i32 (sra GPR:$rs1, (i32 GPR:$rs2))), (SRAW GPR:$rs1, GPR:$rs2)>;
 def : Pat<(i32 (srl GPR:$rs1, (i32 GPR:$rs2))), (SRLW GPR:$rs1, GPR:$rs2)>;
 
-def: Pat<(i64 (sext i32:$rs)), (ADDIW GPR:$rs, 0)>;
+def : Pat<(i32 (shl GPR:$rs1, uimm5i32:$imm)),
+          (SLLIW GPR:$rs1, (i64 (as_i64imm $imm)))>;
+def : Pat<(i32 (sra GPR:$rs1, uimm5i32:$imm)),
+          (SRAIW GPR:$rs1, (i64 (as_i64imm $imm)))>;
+def : Pat<(i32 (srl GPR:$rs1, uimm5i32:$imm)),
+          (SRLIW GPR:$rs1, (i64 (as_i64imm $imm)))>;
+
+def : Pat<(i64 (sext i32:$rs)), (ADDIW GPR:$rs, 0)>;
 }
 
 let Predicates = [HasStdExtMOrZmmul, IsRV64] in {
