Index: llvm/include/llvm/CodeGen/GlobalISel/RegBankSelect.h
===================================================================
--- llvm/include/llvm/CodeGen/GlobalISel/RegBankSelect.h
+++ llvm/include/llvm/CodeGen/GlobalISel/RegBankSelect.h
@@ -407,7 +407,7 @@
     }
   };
 
-private:
+protected:
   /// Helper class used to represent the cost for mapping an instruction.
   /// When mapping an instruction, we may introduce some repairing code.
   /// In most cases, the repairing code is local to the instruction,
@@ -639,6 +639,12 @@
       .set(MachineFunctionProperties::Property::NoPHIs);
   }
 
+  /// Check that our input is fully legal: we require the function to have the
+  /// Legalized property, so it should be.
+  ///
+  /// FIXME: This should be in the MachineVerifier.
+  bool checkFunctionIsLegal(MachineFunction &MF) const;
+
   /// Walk through \p MF and assign a register bank to every virtual register
   /// that are still mapped to nothing.
   /// The target needs to provide a RegisterBankInfo and in particular
@@ -662,6 +668,8 @@
   ///           MIRBuilder.buildInstr(COPY, Tmp, ArgReg)
   ///           inst.getOperand(argument.getOperandNo()).setReg(Tmp)
   /// \endcode
+  bool assignRegisterBanks(MachineFunction &MF);
+
   bool runOnMachineFunction(MachineFunction &MF) override;
 };
 
Index: llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp
===================================================================
--- llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp
+++ llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp
@@ -674,31 +674,7 @@
   return applyMapping(MI, *BestMapping, RepairPts);
 }
 
-bool RegBankSelect::runOnMachineFunction(MachineFunction &MF) {
-  // If the ISel pipeline failed, do not bother running that pass.
-  if (MF.getProperties().hasProperty(
-          MachineFunctionProperties::Property::FailedISel))
-    return false;
-
-  LLVM_DEBUG(dbgs() << "Assign register banks for: " << MF.getName() << '\n');
-  const Function &F = MF.getFunction();
-  Mode SaveOptMode = OptMode;
-  if (F.hasOptNone())
-    OptMode = Mode::Fast;
-  init(MF);
-
-#ifndef NDEBUG
-  // Check that our input is fully legal: we require the function to have the
-  // Legalized property, so it should be.
-  // FIXME: This should be in the MachineVerifier.
-  if (!DisableGISelLegalityCheck)
-    if (const MachineInstr *MI = machineFunctionIsIllegal(MF)) {
-      reportGISelFailure(MF, *TPC, *MORE, "gisel-regbankselect",
-                         "instruction is not legal", *MI);
-      return false;
-    }
-#endif
-
+bool RegBankSelect::assignRegisterBanks(MachineFunction &MF) {
   // 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.
@@ -735,6 +711,43 @@
     }
   }
 
+  return true;
+}
+
+bool RegBankSelect::checkFunctionIsLegal(MachineFunction &MF) const {
+#ifndef NDEBUG
+  if (!DisableGISelLegalityCheck) {
+    if (const MachineInstr *MI = machineFunctionIsIllegal(MF)) {
+      reportGISelFailure(MF, *TPC, *MORE, "gisel-regbankselect",
+                         "instruction is not legal", *MI);
+      return false;
+    }
+  }
+#endif
+
+  return true;
+}
+
+bool RegBankSelect::runOnMachineFunction(MachineFunction &MF) {
+  // If the ISel pipeline failed, do not bother running that pass.
+  if (MF.getProperties().hasProperty(
+          MachineFunctionProperties::Property::FailedISel))
+    return false;
+
+  LLVM_DEBUG(dbgs() << "Assign register banks for: " << MF.getName() << '\n');
+  const Function &F = MF.getFunction();
+  Mode SaveOptMode = OptMode;
+  if (F.hasOptNone())
+    OptMode = Mode::Fast;
+  init(MF);
+
+#ifndef NDEBUG
+  if (!checkFunctionIsLegal(MF))
+    return false;
+#endif
+
+  assignRegisterBanks(MF);
+
   OptMode = SaveOptMode;
   return false;
 }