[MIPS GlobalISel] Improve reg bank handling in MipsInstructionSelector

[Petar.Avramovic]

Summary

Introduce helper methods and refactor pieces of code related to
register banks in MipsInstructionSelector.
Add a few detailed asserts in order to get a better overview
of LLT, register bank combinations that are supported at the moment
and reduce need to look at other files.

[atanasyan]

LGTM with a nit

llvm/lib/Target/Mips/MipsInstructionSelector.cpp

243

We can delete this return and the empty line above. The next return does the task.

[Petar.Avramovic]

Address review comments.

llvm/lib/Target/Mips/MipsInstructionSelector.cpp

public:
   MipsInstructionSelector(const MipsTargetMachine &TM, const MipsSubtarget &STI,
                           const MipsRegisterBankInfo &RBI);
 
   bool select(MachineInstr &I) override;
   static const char *getName() { return DEBUG_TYPE; }
 
 private:
   bool selectImpl(MachineInstr &I, CodeGenCoverage &CoverageInfo) const;
+  bool isRegInGprb(Register Reg, MachineRegisterInfo &MRI) const;
+  bool isRegInFprb(Register Reg, MachineRegisterInfo &MRI) const;
   bool materialize32BitImm(Register DestReg, APInt Imm,
                            MachineIRBuilder &B) const;
   bool selectCopy(MachineInstr &I, MachineRegisterInfo &MRI) const;
   const TargetRegisterClass *
-  getRegClassForTypeOnBank(unsigned OpSize, const RegisterBank &RB,
-                           const RegisterBankInfo &RBI) const;
+  getRegClassForTypeOnBank(Register Reg, MachineRegisterInfo &MRI) const;
   unsigned selectLoadStoreOpCode(MachineInstr &I,
                                  MachineRegisterInfo &MRI) const;
 
   const MipsTargetMachine &TM;
   const MipsSubtarget &STI;
   const MipsInstrInfo &TII;
   const MipsRegisterInfo &TRI;
   const MipsRegisterBankInfo &RBI;

 #include "MipsGenGlobalISel.inc"
 #undef GET_GLOBALISEL_PREDICATES_INIT
 #define GET_GLOBALISEL_TEMPORARIES_INIT
 #include "MipsGenGlobalISel.inc"
 #undef GET_GLOBALISEL_TEMPORARIES_INIT
 {
 }
 
