[GlobalISel] Simplify RegBankSelect

[sebastian-ne]

Summary

Save the instruction list of a block before selecting banks.
This allows to cope with moved instructions, even if they are reordered
or splitted into multiple basic blocks.

[foad]

llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp

46

Don't need this since it's already done in the header.

734–735

Can you initialise the worklist from RPOT and then just have a single worklist loop, instead of two nested loops?

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

774–776

Before this patch, how did these new BBs get regbankselected?

[sebastian-ne]

Fix review comments, use just the worklist instead of a nested loop.

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

774–776

RegBankSelect had logic to detect if applyMapping moved an instruction to a new basic block and then continued from there.

The banks for instructions in RemainderBB and RestoreExecBB are set here, so they don’t need to be selected.

[foad]

OK, so if applyMapping creates new instructions it either has to

1. set banks for the new instructions' operands as they are created, or
2. do something to make sure RegBankSelect visits the new instructions.

[sebastian-ne]

Exactly

[foad]

The low level implementation details look OK to me.

I guess we need to get some agreement about the high level design: is this a good way of handling newly created instructions and BBs? In particular I wonder if there is some way of handling newly created instructions automatically, without having to call setNextInstruction, but I can't quite see how it would work.

llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp

698–702

Could just Worklist.insert(Worklist.end(), ...)?

[sebastian-ne]

Thanks, the Worklist.insert does read better than std::copy.

I guess we need to get some agreement about the high level design: is this a good way of handling newly created instructions and BBs? In particular I wonder if there is some way of handling newly created instructions automatically, without having to call setNextInstruction, but I can't quite see how it would work.

Yes, it gets quite complicated when introducing loops for call instructions because it moves multiple instructions and now also unbundles them (so the I++ used in RegBankSelect points to the wrong instruction – after the bundle).

[aemerson]

In D111223#3045476, @foad wrote:

The low level implementation details look OK to me.

I guess we need to get some agreement about the high level design: is this a good way of handling newly created instructions and BBs? In particular I wonder if there is some way of handling newly created instructions automatically, without having to call setNextInstruction, but I can't quite see how it would work.

Could an observer be installed to watch for newly created instructions?

[sebastian-ne]

I implemented it using an observer and it worked. Then I realized that I never registered the observer anywhere…
So, here’s a simpler implementation that seems to be enough to find all instructions.

The new patch for creating loops around call instructions does not use bundles anymore because they create more problems than they solve (e.g. the legalizer tries and fails to legalize instructions in a bundle). This simplifies the cases that need to be handled.

[foad]

llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp

710

Pass (MBB->instr_rbegin(), MBB->instr_rend()) into the constructor here?

716

I think there's a pop_back_value you can use.

[sebastian-ne]

llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp

710

I just tried and couldn’t get it working. I guess the iterator returns a MachineInstr& and no pointer, but SmallVector<MachineInstr&> doesn’t work.

[sebastian-ne]

Use pop_back_val

[foad]

llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp

710

Ugh. This might work:

SmallVector<MachineInstr *> WorkList(make_pointer_range(reverse(MBB->instrs())));

[sebastian-ne]

The make_pointer_range worked, thanks

[arsenm]

Is the latest diff incomplete or is this really this simple now?

[sebastian-ne]

It’s complete. It worked in all cases I tested, so I hope it really can be that simple :)

llvm/include/llvm/CodeGen/GlobalISel/RegBankSelect.h

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
 #include "llvm/CodeGen/GlobalISel/RegisterBankInfo.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
 #include <cassert>
 #include <cstdint>
