[MachineRegisterInfo] Track register bank for virtual registers.

A virtual register may have either a register bank or a register class.
This is represented by a PointerUnion between the related classes.

Typically, a virtual register went through the following states
regarding register class and register bank:

1. Creation: None is set. Virtual registers are fully generic.
2. Register bank assignment: Register bank is set. Virtual registers
live into a register bank, but we do not know the constraints they need
to fulfil.
3. Instruction selection: Register class is set. Virtual registers are
bound by encoding constraints.

To map these states to GlobalISel, the IRTranslator implements #1,
RegBankSelect #2, and Select #3.

llvm-svn: 265696

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

@@ -16,7 +16,10 @@
 
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/IndexedMap.h"
+#include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/iterator_range.h"
+// PointerUnion needs to have access to the full RegisterBank type.
+#include "llvm/CodeGen/GlobalISel/RegisterBank.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBundle.h"
 #include "llvm/Target/TargetRegisterInfo.h"
@@ -26,6 +29,10 @@
 namespace llvm {
 class PSetIterator;
 
+/// Convenient type to represent either a register class or a register bank.
+typedef PointerUnion<const TargetRegisterClass *, const RegisterBank *>
+    RegClassOrRegBank;
+
 /// MachineRegisterInfo - Keep track of information for virtual and physical
 /// registers, including vreg register classes, use/def chains for registers,
 /// etc.
@@ -55,8 +62,9 @@ private:
   ///
   /// Each element in this list contains the register class of the vreg and the
   /// start of the use/def list for the register.
-  IndexedMap<std::pair<const TargetRegisterClass*, MachineOperand*>,
-             VirtReg2IndexFunctor> VRegInfo;
+  IndexedMap<std::pair<RegClassOrRegBank, MachineOperand *>,
+             VirtReg2IndexFunctor>
+      VRegInfo;
 
   /// RegAllocHints - This vector records register allocation hints for virtual
   /// registers. For each virtual register, it keeps a register and hint type
@@ -564,9 +572,51 @@ public:
   // Virtual Register Info
   //===--------------------------------------------------------------------===//
 
-  /// getRegClass - Return the register class of the specified virtual register.
+  /// Return the register class of the specified virtual register.
+  /// This shouldn't be used directly unless \p Reg has a register class.
+  /// \see getRegClassOrNull when this might happen.
   ///
   const TargetRegisterClass *getRegClass(unsigned Reg) const {
+    assert(VRegInfo[Reg].first.is<const TargetRegisterClass *>() &&
+           "Register class not set, wrong accessor");
+    return VRegInfo[Reg].first.get<const TargetRegisterClass *>();
+  }
+
+  /// Return the register class of \p Reg, or null if Reg has not been assigned
+  /// a register class yet.
+  ///
+  /// \note A null register class can only happen when these two
+  /// conditions are met:
+  /// 1. Generic virtual registers are created.
+  /// 2. The machine function has not completely been through the
+  ///    instruction selection process.
+  /// None of this condition is possible without GlobalISel for now.
+  /// In other words, if GlobalISel is not used or if the query happens after
+  /// the select pass, using getRegClass is safe.
+  const TargetRegisterClass *getRegClassOrNull(unsigned Reg) const {
+    const RegClassOrRegBank &Val = VRegInfo[Reg].first;
+    if (Val.is<const TargetRegisterClass *>())
+      return Val.get<const TargetRegisterClass *>();
+    return nullptr;
+  }
+
+  /// Return the register bank of \p Reg, or null if Reg has not been assigned
+  /// a register bank or has been assigned a register class.
+  /// \note It is possible to get the register bank from the register class via
+  /// RegisterBankInfo::getRegBankFromRegClass.
+  ///
+  const RegisterBank *getRegBankOrNull(unsigned Reg) const {
+    const RegClassOrRegBank &Val = VRegInfo[Reg].first;
+    if (Val.is<const RegisterBank *>())
+      return Val.get<const RegisterBank *>();
+    return nullptr;
+  }
+
+  /// Return the register bank or register class of \p Reg.
+  /// \note Before the register bank gets assigned (i.e., before the
+  /// RegBankSelect pass) \p Reg may not have either.
+  ///
+  const RegClassOrRegBank &getRegClassOrRegBank(unsigned Reg) const {
     return VRegInfo[Reg].first;
   }
 
@@ -574,6 +624,10 @@ public:
   ///
   void setRegClass(unsigned Reg, const TargetRegisterClass *RC);
 
+  /// Set the register bank to \p RegBank for \p Reg.
+  ///
+  void setRegBank(unsigned Reg, const RegisterBank &RegBank);
+
   /// constrainRegClass - Constrain the register class of the specified virtual
   /// register to be a common subclass of RC and the current register class,
   /// but only if the new class has at least MinNumRegs registers.  Return the

llvm/lib/CodeGen/MachineRegisterInfo.cpp

@@ -42,6 +42,11 @@ MachineRegisterInfo::setRegClass(unsigned Reg, const TargetRegisterClass *RC) {
   VRegInfo[Reg].first = RC;
 }
 
+void MachineRegisterInfo::setRegBank(unsigned Reg,
+                                     const RegisterBank &RegBank) {
+  VRegInfo[Reg].first = &RegBank;
+}
+
 const TargetRegisterClass *
 MachineRegisterInfo::constrainRegClass(unsigned Reg,
                                        const TargetRegisterClass *RC,
@@ -121,7 +126,7 @@ MachineRegisterInfo::createGenericVirtualRegister(unsigned Size) {
   unsigned Reg = TargetRegisterInfo::index2VirtReg(getNumVirtRegs());
   VRegInfo.grow(Reg);
   // FIXME: Should we use a dummy register class?
-  VRegInfo[Reg].first = nullptr;
+  VRegInfo[Reg].first = static_cast<TargetRegisterClass *>(nullptr);
   getVRegToSize()[Reg] = Size;
   RegAllocHints.grow(Reg);
   if (TheDelegate)