[llvm] AMDGPU/GlobalISel: AMDGPURegBankSelect (PR #112863)

[Petar Avramovic via llvm-commits]

================
@@ -66,9 +78,183 @@ FunctionPass *llvm::createAMDGPURegBankSelectPass() {
   return new AMDGPURegBankSelect();
 }
 
+class RegBankSelectHelper {
+  MachineIRBuilder &B;
+  MachineRegisterInfo &MRI;
+  AMDGPU::IntrinsicLaneMaskAnalyzer &ILMA;
+  const MachineUniformityInfo &MUI;
+  const RegisterBank *SgprRB;
+  const RegisterBank *VgprRB;
+  const RegisterBank *VccRB;
+
+public:
+  RegBankSelectHelper(MachineIRBuilder &B,
+                      AMDGPU::IntrinsicLaneMaskAnalyzer &ILMA,
+                      const MachineUniformityInfo &MUI,
+                      const RegisterBankInfo &RBI)
+      : B(B), MRI(*B.getMRI()), ILMA(ILMA), MUI(MUI),
+        SgprRB(&RBI.getRegBank(AMDGPU::SGPRRegBankID)),
+        VgprRB(&RBI.getRegBank(AMDGPU::VGPRRegBankID)),
+        VccRB(&RBI.getRegBank(AMDGPU::VCCRegBankID)) {}
+
+  const RegisterBank *getRegBankToAssign(Register Reg) {
+    if (MUI.isUniform(Reg) || ILMA.isS32S64LaneMask(Reg))
+      return SgprRB;
+    if (MRI.getType(Reg) == LLT::scalar(1))
+      return VccRB;
+    return VgprRB;
+  }
+
+  // %rc:RegClass(s32) = G_ ...
+  // ...
+  // %a = G_ ..., %rc
+  // ->
+  // %rb:RegBank(s32) = G_ ...
+  // %rc:RegClass(s32) = COPY %rb
+  // ...
+  // %a = G_ ..., %rb
+  void reAssignRegBankOnDef(MachineInstr &MI, MachineOperand &DefOP,
+                            const RegisterBank *RB) {
+    // Register that already has Register class got it during pre-inst selection
+    // of another instruction. Maybe cross bank copy was required so we insert a
+    // copy that can be removed later. This simplifies post regbanklegalize
+    // combiner and avoids need to special case some patterns.
+    Register Reg = DefOP.getReg();
+    LLT Ty = MRI.getType(Reg);
+    Register NewReg = MRI.createVirtualRegister({RB, Ty});
+    DefOP.setReg(NewReg);
+
+    auto &MBB = *MI.getParent();
+    B.setInsertPt(MBB, MBB.SkipPHIsAndLabels(std::next(MI.getIterator())));
+    B.buildCopy(Reg, NewReg);
+
+    // The problem was discovered for uniform S1 that was used as both
+    // lane mask(vcc) and regular sgpr S1.
+    // - lane-mask(vcc) use was by si_if, this use is divergent and requires
+    //   non-trivial sgpr-S1-to-vcc copy. But pre-inst-selection of si_if sets
+    //   sreg_64_xexec(S1) on def of uniform S1 making it lane-mask.
+    // - the regular sgpr S1(uniform) instruction is now broken since
+    //   it uses sreg_64_xexec(S1) which is divergent.
+
+    // Replace virtual registers with register class on generic instructions
+    // uses with virtual registers with register bank.
+    for (auto &UseMI : make_early_inc_range(MRI.use_instructions(Reg))) {
+      if (UseMI.isPreISelOpcode()) {
+        for (MachineOperand &Op : UseMI.operands()) {
+          if (Op.isReg() && Op.getReg() == Reg)
+            Op.setReg(NewReg);
+        }
+      }
+    }
+  }
+
+  // %a = G_ ..., %rc
+  // ->
+  // %rb:RegBank(s32) = COPY %rc
+  // %a = G_ ..., %rb
+  void constrainRegBankUse(MachineInstr &MI, MachineOperand &UseOP,
+                           const RegisterBank *RB) {
+    Register Reg = UseOP.getReg();
+
+    LLT Ty = MRI.getType(Reg);
+    Register NewReg = MRI.createVirtualRegister({RB, Ty});
+    UseOP.setReg(NewReg);
+
+    if (MI.isPHI()) {
+      auto DefMI = MRI.getVRegDef(Reg)->getIterator();
+      MachineBasicBlock *DefMBB = DefMI->getParent();
+      B.setInsertPt(*DefMBB, DefMBB->SkipPHIsAndLabels(std::next(DefMI)));
+    } else {
+      B.setInstr(MI);
+    }
+
+    B.buildCopy(NewReg, Reg);
+  }
+};
+
 bool AMDGPURegBankSelect::runOnMachineFunction(MachineFunction &MF) {
   if (MF.getProperties().hasProperty(
           MachineFunctionProperties::Property::FailedISel))
     return false;
+
+  // Setup the instruction builder with CSE.
+  const TargetPassConfig &TPC = getAnalysis<TargetPassConfig>();
+  GISelCSEAnalysisWrapper &Wrapper =
+      getAnalysis<GISelCSEAnalysisWrapperPass>().getCSEWrapper();
+  GISelCSEInfo &CSEInfo = Wrapper.get(TPC.getCSEConfig());
+  GISelObserverWrapper Observer;
+  Observer.addObserver(&CSEInfo);
+
+  CSEMIRBuilder B(MF);
+  B.setCSEInfo(&CSEInfo);
+  B.setChangeObserver(Observer);
+
+  RAIIDelegateInstaller DelegateInstaller(MF, &Observer);
+  RAIIMFObserverInstaller MFObserverInstaller(MF, Observer);
+
+  IntrinsicLaneMaskAnalyzer ILMA(MF);
+  MachineUniformityInfo &MUI =
+      getAnalysis<MachineUniformityAnalysisPass>().getUniformityInfo();
+  MachineRegisterInfo &MRI = *B.getMRI();
+  const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
+  RegBankSelectHelper RBSHelper(B, ILMA, MUI, *ST.getRegBankInfo());
+  // Virtual registers at this point don't have register banks.
+  // Virtual registers in def and use operands of already inst-selected
+  // instruction have register class.
+
+  for (MachineBasicBlock &MBB : MF) {
+    for (MachineInstr &MI : MBB) {
+      // Vregs in def and use operands of COPY can have either register class
+      // or bank. If there is neither on vreg in def operand, assign bank.
+      if (MI.isCopy()) {
+        Register DefReg = MI.getOperand(0).getReg();
+        assert(!MRI.getRegBankOrNull(DefReg));
+        if (!MRI.getRegClassOrNull(DefReg))
----------------
petar-avramovic wrote:

Small rework, still want to have both checks. Vregs on def before assigning banks should have either nothing or register class. Assert is there to check that there is no bank already assigned

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