+#include <deque>
 #include <memory>
 
 namespace llvm {
 
 class BlockFrequency;
 class MachineBlockFrequencyInfo;
 class MachineBranchProbabilityInfo;
 class MachineOperand;
private:
   MachineIRBuilder MIRBuilder;
 
   /// Optimization mode of the pass.
   Mode OptMode;
 
   /// Current target configuration. Controls how the pass handles errors.
   const TargetPassConfig *TPC;
 
+  std::deque<MachineBasicBlock *> Worklist;
+  MachineBasicBlock::iterator MII;
+  MachineBasicBlock::iterator End;
+
   /// Assign the register bank of each operand of \p MI.
   /// \return True on success, false otherwise.
   bool assignInstr(MachineInstr &MI);
 
   /// Initialize the field members using \p MF.
   void init(MachineFunction &MF);
 
   /// Check if \p Reg is already assigned what is described by \p ValMapping.
public:
   ///       for each argument in inst
   ///         if (CurRegBank != argument.RegBank)
   ///           ArgReg = argument.getReg()
   ///           Tmp = MRI.createNewVirtual(MRI.getSize(ArgReg), CurRegBank)
   ///           MIRBuilder.buildInstr(COPY, Tmp, ArgReg)
   ///           inst.getOperand(argument.getOperandNo()).setReg(Tmp)
   /// \endcode
   bool runOnMachineFunction(MachineFunction &MF) override;
+
+  /// Add a newly created basic block at the front of the worklist.
+  void prependBBToWorkQueue(MachineBasicBlock *BB);
+  /// Set the next instruction that should be assigned register banks.
+  void setNextInstruction(MachineBasicBlock::iterator MII);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_GLOBALISEL_REGBANKSELECT_H

llvm/include/llvm/CodeGen/GlobalISel/RegisterBankInfo.h

 #include <initializer_list>
 #include <memory>
 
 namespace llvm {
 
 class MachineInstr;
 class MachineRegisterInfo;
 class raw_ostream;
+class RegBankSelect;
 class RegisterBank;
 class TargetInstrInfo;
 class TargetRegisterClass;
 class TargetRegisterInfo;
 
 /// Holds all the information related to register banks.
 class RegisterBankInfo {
 public:
public:
   /// If the mapping of one of the operand spans several registers, this
   /// method will abort as this is not like a default mapping anymore.
   ///
   /// \pre For OpIdx in {0..\p OpdMapper.getMI().getNumOperands())
   ///        the range OpdMapper.getVRegs(OpIdx) is empty or of size 1.
   static void applyDefaultMapping(const OperandsMapper &OpdMapper);
 
   /// See ::applyMapping.
-  virtual void applyMappingImpl(const OperandsMapper &OpdMapper) const {
+  virtual void applyMappingImpl(RegBankSelect &RegBankSelectPass,
+                                const OperandsMapper &OpdMapper) const {
     llvm_unreachable("The target has to implement that part");
   }
 
 public:
   virtual ~RegisterBankInfo() = default;
 
   /// Get the register bank identified by \p ID.
   const RegisterBank &getRegBank(unsigned ID) const {
public:
   /// Apply \p OpdMapper.getInstrMapping() to \p OpdMapper.getMI().
   /// After this call \p OpdMapper.getMI() may not be valid anymore.
   /// \p OpdMapper.getInstrMapping().getID() carries the information of
   /// what has been chosen to map \p OpdMapper.getMI(). This ID is set
   /// by the various getInstrXXXMapping method.
   ///
   /// Therefore, getting the mapping and applying it should be kept in
   /// sync.
-  void applyMapping(const OperandsMapper &OpdMapper) const {
+  void applyMapping(RegBankSelect &RegBankSelectPass,
+                    const OperandsMapper &OpdMapper) const {
     // The only mapping we know how to handle is the default mapping.
     if (OpdMapper.getInstrMapping().getID() == DefaultMappingID)
       return applyDefaultMapping(OpdMapper);
     // For other mapping, the target needs to do the right thing.
     // If that means calling applyDefaultMapping, fine, but this
     // must be explicitly stated.
-    applyMappingImpl(OpdMapper);
+    applyMappingImpl(RegBankSelectPass, OpdMapper);
   }
 
   /// Get the size in bits of \p Reg.
   /// Utility method to get the size of any registers. Unlike
   /// MachineRegisterInfo::getSize, the register does not need to be a
   /// virtual register.
   ///
   /// \pre \p Reg != 0 (NoRegister).

llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp

 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <limits>
 #include <memory>
 #include <utility>
 
 #define DEBUG_TYPE "regbankselect"
 
 using namespace llvm;
 
 static cl::opt<RegBankSelect::Mode> RegBankSelectMode(
case RepairingPlacement::Insert:
       break;
     default:
       llvm_unreachable("Other kind should not happen");
     }
   }
 
   // Second, rewrite the instruction.
   LLVM_DEBUG(dbgs() << "Actual mapping of the operands: " << OpdMapper << '\n');
-  RBI->applyMapping(OpdMapper);
+  RBI->applyMapping(*this, OpdMapper);
 
   return true;
 }
 
 bool RegBankSelect::assignInstr(MachineInstr &MI) {
   LLVM_DEBUG(dbgs() << "Assign: " << MI);
 
   unsigned Opc = MI.getOpcode();
if (const MachineInstr *MI = machineFunctionIsIllegal(MF)) {
       return false;
     }
 #endif
 
   // Walk the function and assign register banks to all operands.
   // Use a RPOT to make sure all registers are assigned before we choose
   // the best mapping of the current instruction.
   ReversePostOrderTraversal<MachineFunction*> RPOT(&MF);
-  for (MachineBasicBlock *MBB : RPOT) {
+  Worklist.insert(Worklist.end(), RPOT.begin(), RPOT.end());
+  while (!Worklist.empty()) {
+    MachineBasicBlock *MBB = Worklist.front();
+    Worklist.pop_front();
+
     // Set a sensible insertion point so that subsequent calls to
     // MIRBuilder.
     MIRBuilder.setMBB(*MBB);
-    for (MachineBasicBlock::iterator MII = MBB->begin(), End = MBB->end();
-         MII != End;) {
+    End = MBB->end();
+    for (MII = MBB->begin(); MII != End;) {
       // MI might be invalidated by the assignment, so move the
       // iterator before hand.
       MachineInstr &MI = *MII++;
 
       // Ignore target-specific post-isel instructions: they should use proper
       // regclasses.
       if (isTargetSpecificOpcode(MI.getOpcode()) && !MI.isPreISelOpcode())
         continue;
 
       // Ignore inline asm instructions: they should use physical
       // registers/regclasses
       if (MI.isInlineAsm())
         continue;
 
       // Ignore debug info.
       if (MI.isDebugInstr())
         continue;
 
       // Ignore IMPLICIT_DEF which must have a regclass.
       if (MI.isImplicitDef())
         continue;
 
       if (!assignInstr(MI)) {
         reportGISelFailure(MF, *TPC, *MORE, "gisel-regbankselect",
                            "unable to map instruction", MI);
         return false;
       }
-
-      // It's possible the mapping changed control flow, and moved the following
-      // instruction to a new block, so figure out the new parent.
-      if (MII != End) {
-        MachineBasicBlock *NextInstBB = MII->getParent();
-        if (NextInstBB != MBB) {
-          LLVM_DEBUG(dbgs() << "Instruction mapping changed control flow\n");
-          MBB = NextInstBB;
-          MIRBuilder.setMBB(*MBB);
-          End = MBB->end();
-        }
-      }
     }
   }
 
   OptMode = SaveOptMode;
   return false;
 }
 
 //------------------------------------------------------------------------------
 //                  Helper Classes Implementation
if (*this == ImpossibleCost()) {
     return;
   }
   if (isSaturated()) {
     OS << "saturated";
     return;
   }
   OS << LocalFreq << " * " << LocalCost << " + " << NonLocalCost;
 }
+
+void RegBankSelect::prependBBToWorkQueue(MachineBasicBlock *BB) {
+  Worklist.push_front(BB);
+}
+
+void RegBankSelect::setNextInstruction(MachineBasicBlock::iterator MII) {
+  this->MII = MII;
+  MachineBasicBlock *NextInstBB = MII->getParent();
+  End = NextInstBB->end();
+  MIRBuilder.setMBB(*NextInstBB);
+}

llvm/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.h

 
 #define GET_TARGET_REGBANK_CLASS
 #include "AArch64GenRegisterBank.inc"
 };
 
 /// This class provides the information for the target register banks.
 class AArch64RegisterBankInfo final : public AArch64GenRegisterBankInfo {
   /// See RegisterBankInfo::applyMapping.
-  void applyMappingImpl(const OperandsMapper &OpdMapper) const override;
+  void applyMappingImpl(RegBankSelect &RegBankSelectPass,
+                        const OperandsMapper &OpdMapper) const override;
 
   /// Get an instruction mapping where all the operands map to
   /// the same register bank and have similar size.
   ///
   /// \pre MI.getNumOperands() <= 3
   ///
   /// \return An InstructionMappings with a statically allocated
   /// OperandsMapping.

llvm/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.cpp

AArch64RegisterBankInfo::getInstrAlternativeMappings(
   }
   default:
     break;
   }
   return RegisterBankInfo::getInstrAlternativeMappings(MI);
 }
 
 void AArch64RegisterBankInfo::applyMappingImpl(
-    const OperandsMapper &OpdMapper) const {
+    RegBankSelect &RegBankSelectPass, const OperandsMapper &OpdMapper) const {
   switch (OpdMapper.getMI().getOpcode()) {
   case TargetOpcode::G_OR:
   case TargetOpcode::G_BITCAST:
   case TargetOpcode::G_LOAD:
     // Those ID must match getInstrAlternativeMappings.
     assert((OpdMapper.getInstrMapping().getID() >= 1 &&
             OpdMapper.getInstrMapping().getID() <= 4) &&
            "Don't know how to handle that ID");

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.h

public:
   bool buildVCopy(MachineIRBuilder &B, Register DstReg, Register SrcReg) const;
 
   bool collectWaterfallOperands(
     SmallSet<Register, 4> &SGPROperandRegs,
     MachineInstr &MI,
     MachineRegisterInfo &MRI,
     ArrayRef<unsigned> OpIndices) const;
 
-  bool executeInWaterfallLoop(
+  bool executeInWaterfallLoop(RegBankSelect &RegBankSelectPass,
                               MachineIRBuilder &B,
                               iterator_range<MachineBasicBlock::iterator> Range,
                               SmallSet<Register, 4> &SGPROperandRegs,
                               MachineRegisterInfo &MRI) const;
 
-  bool executeInWaterfallLoop(MachineIRBuilder &B,
-                              MachineInstr &MI,
+  bool executeInWaterfallLoop(RegBankSelect &RegBankSelectPass,
+                              MachineIRBuilder &B, MachineInstr &MI,
                               MachineRegisterInfo &MRI,
                               ArrayRef<unsigned> OpIndices) const;
-  bool executeInWaterfallLoop(MachineInstr &MI,
-                              MachineRegisterInfo &MRI,
+  bool executeInWaterfallLoop(RegBankSelect &RegBankSelectPass,
+                              MachineInstr &MI, MachineRegisterInfo &MRI,
                               ArrayRef<unsigned> OpIndices) const;
 
   void constrainOpWithReadfirstlane(MachineInstr &MI, MachineRegisterInfo &MRI,
                                     unsigned OpIdx) const;
   bool applyMappingDynStackAlloc(MachineInstr &MI,
                                  const OperandsMapper &OpdMapper,
                                  MachineRegisterInfo &MRI) const;
   bool applyMappingLoad(MachineInstr &MI,
                         const OperandsMapper &OpdMapper,
                         MachineRegisterInfo &MRI) const;
-  bool
-  applyMappingImage(MachineInstr &MI,
+  bool applyMappingImage(RegBankSelect &RegBankSelectPass, MachineInstr &MI,
                          const OperandsMapper &OpdMapper,
                          MachineRegisterInfo &MRI, int RSrcIdx) const;
-  bool applyMappingSBufferLoad(const OperandsMapper &OpdMapper) const;
+  bool applyMappingSBufferLoad(RegBankSelect &RegBankSelectPass,
+                               const OperandsMapper &OpdMapper) const;
 
   bool applyMappingBFE(const OperandsMapper &OpdMapper, bool Signed) const;
 
   Register handleD16VData(MachineIRBuilder &B, MachineRegisterInfo &MRI,
                           Register Reg) const;
 
   std::pair<Register, unsigned>
   splitBufferOffsets(MachineIRBuilder &B, Register Offset) const;
 
-  MachineInstr *selectStoreIntrinsic(MachineIRBuilder &B,
+  MachineInstr *selectStoreIntrinsic(RegBankSelect &RegBankSelectPass,
+                                     MachineIRBuilder &B,
                                      MachineInstr &MI) const;
 
   /// See RegisterBankInfo::applyMapping.
-  void applyMappingImpl(const OperandsMapper &OpdMapper) const override;
+  void applyMappingImpl(RegBankSelect &RegBankSelectPass,
+                        const OperandsMapper &OpdMapper) const override;
 
   const ValueMapping *getValueMappingForPtr(const MachineRegisterInfo &MRI,
                                             Register Ptr) const;
 
   const RegisterBankInfo::InstructionMapping &
   getInstrMappingForLoad(const MachineInstr &MI) const;
 
   unsigned getRegBankID(Register Reg, const MachineRegisterInfo &MRI,

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

 #include "AMDGPUGlobalISelUtils.h"
 #include "AMDGPUInstrInfo.h"
 #include "GCNSubtarget.h"
 #include "SIMachineFunctionInfo.h"
 #include "SIRegisterInfo.h"
 #include "llvm/CodeGen/GlobalISel/LegalizerHelper.h"
 #include "llvm/CodeGen/GlobalISel/MIPatternMatch.h"
 #include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
+#include "llvm/CodeGen/GlobalISel/RegBankSelect.h"
 #include "llvm/CodeGen/GlobalISel/RegisterBank.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
 
 #define GET_TARGET_REGBANK_IMPL
 #include "AMDGPUGenRegisterBank.inc"
 
 // This file will be TableGen'ed at some point.
 #include "AMDGPUGenRegisterBankInfo.def"

 ///   SGPR = read SGPR value for current lane from VGPR
 ///   VGPRResult[Lane] = use_op SGPR
 /// }
 /// Restore Execution Mask
 ///
 /// There is additional complexity to try for compare values to identify the
 /// unique values used.
 bool AMDGPURegisterBankInfo::executeInWaterfallLoop(
-  MachineIRBuilder &B,
+    RegBankSelect &RegBankSelectPass, MachineIRBuilder &B,
     iterator_range<MachineBasicBlock::iterator> Range,
-  SmallSet<Register, 4> &SGPROperandRegs,
-  MachineRegisterInfo &MRI) const {
+    SmallSet<Register, 4> &SGPROperandRegs, MachineRegisterInfo &MRI) const {
   SmallVector<Register, 4> ResultRegs;
   SmallVector<Register, 4> InitResultRegs;
   SmallVector<Register, 4> PhiRegs;
 
   // Track use registers which have already been expanded with a readfirstlane
   // sequence. This may have multiple uses if moving a sequence.
   DenseMap<Register, Register> WaterfalledRegMap;
 
#endif
   MachineBasicBlock *LoopBB = MF->CreateMachineBasicBlock();
   MachineBasicBlock *RemainderBB = MF->CreateMachineBasicBlock();
   MachineBasicBlock *RestoreExecBB = MF->CreateMachineBasicBlock();
   MachineFunction::iterator MBBI(MBB);
   ++MBBI;
   MF->insert(MBBI, LoopBB);
   MF->insert(MBBI, RestoreExecBB);
   MF->insert(MBBI, RemainderBB);
+  RegBankSelectPass.prependBBToWorkQueue(RestoreExecBB);
+  RegBankSelectPass.prependBBToWorkQueue(RemainderBB);
+  RegBankSelectPass.prependBBToWorkQueue(LoopBB);
 
   LoopBB->addSuccessor(RestoreExecBB);
   LoopBB->addSuccessor(LoopBB);
 
   // Move the rest of the block into a new block.
   RemainderBB->transferSuccessorsAndUpdatePHIs(&MBB);
   RemainderBB->splice(RemainderBB->begin(), &MBB, Range.end(), MBB.end());
 
#endif
   B.buildInstr(MovTermOpc)
     .addDef(ExecReg)
     .addReg(SaveExecReg);
 
   // Set the insert point after the original instruction, so any new
   // instructions will be in the remainder.
   B.setInsertPt(*RemainderBB, RemainderBB->begin());
 
+  RegBankSelectPass.setNextInstruction(NewBegin);
+
   return true;
 }
 
 // Return any unique registers used by \p MI at \p OpIndices that need to be
 // handled in a waterfall loop. Returns these registers in \p
 // SGPROperandRegs. Returns true if there are any operands to handle and a
 // waterfall loop is necessary.
 bool AMDGPURegisterBankInfo::collectWaterfallOperands(
   SmallSet<Register, 4> &SGPROperandRegs, MachineInstr &MI,
   MachineRegisterInfo &MRI, ArrayRef<unsigned> OpIndices) const {
   for (unsigned Op : OpIndices) {
     assert(MI.getOperand(Op).isUse());
     Register Reg = MI.getOperand(Op).getReg();
     const RegisterBank *OpBank = getRegBank(Reg, MRI, *TRI);
     if (OpBank->getID() != AMDGPU::SGPRRegBankID)
       SGPROperandRegs.insert(Reg);
   }
 
   // No operands need to be replaced, so no need to loop.
   return !SGPROperandRegs.empty();
 }
 
 bool AMDGPURegisterBankInfo::executeInWaterfallLoop(
-  MachineIRBuilder &B, MachineInstr &MI, MachineRegisterInfo &MRI,
-  ArrayRef<unsigned> OpIndices) const {
+    RegBankSelect &RegBankSelectPass, MachineIRBuilder &B, MachineInstr &MI,
+    MachineRegisterInfo &MRI, ArrayRef<unsigned> OpIndices) const {
   // Use a set to avoid extra readfirstlanes in the case where multiple operands
   // are the same register.
   SmallSet<Register, 4> SGPROperandRegs;
 
   if (!collectWaterfallOperands(SGPROperandRegs, MI, MRI, OpIndices))
     return false;
 
   MachineBasicBlock::iterator I = MI.getIterator();
-  return executeInWaterfallLoop(B, make_range(I, std::next(I)),
-                                SGPROperandRegs, MRI);
+  return executeInWaterfallLoop(
+      RegBankSelectPass, B, make_range(I, std::next(I)), SGPROperandRegs, MRI);
 }
 
 bool AMDGPURegisterBankInfo::executeInWaterfallLoop(
-  MachineInstr &MI, MachineRegisterInfo &MRI,
-  ArrayRef<unsigned> OpIndices) const {
+    RegBankSelect &RegBankSelectPass, MachineInstr &MI,
+    MachineRegisterInfo &MRI, ArrayRef<unsigned> OpIndices) const {
   MachineIRBuilder B(MI);
-  return executeInWaterfallLoop(B, MI, MRI, OpIndices);
+  return executeInWaterfallLoop(RegBankSelectPass, B, MI, MRI, OpIndices);
 }
 
 // Legalize an operand that must be an SGPR by inserting a readfirstlane.
 void AMDGPURegisterBankInfo::constrainOpWithReadfirstlane(
     MachineInstr &MI, MachineRegisterInfo &MRI, unsigned OpIdx) const {
   Register Reg = MI.getOperand(OpIdx).getReg();
   const RegisterBank *Bank = getRegBank(Reg, MRI, *TRI);
   if (Bank == &AMDGPU::SGPRRegBank)
if (Alignment > TFI.getStackAlign()) {
     B.buildPtrAdd(Dst, SPCopy, ScaledSize);
   }
 
   MI.eraseFromParent();
   return true;
 }
 
 bool AMDGPURegisterBankInfo::applyMappingImage(
-    MachineInstr &MI, const AMDGPURegisterBankInfo::OperandsMapper &OpdMapper,
+    RegBankSelect &RegBankSelectPass, MachineInstr &MI,
+    const AMDGPURegisterBankInfo::OperandsMapper &OpdMapper,
     MachineRegisterInfo &MRI, int RsrcIdx) const {
   const int NumDefs = MI.getNumExplicitDefs();
 
   // The reported argument index is relative to the IR intrinsic call arguments,
   // so we need to shift by the number of defs and the intrinsic ID.
   RsrcIdx += NumDefs + 1;
 
   // Insert copies to VGPR arguments.
   applyDefaultMapping(OpdMapper);
 
   // Fixup any SGPR arguments.
   SmallVector<unsigned, 4> SGPRIndexes;
   for (int I = NumDefs, NumOps = MI.getNumOperands(); I != NumOps; ++I) {
     if (!MI.getOperand(I).isReg())
       continue;
 
     // If this intrinsic has a sampler, it immediately follows rsrc.
     if (I == RsrcIdx || I == RsrcIdx + 1)
       SGPRIndexes.push_back(I);
   }
 
-  executeInWaterfallLoop(MI, MRI, SGPRIndexes);
+  executeInWaterfallLoop(RegBankSelectPass, MI, MRI, SGPRIndexes);
   return true;
 }
 
 static Register getSrcRegIgnoringCopies(const MachineRegisterInfo &MRI,
                                         Register Reg) {
   MachineInstr *Def = getDefIgnoringCopies(Reg, MRI);
   if (!Def)
     return Reg;
static unsigned setBufferOffsets(MachineIRBuilder &B,
   }
 
   SOffsetReg = B.buildConstant(S32, 0).getReg(0);
   B.getMRI()->setRegBank(SOffsetReg, AMDGPU::SGPRRegBank);
   return 0;
 }
 
 bool AMDGPURegisterBankInfo::applyMappingSBufferLoad(
-  const OperandsMapper &OpdMapper) const {
+    RegBankSelect &RegBankSelectPass, const OperandsMapper &OpdMapper) const {
   MachineInstr &MI = OpdMapper.getMI();
   MachineRegisterInfo &MRI = OpdMapper.getMRI();
 
   const LLT S32 = LLT::scalar(32);
   Register Dst = MI.getOperand(0).getReg();
   LLT Ty = MRI.getType(Dst);
 
   const RegisterBank *RSrcBank =
if (RSrcBank != &AMDGPU::SGPRRegBank) {
     // Remove the original instruction to avoid potentially confusing the
     // waterfall loop logic.
     B.setInstr(*Span.begin());
     MI.eraseFromParent();
 
     SmallSet<Register, 4> OpsToWaterfall;
 
     OpsToWaterfall.insert(RSrc);
-    executeInWaterfallLoop(B, make_range(Span.begin(), Span.end()),
-                           OpsToWaterfall, MRI);
+    executeInWaterfallLoop(RegBankSelectPass, B,
+                           make_range(Span.begin(), Span.end()), OpsToWaterfall,
+                           MRI);
   }
 
   if (NumLoads != 1) {
     if (Ty.isVector())
       B.buildConcatVectors(Dst, LoadParts);
     else
       B.buildMerge(Dst, LoadParts);
   }

 static unsigned extractCPol(unsigned CachePolicy) {
   return CachePolicy & AMDGPU::CPol::ALL;
 }
 
 static unsigned extractSWZ(unsigned CachePolicy) {
   return (CachePolicy >> 3) & 1;
 }
 
-
 MachineInstr *
-AMDGPURegisterBankInfo::selectStoreIntrinsic(MachineIRBuilder &B,
+AMDGPURegisterBankInfo::selectStoreIntrinsic(RegBankSelect &RegBankSelectPass,
+                                             MachineIRBuilder &B,
                                              MachineInstr &MI) const {
-   MachineRegisterInfo &MRI = *B.getMRI();
-  executeInWaterfallLoop(B, MI, MRI, {2, 4});
+  executeInWaterfallLoop(RegBankSelectPass, B, MI, MRI, {2, 4});
 
   // FIXME: DAG lowering brokenly changes opcode based on FP vs. integer.
 
   Register VData = MI.getOperand(1).getReg();
   LLT Ty = MRI.getType(VData);
 
   int EltSize = Ty.getScalarSizeInBits();
   int Size = Ty.getSizeInBits();
bool AMDGPURegisterBankInfo::foldInsertEltToCmpSelect(
 
   MRI.setRegBank(MI.getOperand(0).getReg(), DstBank);
   MI.eraseFromParent();
 
   return true;
 }
 
 void AMDGPURegisterBankInfo::applyMappingImpl(
-    const OperandsMapper &OpdMapper) const {
+    RegBankSelect &RegBankSelectPass, const OperandsMapper &OpdMapper) const {
   MachineInstr &MI = OpdMapper.getMI();
   unsigned Opc = MI.getOpcode();
   MachineRegisterInfo &MRI = OpdMapper.getMRI();
   switch (Opc) {
   case AMDGPU::G_PHI: {
     Register DstReg = MI.getOperand(0).getReg();
     LLT DstTy = MRI.getType(DstReg);
     if (DstTy != LLT::scalar(1))
case AMDGPU::G_EXTRACT_VECTOR_ELT: {
     // actual indexing will be done on the SGPR source vector, which will
     // produce a scalar result. We need to copy to the VGPR result inside the
     // waterfall loop.
     const bool NeedCopyToVGPR = DstBank == &AMDGPU::VGPRRegBank &&
                                 SrcBank == &AMDGPU::SGPRRegBank;
     if (DstRegs.empty()) {
       applyDefaultMapping(OpdMapper);
 
-      executeInWaterfallLoop(MI, MRI, { 2 });
+      executeInWaterfallLoop(RegBankSelectPass, MI, MRI, {2});
 
       if (NeedCopyToVGPR) {
         // We don't want a phi for this temporary reg.
         Register TmpReg = MRI.createGenericVirtualRegister(DstTy);
         MRI.setRegBank(TmpReg, AMDGPU::SGPRRegBank);
         MI.getOperand(0).setReg(TmpReg);
         B.setInsertPt(*MI.getParent(), ++MI.getIterator());
 
if (!collectWaterfallOperands(OpsToWaterfall, MI, MRI, { 2 })) {
       MI.eraseFromParent();
       return;
     }
 
     // Remove the original instruction to avoid potentially confusing the
     // waterfall loop logic.
     B.setInstr(*Span.begin());
     MI.eraseFromParent();
-    executeInWaterfallLoop(B, make_range(Span.begin(), Span.end()),
-                           OpsToWaterfall, MRI);
+    executeInWaterfallLoop(RegBankSelectPass, B,
+                           make_range(Span.begin(), Span.end()), OpsToWaterfall,
+                           MRI);
 
     if (NeedCopyToVGPR) {
       MachineBasicBlock *LoopBB = Extract1->getParent();
       Register TmpReg0 = MRI.createGenericVirtualRegister(S32);
       Register TmpReg1 = MRI.createGenericVirtualRegister(S32);
       MRI.setRegBank(TmpReg0, AMDGPU::SGPRRegBank);
       MRI.setRegBank(TmpReg1, AMDGPU::SGPRRegBank);
 
bool ShouldMoveIndexIntoLoop = IdxBank != &AMDGPU::SGPRRegBank &&
       ConstOffset < VecTy.getNumElements();
 
     // Move the base register. We'll re-insert the add later.
     if (ShouldMoveIndexIntoLoop)
       MI.getOperand(3).setReg(BaseIdxReg);
 
 
     if (InsRegs.empty()) {
-      executeInWaterfallLoop(MI, MRI, { 3 });
+      executeInWaterfallLoop(RegBankSelectPass, MI, MRI, {3});
 
       // Re-insert the constant offset add inside the waterfall loop.
       if (ShouldMoveIndexIntoLoop) {
         MachineIRBuilder B(MI);
         reinsertVectorIndexAdd(B, MI, 3, ConstOffset);
       }
 
       return;
if (!collectWaterfallOperands(OpsToWaterfall, MI, MRI, { 3 })) {
       return;
     }
 
     B.setInstr(*Span.begin());
     MI.eraseFromParent();
 
     // Figure out the point after the waterfall loop before mangling the control
     // flow.
-    executeInWaterfallLoop(B, make_range(Span.begin(), Span.end()),
-                           OpsToWaterfall, MRI);
+    executeInWaterfallLoop(RegBankSelectPass, B,
+                           make_range(Span.begin(), Span.end()), OpsToWaterfall,
+                           MRI);
 
     // The insertion point is now right after the original instruction.
     //
     // Keep the bitcast to the original vector type out of the loop. Doing this
     // saved an extra phi we don't need inside the loop.
     B.buildBitcast(DstReg, InsHi);
 
     // Re-insert the constant offset add inside the waterfall loop.
void AMDGPURegisterBankInfo::applyMappingImpl(
   case AMDGPU::G_AMDGPU_BUFFER_STORE:
   case AMDGPU::G_AMDGPU_BUFFER_STORE_BYTE:
   case AMDGPU::G_AMDGPU_BUFFER_STORE_SHORT:
   case AMDGPU::G_AMDGPU_BUFFER_STORE_FORMAT:
   case AMDGPU::G_AMDGPU_BUFFER_STORE_FORMAT_D16:
   case AMDGPU::G_AMDGPU_TBUFFER_STORE_FORMAT:
   case AMDGPU::G_AMDGPU_TBUFFER_STORE_FORMAT_D16: {
     applyDefaultMapping(OpdMapper);
-    executeInWaterfallLoop(MI, MRI, {1, 4});
+    executeInWaterfallLoop(RegBankSelectPass, MI, MRI, {1, 4});
     return;
   }
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_SWAP:
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_ADD:
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_SUB:
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_SMIN:
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_UMIN:
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_SMAX:
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_UMAX:
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_AND:
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_OR:
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_XOR:
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_INC:
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_DEC: {
     applyDefaultMapping(OpdMapper);
-    executeInWaterfallLoop(MI, MRI, {2, 5});
+    executeInWaterfallLoop(RegBankSelectPass, MI, MRI, {2, 5});
     return;
   }
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_FADD:
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_FMIN:
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_FMAX: {
     applyDefaultMapping(OpdMapper);
-    executeInWaterfallLoop(MI, MRI, {2, 5});
+    executeInWaterfallLoop(RegBankSelectPass, MI, MRI, {2, 5});
     return;
   }
   case AMDGPU::G_AMDGPU_BUFFER_ATOMIC_CMPSWAP: {
     applyDefaultMapping(OpdMapper);
-    executeInWaterfallLoop(MI, MRI, {3, 6});
+    executeInWaterfallLoop(RegBankSelectPass, MI, MRI, {3, 6});
     return;
   }
   case AMDGPU::G_AMDGPU_S_BUFFER_LOAD: {
-    applyMappingSBufferLoad(OpdMapper);
+    applyMappingSBufferLoad(RegBankSelectPass, OpdMapper);
     return;
   }
   case AMDGPU::G_INTRINSIC: {
     switch (MI.getIntrinsicID()) {
     case Intrinsic::amdgcn_readlane: {
       substituteSimpleCopyRegs(OpdMapper, 2);
 
       assert(OpdMapper.getVRegs(0).empty());
void AMDGPURegisterBankInfo::applyMappingImpl(
   case AMDGPU::G_AMDGPU_INTRIN_IMAGE_LOAD:
   case AMDGPU::G_AMDGPU_INTRIN_IMAGE_STORE: {
     const AMDGPU::RsrcIntrinsic *RSrcIntrin
       = AMDGPU::lookupRsrcIntrinsic(MI.getIntrinsicID());
     assert(RSrcIntrin && RSrcIntrin->IsImage);
     // Non-images can have complications from operands that allow both SGPR
     // and VGPR. For now it's too complicated to figure out the final opcode
     // to derive the register bank from the MCInstrDesc.
-    applyMappingImage(MI, OpdMapper, MRI, RSrcIntrin->RsrcArg);
+    applyMappingImage(RegBankSelectPass, MI, OpdMapper, MRI,
+                      RSrcIntrin->RsrcArg);
     return;
   }
   case AMDGPU::G_AMDGPU_INTRIN_BVH_INTERSECT_RAY: {
     unsigned N = MI.getNumExplicitOperands() - 2;
     applyDefaultMapping(OpdMapper);
-    executeInWaterfallLoop(MI, MRI, { N });
+    executeInWaterfallLoop(RegBankSelectPass, MI, MRI, {N});
     return;
   }
   case AMDGPU::G_INTRINSIC_W_SIDE_EFFECTS: {
     auto IntrID = MI.getIntrinsicID();
     switch (IntrID) {
     case Intrinsic::amdgcn_ds_ordered_add:
     case Intrinsic::amdgcn_ds_ordered_swap: {
       // This is only allowed to execute with 1 lane, so readfirstlane is safe.
case AMDGPU::G_INTRINSIC_W_SIDE_EFFECTS: {
     }
     default: {
       if (const AMDGPU::RsrcIntrinsic *RSrcIntrin =
               AMDGPU::lookupRsrcIntrinsic(IntrID)) {
         // Non-images can have complications from operands that allow both SGPR
         // and VGPR. For now it's too complicated to figure out the final opcode
         // to derive the register bank from the MCInstrDesc.
         if (RSrcIntrin->IsImage) {
-          applyMappingImage(MI, OpdMapper, MRI, RSrcIntrin->RsrcArg);
+          applyMappingImage(RegBankSelectPass, MI, OpdMapper, MRI,
+                            RSrcIntrin->RsrcArg);
           return;
         }
       }
 
       break;
     }
     }
     break;

llvm/lib/Target/Mips/MipsRegisterBankInfo.h

public:
 
   const InstructionMapping &
   getInstrMapping(const MachineInstr &MI) const override;
 
   /// Here we have to narrowScalar s64 operands to s32, combine away G_MERGE or
   /// G_UNMERGE and erase instructions that became dead in the process. We
   /// manually assign bank to def operand of all new instructions that were
   /// created in the process since they will not end up in RegBankSelect loop.
-  void applyMappingImpl(const OperandsMapper &OpdMapper) const override;
+  void applyMappingImpl(RegBankSelect &RegBankSelectPass,
+                        const OperandsMapper &OpdMapper) const override;
 
   /// RegBankSelect determined that s64 operand is better to be split into two
   /// s32 operands in gprb. Here we manually set register banks of def operands
   /// of newly created instructions since they will not get regbankselected.
   void setRegBank(MachineInstr &MI, MachineRegisterInfo &MRI) const;
 
 private:
   /// Some instructions are used with both floating point and integer operands.

llvm/lib/Target/Mips/MipsRegisterBankInfo.cpp

combineAwayG_UNMERGE_VALUES(LegalizationArtifactCombiner &ArtCombiner,
   SmallVector<MachineInstr *, 2> DeadInstrs;
   ArtCombiner.tryCombineUnmergeValues(MI, DeadInstrs,
                                       UpdatedDefs, Observer);
   for (MachineInstr *DeadMI : DeadInstrs)
     DeadMI->eraseFromParent();
 }
 
 void MipsRegisterBankInfo::applyMappingImpl(
-    const OperandsMapper &OpdMapper) const {
+    RegBankSelect &RegBankSelectPass, const OperandsMapper &OpdMapper) const {
   MachineInstr &MI = OpdMapper.getMI();
   InstListTy NewInstrs;
   MachineFunction *MF = MI.getMF();
   MachineRegisterInfo &MRI = OpdMapper.getMRI();
   const LegalizerInfo &LegInfo = *MF->getSubtarget().getLegalizerInfo();
 
   InstManager NewInstrObserver(NewInstrs);
   MachineIRBuilder B(MI, NewInstrObserver);

llvm/lib/Target/X86/X86RegisterBankInfo.h

public:
 
   const RegisterBank &getRegBankFromRegClass(const TargetRegisterClass &RC,
                                              LLT) const override;
 
   InstructionMappings
   getInstrAlternativeMappings(const MachineInstr &MI) const override;
 
   /// See RegisterBankInfo::applyMapping.
-  void applyMappingImpl(const OperandsMapper &OpdMapper) const override;
+  void applyMappingImpl(RegBankSelect &RegBankSelectPass,
+                        const OperandsMapper &OpdMapper) const override;
 
   const InstructionMapping &
   getInstrMapping(const MachineInstr &MI) const override;
 };
 
 } // namespace llvm
 #endif

llvm/lib/Target/X86/X86RegisterBankInfo.cpp

X86RegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
   if (!getInstrValueMapping(MI, OpRegBankIdx, OpdsMapping))
     return getInvalidInstructionMapping();
 
   return getInstructionMapping(DefaultMappingID, /* Cost */ 1,
                                getOperandsMapping(OpdsMapping), NumOperands);
 }
 
 void X86RegisterBankInfo::applyMappingImpl(
-    const OperandsMapper &OpdMapper) const {
+    RegBankSelect &RegBankSelectPass, const OperandsMapper &OpdMapper) const {
   return applyDefaultMapping(OpdMapper);
 }
 
 RegisterBankInfo::InstructionMappings
 X86RegisterBankInfo::getInstrAlternativeMappings(const MachineInstr &MI) const {
 
   const MachineFunction &MF = *MI.getParent()->getParent();
   const TargetSubtargetInfo &STI = MF.getSubtarget();