+bool MipsInstructionSelector::isRegInGprb(Register Reg,
+                                          MachineRegisterInfo &MRI) const {
+  return RBI.getRegBank(Reg, MRI, TRI)->getID() == Mips::GPRBRegBankID;
+}
+
+bool MipsInstructionSelector::isRegInFprb(Register Reg,
+                                          MachineRegisterInfo &MRI) const {
+  return RBI.getRegBank(Reg, MRI, TRI)->getID() == Mips::FPRBRegBankID;
+}
+
 bool MipsInstructionSelector::selectCopy(MachineInstr &I,
                                          MachineRegisterInfo &MRI) const {
   Register DstReg = I.getOperand(0).getReg();
   if (Register::isPhysicalRegister(DstReg))
     return true;
 
-  const RegisterBank *RegBank = RBI.getRegBank(DstReg, MRI, TRI);
-  const unsigned DstSize = MRI.getType(DstReg).getSizeInBits();
-
-  const TargetRegisterClass *RC = &Mips::GPR32RegClass;
-  if (RegBank->getID() == Mips::FPRBRegBankID) {
-    if (DstSize == 32)
-      RC = &Mips::FGR32RegClass;
-    else if (DstSize == 64)
-      RC = STI.isFP64bit() ? &Mips::FGR64RegClass : &Mips::AFGR64RegClass;
-    else
-      llvm_unreachable("Unsupported destination size");
-  }
+  const TargetRegisterClass *RC = getRegClassForTypeOnBank(DstReg, MRI);
   if (!RBI.constrainGenericRegister(DstReg, *RC, MRI)) {
     LLVM_DEBUG(dbgs() << "Failed to constrain " << TII.getName(I.getOpcode())
                       << " operand\n");
     return false;
   }
   return true;
 }
 
 const TargetRegisterClass *MipsInstructionSelector::getRegClassForTypeOnBank(
-    unsigned OpSize, const RegisterBank &RB,
-    const RegisterBankInfo &RBI) const {
-  if (RB.getID() == Mips::GPRBRegBankID)
+    Register Reg, MachineRegisterInfo &MRI) const {
+  const LLT Ty = MRI.getType(Reg);
+  const unsigned TySize = Ty.getSizeInBits();
+
+  if (isRegInGprb(Reg, MRI)) {
+    assert((Ty.isScalar() || Ty.isPointer()) && TySize == 32 &&
+           "Register class not available for LLT, register bank combination");
     return &Mips::GPR32RegClass;
+  }
 
-  if (RB.getID() == Mips::FPRBRegBankID)
-    return OpSize == 32
-               ? &Mips::FGR32RegClass
-               : STI.hasMips32r6() || STI.isFP64bit() ? &Mips::FGR64RegClass
-                                                      : &Mips::AFGR64RegClass;
+  if (isRegInFprb(Reg, MRI)) {
+    if (Ty.isScalar()) {
+      assert((TySize == 32 || TySize == 64) &&
+             "Register class not available for LLT, register bank combination");
+      if (TySize == 32)
+        return &Mips::FGR32RegClass;
+      return STI.isFP64bit() ? &Mips::FGR64RegClass : &Mips::AFGR64RegClass;
+    }
+  }
 
-  llvm_unreachable("getRegClassForTypeOnBank can't find register class.");
-  return nullptr;
+  llvm_unreachable("Unsupported register bank.");
 }
 
 bool MipsInstructionSelector::materialize32BitImm(Register DestReg, APInt Imm,
                                                   MachineIRBuilder &B) const {
   assert(Imm.getBitWidth() == 32 && "Unsupported immediate size.");
   // Ori zero extends immediate. Used for values with zeros in high 16 bits.
   if (Imm.getHiBits(16).isNullValue()) {
     MachineInstr *Inst = B.buildInstr(Mips::ORi, {DestReg}, {Register(Mips::ZERO)})
MachineInstr *ORi = B.buildInstr(Mips::ORi, {DestReg}, {LUiReg})
                           .addImm(Imm.getLoBits(16).getLimitedValue());
   if (!constrainSelectedInstRegOperands(*LUi, TII, TRI, RBI))
     return false;
   if (!constrainSelectedInstRegOperands(*ORi, TII, TRI, RBI))
     return false;
   return true;
 }
 
-/// Returning Opc indicates that we failed to select MIPS instruction opcode.
+/// When I.getOpcode() is returned, we failed to select MIPS instruction opcode.
 unsigned
 MipsInstructionSelector::selectLoadStoreOpCode(MachineInstr &I,
                                                MachineRegisterInfo &MRI) const {
-  STI.getRegisterInfo();
-  const Register DestReg = I.getOperand(0).getReg();
-  const unsigned RegBank = RBI.getRegBank(DestReg, MRI, TRI)->getID();
+  const Register ValueReg = I.getOperand(0).getReg();
+  const LLT Ty = MRI.getType(ValueReg);
+  const unsigned TySize = Ty.getSizeInBits();
   const unsigned MemSizeInBytes = (*I.memoperands_begin())->getSize();
   unsigned Opc = I.getOpcode();
   const bool isStore = Opc == TargetOpcode::G_STORE;
-  if (RegBank == Mips::GPRBRegBankID) {
+
+  if (isRegInGprb(ValueReg, MRI)) {
+    assert(((Ty.isScalar() && TySize == 32) ||
+            (Ty.isPointer() && TySize == 32 && MemSizeInBytes == 4)) &&
+           "Unsupported register bank, LLT, MemSizeInBytes combination");
     if (isStore)
       switch (MemSizeInBytes) {
       case 4:
         return Mips::SW;
       case 2:
         return Mips::SH;
       case 1:
         return Mips::SB;
else
         return Opc == TargetOpcode::G_SEXTLOAD ? Mips::LH : Mips::LHu;
       case 1:
         return Opc == TargetOpcode::G_SEXTLOAD ? Mips::LB : Mips::LBu;
       default:
         return Opc;
       }
   }
 
-  if (RegBank == Mips::FPRBRegBankID) {
-    switch (MemSizeInBytes) {
-    case 4:
+  if (isRegInFprb(ValueReg, MRI)) {
+    if (Ty.isScalar()) {
+      assert(((TySize == 32 && MemSizeInBytes == 4) ||
+              (TySize == 64 && MemSizeInBytes == 8)) &&
+             "Unsupported register bank, LLT, MemSizeInBytes combination");
+
+      if (MemSizeInBytes == 4)
         return isStore ? Mips::SWC1 : Mips::LWC1;
-    case 8:
+
       if (STI.isFP64bit())
         return isStore ? Mips::SDC164 : Mips::LDC164;
-      else
-        return isStore ? Mips::SDC1 : Mips::LDC1;
-    case 16: {
+    }
+
+    if (Ty.isVector()) {
       assert(STI.hasMSA() && "Vector instructions require target with MSA.");
-      const unsigned VectorElementSizeInBytes =
-          MRI.getType(DestReg).getElementType().getSizeInBytes();
-      if (VectorElementSizeInBytes == 1)
+      assert((TySize == 128 && MemSizeInBytes == 16) &&
+             "Unsupported register bank, LLT, MemSizeInBytes combination");
+      switch (Ty.getElementType().getSizeInBits()) {
+      case 8:
         return isStore ? Mips::ST_B : Mips::LD_B;
-      if (VectorElementSizeInBytes == 2)
+      case 16:
         return isStore ? Mips::ST_H : Mips::LD_H;
-      if (VectorElementSizeInBytes == 4)
+      case 32:
         return isStore ? Mips::ST_W : Mips::LD_W;
-      if (VectorElementSizeInBytes == 8)
+      case 64:
         return isStore ? Mips::ST_D : Mips::LD_D;
-      return Opc;
-    }
       default:
         return Opc;
       }
     }
+
+    return Opc;
+  }
+
   return Opc;
 }
 
 bool MipsInstructionSelector::select(MachineInstr &I) {
 
   MachineBasicBlock &MBB = *I.getParent();
   MachineFunction &MF = *MBB.getParent();
   MachineRegisterInfo &MRI = MF.getRegInfo();
 
   if (!isPreISelGenericOpcode(I.getOpcode())) {
     if (I.isCopy())
       return selectCopy(I, MRI);
 
     return true;
   }
 
   if (I.getOpcode() == Mips::G_MUL &&
-      (RBI.getRegBank(I.getOperand(0).getReg(), MRI, TRI)->getID() ==
-       Mips::GPRBRegBankID)) {
+      isRegInGprb(I.getOperand(0).getReg(), MRI)) {
     MachineInstr *Mul = BuildMI(MBB, I, I.getDebugLoc(), TII.get(Mips::MUL))
                             .add(I.getOperand(0))
                             .add(I.getOperand(1))
                             .add(I.getOperand(2));
     if (!constrainSelectedInstRegOperands(*Mul, TII, TRI, RBI))
       return false;
     Mul->getOperand(3).setIsDead(true);
     Mul->getOperand(4).setIsDead(true);
bool MipsInstructionSelector::select(MachineInstr &I) {
   }
   case G_BRINDIRECT: {
     MI = BuildMI(MBB, I, I.getDebugLoc(), TII.get(Mips::PseudoIndirectBranch))
              .add(I.getOperand(0));
     break;
   }
   case G_PHI: {
     const Register DestReg = I.getOperand(0).getReg();
-    const unsigned OpSize = MRI.getType(DestReg).getSizeInBits();
 
     const TargetRegisterClass *DefRC = nullptr;
     if (Register::isPhysicalRegister(DestReg))
       DefRC = TRI.getRegClass(DestReg);
     else
-      DefRC = getRegClassForTypeOnBank(OpSize,
-                                       *RBI.getRegBank(DestReg, MRI, TRI), RBI);
+      DefRC = getRegClassForTypeOnBank(DestReg, MRI);
 
     I.setDesc(TII.get(TargetOpcode::PHI));
     return RBI.constrainGenericRegister(DestReg, *DefRC, MRI);
   }
   case G_STORE:
   case G_LOAD:
   case G_ZEXTLOAD:
   case G_SEXTLOAD: {
case G_SELECT: {
     MI = BuildMI(MBB, I, I.getDebugLoc(), TII.get(Mips::MOVN_I_I))
              .add(I.getOperand(0))
              .add(I.getOperand(2))
              .add(I.getOperand(1))
              .add(I.getOperand(3));
     break;
   }
   case G_IMPLICIT_DEF: {
+    Register Dst = I.getOperand(0).getReg();
     MI = BuildMI(MBB, I, I.getDebugLoc(), TII.get(Mips::IMPLICIT_DEF))
-             .add(I.getOperand(0));
+             .addDef(Dst);
 
     // Set class based on register bank, there can be fpr and gpr implicit def.
-    MRI.setRegClass(MI->getOperand(0).getReg(),
-                    getRegClassForTypeOnBank(
-                        MRI.getType(I.getOperand(0).getReg()).getSizeInBits(),
-                        *RBI.getRegBank(I.getOperand(0).getReg(), MRI, TRI),
-                        RBI));
+    MRI.setRegClass(Dst, getRegClassForTypeOnBank(Dst, MRI));
     break;
   }
   case G_CONSTANT: {
     MachineIRBuilder B(I);
     if (!materialize32BitImm(I.getOperand(0).getReg(),
                              I.getOperand(1).getCImm()->getValue(), B))
       return false;