AMDGPU/GlobalISel: Implement applyMappingImpl less incorrectly

[arsenm]

Summary

We're checking the current register bank of the registers in the
instruction, but the mapping may have inserted cross bank copies and
is expecting to replace the registers.

We mostly get away with this currently, because VGPR->SGPR copies are
illegal, and we assume this won't happen. In a future change, we'll
start relying on more cross register bank copies being inserted, and
this starts to break down.

[kerbowa]

LGTM

[arsenm]

5eed4e2664aa7187f46eb12e45a376d7ab7dd308

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

bool AMDGPURegisterBankInfo::applyMappingWideLoad(MachineInstr &MI,
   if (LoadSize <= MaxNonSmrdLoadSize)
     return false;
 
   SmallVector<unsigned, 16> DefRegs(OpdMapper.getVRegs(0));
   SmallVector<unsigned, 1> SrcRegs(OpdMapper.getVRegs(1));
 
   // If the pointer is an SGPR, we have nothing to do.
   if (SrcRegs.empty()) {
-    Register PtrReg = MI.getOperand(1).getReg();
-    const RegisterBank *PtrBank = getRegBank(PtrReg, MRI, *TRI);
+    const RegisterBank *PtrBank =
+      OpdMapper.getInstrMapping().getOperandMapping(1).BreakDown[0].RegBank;
     if (PtrBank == &AMDGPU::SGPRRegBank)
       return false;
-    SrcRegs.push_back(PtrReg);
+    SrcRegs.push_back(MI.getOperand(1).getReg());
   }
 
   assert(LoadSize % MaxNonSmrdLoadSize == 0);
 
   // We want to get the repair instruction now, because it will help us
   // determine which instruction the legalizer inserts that will also
   // write to DstReg.
   MachineInstr *RepairInst = getOtherVRegDef(MRI, DstReg, MI);
void AMDGPURegisterBankInfo::applyMappingImpl(
   case AMDGPU::G_ADD:
   case AMDGPU::G_SUB:
   case AMDGPU::G_MUL: {
     Register DstReg = MI.getOperand(0).getReg();
     LLT DstTy = MRI.getType(DstReg);
     if (DstTy != LLT::scalar(16))
       break;
 
-    const RegisterBank *DstBank = getRegBank(DstReg, MRI, *TRI);
+    const RegisterBank *DstBank =
+      OpdMapper.getInstrMapping().getOperandMapping(0).BreakDown[0].RegBank;
     if (DstBank == &AMDGPU::VGPRRegBank)
       break;
 
     // 16-bit operations are VALU only, but can be promoted to 32-bit SALU.
     MachineFunction *MF = MI.getParent()->getParent();
     MachineIRBuilder B(MI);
     ApplyRegBankMapping ApplySALU(MRI, &AMDGPU::SGPRRegBank);
     GISelObserverWrapper Observer(&ApplySALU);
     LegalizerHelper Helper(*MF, Observer, B);
 
     if (Helper.widenScalar(MI, 0, LLT::scalar(32)) !=
         LegalizerHelper::Legalized)
       llvm_unreachable("widen scalar should have succeeded");
     return;
   }
   case AMDGPU::G_SMIN:
   case AMDGPU::G_SMAX:
   case AMDGPU::G_UMIN:
   case AMDGPU::G_UMAX: {
     Register DstReg = MI.getOperand(0).getReg();
-    const RegisterBank *DstBank = getRegBank(DstReg, MRI, *TRI);
+    const RegisterBank *DstBank =
+      OpdMapper.getInstrMapping().getOperandMapping(0).BreakDown[0].RegBank;
     if (DstBank == &AMDGPU::VGPRRegBank)
       break;
 
     MachineFunction *MF = MI.getParent()->getParent();
     MachineIRBuilder B(MI);
     ApplyRegBankMapping ApplySALU(MRI, &AMDGPU::SGPRRegBank);
     GISelObserverWrapper Observer(&ApplySALU);
     LegalizerHelper Helper(*MF, Observer, B);
void AMDGPURegisterBankInfo::applyMappingImpl(
   }
   case AMDGPU::G_SEXT:
   case AMDGPU::G_ZEXT: {
     Register SrcReg = MI.getOperand(1).getReg();
     LLT SrcTy = MRI.getType(SrcReg);
     bool Signed = Opc == AMDGPU::G_SEXT;
 
     MachineIRBuilder B(MI);
-    const RegisterBank *SrcBank = getRegBank(SrcReg, MRI, *TRI);
+    const RegisterBank *SrcBank =
+      OpdMapper.getInstrMapping().getOperandMapping(1).BreakDown[0].RegBank;
 
     Register DstReg = MI.getOperand(0).getReg();
     LLT DstTy = MRI.getType(DstReg);
     if (DstTy.isScalar() &&
         SrcBank != &AMDGPU::SGPRRegBank &&
         SrcBank != &AMDGPU::SCCRegBank &&
         SrcBank != &AMDGPU::VCCRegBank &&
         // FIXME: Should handle any type that round to s64 when irregular
case AMDGPU::G_BUILD_VECTOR_TRUNC: {
     LLT DstTy = MRI.getType(DstReg);
     if (DstTy != LLT::vector(2, 16))
       break;
 
     assert(MI.getNumOperands() == 3 && OpdMapper.getVRegs(0).empty());
     substituteSimpleCopyRegs(OpdMapper, 1);
     substituteSimpleCopyRegs(OpdMapper, 2);
 
-    const RegisterBank *DstBank = getRegBank(DstReg, MRI, *TRI);
+    const RegisterBank *DstBank =
+      OpdMapper.getInstrMapping().getOperandMapping(0).BreakDown[0].RegBank;
     if (DstBank == &AMDGPU::SGPRRegBank)
       break; // Can use S_PACK_* instructions.
 
     MachineIRBuilder B(MI);
 
     Register Lo = MI.getOperand(1).getReg();
     Register Hi = MI.getOperand(2).getReg();
     const LLT S32 = LLT::scalar(32);
 
-    const RegisterBank *BankLo = getRegBank(Lo, MRI, *TRI);
-    const RegisterBank *BankHi = getRegBank(Hi, MRI, *TRI);
+    const RegisterBank *BankLo =
+      OpdMapper.getInstrMapping().getOperandMapping(1).BreakDown[0].RegBank;
+    const RegisterBank *BankHi =
+      OpdMapper.getInstrMapping().getOperandMapping(2).BreakDown[0].RegBank;
 
     Register ZextLo;
     Register ShiftHi;
 
     if (Opc == AMDGPU::G_BUILD_VECTOR) {
       ZextLo = B.buildZExt(S32, Lo).getReg(0);
       MRI.setRegBank(ZextLo, *BankLo);
 
case AMDGPU::G_EXTRACT_VECTOR_ELT: {
     auto IdxHi = B.buildAdd(S32, IdxLo, One);
     B.buildExtractVectorElement(DstRegs[0], CastSrc, IdxLo);
     B.buildExtractVectorElement(DstRegs[1], CastSrc, IdxHi);
 
     const ValueMapping &DstMapping
       = OpdMapper.getInstrMapping().getOperandMapping(0);
 
     // FIXME: Should be getting from mapping or not?
-    const RegisterBank *SrcBank = getRegBank(SrcReg, MRI, *TRI);
+    const RegisterBank *SrcBank =
+      OpdMapper.getInstrMapping().getOperandMapping(1).BreakDown[0].RegBank;
     MRI.setRegBank(DstReg, *DstMapping.BreakDown[0].RegBank);
     MRI.setRegBank(CastSrc.getReg(0), *SrcBank);
     MRI.setRegBank(One.getReg(0), AMDGPU::SGPRRegBank);
     MRI.setRegBank(IdxLo.getReg(0), AMDGPU::SGPRRegBank);
     MRI.setRegBank(IdxHi.getReg(0), AMDGPU::SGPRRegBank);
 
     SmallSet<Register, 4> OpsToWaterfall;
     if (!collectWaterfallOperands(OpsToWaterfall, MI, MRI, { 2 })) {
case AMDGPU::G_INSERT_VECTOR_ELT: {
     // Compute 32-bit element indices, (2 * OrigIdx, 2 * OrigIdx + 1).
     auto IdxLo = B.buildShl(S32, IdxReg, One);
     auto IdxHi = B.buildAdd(S32, IdxLo, One);
 
     auto InsLo = B.buildInsertVectorElement(Vec32, CastSrc, InsRegs[0], IdxLo);
     auto InsHi = B.buildInsertVectorElement(Vec32, InsLo, InsRegs[1], IdxHi);
     B.buildBitcast(DstReg, InsHi);
 
-    const RegisterBank *DstBank = getRegBank(DstReg, MRI, *TRI);
-    const RegisterBank *SrcBank = getRegBank(SrcReg, MRI, *TRI);
-    const RegisterBank *InsSrcBank = getRegBank(InsReg, MRI, *TRI);
+    const RegisterBank *DstBank =
+      OpdMapper.getInstrMapping().getOperandMapping(0).BreakDown[0].RegBank;
+    const RegisterBank *SrcBank =
+      OpdMapper.getInstrMapping().getOperandMapping(1).BreakDown[0].RegBank;
+    const RegisterBank *InsSrcBank =
+      OpdMapper.getInstrMapping().getOperandMapping(2).BreakDown[0].RegBank;
 
     MRI.setRegBank(InsReg, *InsSrcBank);
     MRI.setRegBank(CastSrc.getReg(0), *SrcBank);
     MRI.setRegBank(InsLo.getReg(0), *DstBank);
     MRI.setRegBank(InsHi.getReg(0), *DstBank);
     MRI.setRegBank(One.getReg(0), AMDGPU::SGPRRegBank);
     MRI.setRegBank(IdxLo.getReg(0), AMDGPU::SGPRRegBank);
     MRI.setRegBank(IdxHi.getReg(0), AMDGPU::SGPRRegBank);