[llvm] [NewPM][CodeGen] Port VirtRegMap to NPM (PR #109936)

[Christudasan Devadasan via llvm-commits]

================
@@ -30,171 +31,210 @@ class MachineRegisterInfo;
 class raw_ostream;
 class TargetInstrInfo;
 
-  class VirtRegMap : public MachineFunctionPass {
-    MachineRegisterInfo *MRI = nullptr;
-    const TargetInstrInfo *TII = nullptr;
-    const TargetRegisterInfo *TRI = nullptr;
-    MachineFunction *MF = nullptr;
-
-    /// Virt2PhysMap - This is a virtual to physical register
-    /// mapping. Each virtual register is required to have an entry in
-    /// it; even spilled virtual registers (the register mapped to a
-    /// spilled register is the temporary used to load it from the
-    /// stack).
-    IndexedMap<MCRegister, VirtReg2IndexFunctor> Virt2PhysMap;
+class VirtRegMap {
+  MachineRegisterInfo *MRI = nullptr;
+  const TargetInstrInfo *TII = nullptr;
+  const TargetRegisterInfo *TRI = nullptr;
+  MachineFunction *MF = nullptr;
+
+  /// Virt2PhysMap - This is a virtual to physical register
+  /// mapping. Each virtual register is required to have an entry in
+  /// it; even spilled virtual registers (the register mapped to a
+  /// spilled register is the temporary used to load it from the
+  /// stack).
+  IndexedMap<MCRegister, VirtReg2IndexFunctor> Virt2PhysMap;
+
+  /// Virt2StackSlotMap - This is virtual register to stack slot
+  /// mapping. Each spilled virtual register has an entry in it
+  /// which corresponds to the stack slot this register is spilled
+  /// at.
+  IndexedMap<int, VirtReg2IndexFunctor> Virt2StackSlotMap;
+
+  /// Virt2SplitMap - This is virtual register to splitted virtual register
+  /// mapping.
+  IndexedMap<Register, VirtReg2IndexFunctor> Virt2SplitMap;
+
+  /// Virt2ShapeMap - For X86 AMX register whose register is bound shape
+  /// information.
+  DenseMap<Register, ShapeT> Virt2ShapeMap;
+
+  /// createSpillSlot - Allocate a spill slot for RC from MFI.
+  unsigned createSpillSlot(const TargetRegisterClass *RC);
+
+public:
+  static constexpr int NO_STACK_SLOT = INT_MAX;
+
+  VirtRegMap() : Virt2StackSlotMap(NO_STACK_SLOT) {}
+  VirtRegMap(const VirtRegMap &) = delete;
+  VirtRegMap &operator=(const VirtRegMap &) = delete;
+  VirtRegMap(VirtRegMap &&) = default;
+
+  void init(MachineFunction &MF);
+
+  MachineFunction &getMachineFunction() const {
+    assert(MF && "getMachineFunction called before runOnMachineFunction");
+    return *MF;
+  }
 
-    /// Virt2StackSlotMap - This is virtual register to stack slot
-    /// mapping. Each spilled virtual register has an entry in it
-    /// which corresponds to the stack slot this register is spilled
-    /// at.
-    IndexedMap<int, VirtReg2IndexFunctor> Virt2StackSlotMap;
+  MachineRegisterInfo &getRegInfo() const { return *MRI; }
+  const TargetRegisterInfo &getTargetRegInfo() const { return *TRI; }
 
-    /// Virt2SplitMap - This is virtual register to splitted virtual register
-    /// mapping.
-    IndexedMap<Register, VirtReg2IndexFunctor> Virt2SplitMap;
+  void grow();
 
-    /// Virt2ShapeMap - For X86 AMX register whose register is bound shape
-    /// information.
-    DenseMap<Register, ShapeT> Virt2ShapeMap;
+  /// returns true if the specified virtual register is
+  /// mapped to a physical register
+  bool hasPhys(Register virtReg) const { return getPhys(virtReg).isValid(); }
 
-    /// createSpillSlot - Allocate a spill slot for RC from MFI.
-    unsigned createSpillSlot(const TargetRegisterClass *RC);
+  /// returns the physical register mapped to the specified
+  /// virtual register
+  MCRegister getPhys(Register virtReg) const {
+    assert(virtReg.isVirtual());
+    return Virt2PhysMap[virtReg];
+  }
 
-  public:
-    static char ID;
+  /// creates a mapping for the specified virtual register to
+  /// the specified physical register
+  void assignVirt2Phys(Register virtReg, MCPhysReg physReg);
 
-    static constexpr int NO_STACK_SLOT = INT_MAX;
+  bool isShapeMapEmpty() const { return Virt2ShapeMap.empty(); }
 
-    VirtRegMap() : MachineFunctionPass(ID), Virt2StackSlotMap(NO_STACK_SLOT) {}
-    VirtRegMap(const VirtRegMap &) = delete;
-    VirtRegMap &operator=(const VirtRegMap &) = delete;
+  bool hasShape(Register virtReg) const {
+    return Virt2ShapeMap.contains(virtReg);
+  }
 
-    bool runOnMachineFunction(MachineFunction &MF) override;
+  ShapeT getShape(Register virtReg) const {
+    assert(virtReg.isVirtual());
+    return Virt2ShapeMap.lookup(virtReg);
+  }
 
-    void getAnalysisUsage(AnalysisUsage &AU) const override {
-      AU.setPreservesAll();
-      MachineFunctionPass::getAnalysisUsage(AU);
-    }
+  void assignVirt2Shape(Register virtReg, ShapeT shape) {
+    Virt2ShapeMap[virtReg] = shape;
+  }
 
-    MachineFunction &getMachineFunction() const {
-      assert(MF && "getMachineFunction called before runOnMachineFunction");
-      return *MF;
-    }
+  /// clears the specified virtual register's, physical
+  /// register mapping
+  void clearVirt(Register virtReg) {
+    assert(virtReg.isVirtual());
+    assert(Virt2PhysMap[virtReg] &&
+           "attempt to clear a not assigned virtual register");
+    Virt2PhysMap[virtReg] = MCRegister();
+  }
 
-    MachineRegisterInfo &getRegInfo() const { return *MRI; }
-    const TargetRegisterInfo &getTargetRegInfo() const { return *TRI; }
+  /// clears all virtual to physical register mappings
+  void clearAllVirt() {
+    Virt2PhysMap.clear();
+    grow();
+  }
 
-    void grow();
+  /// returns true if VirtReg is assigned to its preferred physreg.
+  bool hasPreferredPhys(Register VirtReg) const;
 
-    /// returns true if the specified virtual register is
-    /// mapped to a physical register
-    bool hasPhys(Register virtReg) const { return getPhys(virtReg).isValid(); }
+  /// returns true if VirtReg has a known preferred register.
+  /// This returns false if VirtReg has a preference that is a virtual
+  /// register that hasn't been assigned yet.
+  bool hasKnownPreference(Register VirtReg) const;
 
-    /// returns the physical register mapped to the specified
-    /// virtual register
-    MCRegister getPhys(Register virtReg) const {
-      assert(virtReg.isVirtual());
-      return Virt2PhysMap[virtReg];
+  /// records virtReg is a split live interval from SReg.
+  void setIsSplitFromReg(Register virtReg, Register SReg) {
+    Virt2SplitMap[virtReg] = SReg;
+    if (hasShape(SReg)) {
+      Virt2ShapeMap[virtReg] = getShape(SReg);
     }
+  }
 
-    /// creates a mapping for the specified virtual register to
-    /// the specified physical register
-    void assignVirt2Phys(Register virtReg, MCPhysReg physReg);
+  /// returns the live interval virtReg is split from.
+  Register getPreSplitReg(Register virtReg) const {
+    return Virt2SplitMap[virtReg];
+  }
 
-    bool isShapeMapEmpty() const { return Virt2ShapeMap.empty(); }
+  /// getOriginal - Return the original virtual register that VirtReg descends
+  /// from through splitting.
+  /// A register that was not created by splitting is its own original.
+  /// This operation is idempotent.
+  Register getOriginal(Register VirtReg) const {
+    Register Orig = getPreSplitReg(VirtReg);
+    return Orig ? Orig : VirtReg;
+  }
 
-    bool hasShape(Register virtReg) const {
-      return Virt2ShapeMap.contains(virtReg);
-    }
+  /// returns true if the specified virtual register is not
+  /// mapped to a stack slot or rematerialized.
+  bool isAssignedReg(Register virtReg) const {
+    if (getStackSlot(virtReg) == NO_STACK_SLOT)
+      return true;
+    // Split register can be assigned a physical register as well as a
+    // stack slot or remat id.
+    return (Virt2SplitMap[virtReg] && Virt2PhysMap[virtReg]);
+  }
 
-    ShapeT getShape(Register virtReg) const {
-      assert(virtReg.isVirtual());
-      return Virt2ShapeMap.lookup(virtReg);
-    }
+  /// returns the stack slot mapped to the specified virtual
+  /// register
+  int getStackSlot(Register virtReg) const {
+    assert(virtReg.isVirtual());
+    return Virt2StackSlotMap[virtReg];
+  }
 
-    void assignVirt2Shape(Register virtReg, ShapeT shape) {
-      Virt2ShapeMap[virtReg] = shape;
-    }
+  /// create a mapping for the specifed virtual register to
+  /// the next available stack slot
+  int assignVirt2StackSlot(Register virtReg);
 
-    /// clears the specified virtual register's, physical
-    /// register mapping
-    void clearVirt(Register virtReg) {
-      assert(virtReg.isVirtual());
-      assert(Virt2PhysMap[virtReg] &&
-             "attempt to clear a not assigned virtual register");
-      Virt2PhysMap[virtReg] = MCRegister();
-    }
+  /// create a mapping for the specified virtual register to
+  /// the specified stack slot
+  void assignVirt2StackSlot(Register virtReg, int SS);
 
-    /// clears all virtual to physical register mappings
-    void clearAllVirt() {
-      Virt2PhysMap.clear();
-      grow();
-    }
+  void print(raw_ostream &OS, const Module *M = nullptr) const;
+  void dump() const;
+};
 
-    /// returns true if VirtReg is assigned to its preferred physreg.
-    bool hasPreferredPhys(Register VirtReg) const;
+inline raw_ostream &operator<<(raw_ostream &OS, const VirtRegMap &VRM) {
+  VRM.print(OS);
+  return OS;
+}
 
-    /// returns true if VirtReg has a known preferred register.
-    /// This returns false if VirtReg has a preference that is a virtual
-    /// register that hasn't been assigned yet.
-    bool hasKnownPreference(Register VirtReg) const;
+class VirtRegMapWrapperPass : public MachineFunctionPass {
----------------
cdevadas wrote:

Unfortunately, most of the analysis passes when ported to NPM had ended up with PassPass in the initialize* functions. We should have followed 'WrapperLegacy' everywhere.

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