[CodeGen] Use delegate to notify targets when virtual registers are created

[cdevadas]

Summary

This will help targets to customize certain codegen decisions based on
the virtual registers involved in special operations. This patch also
extends the existing delegate in MRI to start supporting multicast.

[cdevadas]

See the use case in D134951.
Testcase: llvm/test/CodeGen/AMDGPU/wwm-register-spill-during-regalloc.ll

[cdevadas]

Ping.

[qcolombet]

Hi,

Given this is completely target specific, have you considered implementing that in your target specific MachineFunctionInfo pass?
Also, could your use case be supported by the existing target flags that can be attached to MachineOperands?

If we were to go down adding target flags for virtual registers, there are a few things that we would need to consider like what do we do with the flags when we merge live-ranges (e.g., in the coalescer), could a flag become invalid (e.g., when we do some peephole), etc.

I feel this patch is creating the llvm equivalent of metadata with all the problems we have with them and I would rather we avoid that.

Cheers,
-Quentin

[arsenm]

In D134950#3837377, @qcolombet wrote:

Hi,

Given this is completely target specific, have you considered implementing that in your target specific MachineFunctionInfo pass?
Also, could your use case be supported by the existing target flags that can be attached to MachineOperands?

MachineFunctionInfo isn't really a pass, it's a dumping ground for side data. The primary use of this is when new spill instructions are introduced, so there's no existing MachineOperand to attach this to. Plus a pass to re-add these implies the MIR was in an incorrect state and needed a fixup, which rarely works out well.

If we were to go down adding target flags for virtual registers, there are a few things that we would need to consider like what do we do with the flags when we merge live-ranges (e.g., in the coalescer), could a flag become invalid (e.g., when we do some peephole), etc.

Really what we want is for optimizations to not touch these at all. Coalescing shouldn't merge ranges with different flags.

I feel this patch is creating the llvm equivalent of metadata with all the problems we have with them and I would rather we avoid that.

The alternative is to make everything more aware of the cross lane nature of VGPRs, and I don't really know what that would look like

[qcolombet]

Really what we want is for optimizations to not touch these at all. Coalescing shouldn't merge ranges with different flags.

The problem is it is true for your use case, but it may not be true if other targets start using this.

Let me digest the information and come back to you.

[cdevadas]

The problem is it is true for your use case, but it may not be true if other targets start using this.

Let me digest the information and come back to you.

@qcolombet, much appreciate it if you could give feedback on the vreg flags.

[qcolombet]

much appreciate it if you could give feedback on the vreg flags.

I think that could be fine as long as we specifically say that they can be dropped inadvertently, i.e., you can't rely on them.

That said, I believe that if we put such language, you wouldn't be able to do what you want since in theory the information could be inaccurate and at this point, having this information or no information is pretty much the same. (I would actually argue that no information is probably better!)

[arsenm]

In D134950#3854402, @qcolombet wrote:

much appreciate it if you could give feedback on the vreg flags.

I think that could be fine as long as we specifically say that they can be dropped inadvertently, i.e., you can't rely on them.

That said, I believe that if we put such language, you wouldn't be able to do what you want since in theory the information could be inaccurate and at this point, having this information or no information is pretty much the same. (I would actually argue that no information is probably better!)

They really can't be dropped. Missing information will be just as wrong. What we really want is to establish a dynamic value class. It's like a register class, except it can't map one to one with a physical register since in other contexts the same physical register could be used differently.

[cdevadas]

@qcolombet, you requested a new revision earlier. Let me know if there is any change still required before you can proceed with the review.

[qcolombet]

Hi all,

I've thought a little bit more about this.
Could we handle this with the existing delegate mechanism? (Look at how MRI_NoteNewVirtualRegister works.)

That way, you could do whatever target specific thing you want and it won't spill in the generic code.
To do that we may have to lift the limitation that we have only one delegate at a time (and we may need to add a few more hooks to notify the delegate on the events you are interested in.)

How does that sounds?

Cheers,
-Quentin

[cdevadas]

In D134950#3869851, @qcolombet wrote:

Hi all,

I've thought a little bit more about this.
Could we handle this with the existing delegate mechanism? (Look at how MRI_NoteNewVirtualRegister works.)

That way, you could do whatever target specific thing you want and it won't spill in the generic code.
To do that we may have to lift the limitation that we have only one delegate at a time (and we may need to add a few more hooks to notify the delegate on the events you are interested in.)

How does that sounds?

Cheers,
-Quentin

The delegate is not used in any target-specific passes so far. The specific cases they are currently used for, are looking reasonable. But I'm not really sure the delegate would be a good idea to handle the flags here. It seems like a more complicated way to propagate the flags.

[qcolombet]

Hi,

It seems like a more complicated way to propagate the flags.

It is definitely more complicated, but it doesn't bake any target specific knowledge in the generic code. In particular, in your case, you don't need any merging of the flags during coalescing and whatnot, but that may not be true for other targets if they want to adopt it and I'm not willing to open that door at the moment. Hence the suggested trade-off.

What I'm suggesting is:

• Create your own pass that would attach your delegate to MRI before the point you need it.
• Do whatever you want with your delegate.
• Create another pass to unregister your delegate (you may even leave it there if it doesn't hurt anything.)

Cheers,
-Quentin

[cdevadas]

Used the delegate method to notify targets when a new virtual register is created as well as cloned.

[qcolombet]

I think we're going in the right direction, thanks for your patience with all this.

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

63

I would go with something simpler: SmallPtrSet<Delegate *>.

158

I don't think we need that method and as a result we can ditch the DelegateReceiver altogether.

Essentially when an event occurs, we need to notify all the delegates.

If you still need a getter for your target specific use case, you can keep a hold on your custom delegate in your custom MachineFunctionInfo class.

Alternatively, I guess we could introduce the notion of ID (what you did with DelegateReceiver) but that must not be an enum. For instance, we could use something similar to Pass::ID for that, but for now if the MachineFunctionInfo solution works for you, let's not spend too much time on that.

For the record, the rationale for being against the enum is if people use more than one custom delegate, then the developer's experience is not great because we would have only one ID (RESERVED) to "register" it or they would need to modify the DelegateReceiver to add enum values.

165

After you ditch the DR, we would just need to check that the given delegate is already in the set.

169

Rename this in addDelegate and ditch the DelegateReceiver argument.

llvm/lib/CodeGen/MachineRegisterInfo.cpp

152

By moving this call here, the delegate may miss some important information (like the register class).
I wonder if it would make sense to send two "notifications": one when creating the incomplete register, like here, and one when the register is fully done (like what you did with the clone method).

Otherwise, we'll probably need to stick to having to send the notification just once where it originally was (i.e., within MachineRegisterInfo::createVirtualRegister, MachineRegisterInfo::createGenericVirtualRegister, and MachineRegisterInfo::cloneVirtualRegister).

What do you think?

179

We'll need to notify all delegates here, not just the RESERVED one.

[cdevadas]

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

158

Will remove the enum class and can use just SmallPtrSet.

169

Sure

llvm/lib/CodeGen/MachineRegisterInfo.cpp

152

It makes sense to notify just like the earlier code. I will move it.

[cdevadas]

Addressed review comments.

[qcolombet]

LGTM, minor fixes inlined.

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

60

Have the default implementation call MRI_NoteNewVirtualRegister and ditch the call to MRI_NoteNewVirtualRegister from cloneVirtualRegister.

60

Nit: I found ClonedReg to be ambiguous here: is it the result or the source of the cloning process?

Maybe use NewReg and SrcReg? (Ditto for noteCloneVirtualRegister)

llvm/lib/CodeGen/MachineRegisterInfo.cpp

179

Ditch this one (to make sure we don't double notify delegates that listen to both events.)

[cdevadas]

Thanks. Suggestions incorporated.

[cdevadas]

Used auto * Delegate in the range-based loops.

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

public:
   class Delegate {
     virtual void anchor();
 
   public:
     virtual ~Delegate() = default;
 
     virtual void MRI_NoteNewVirtualRegister(Register Reg) = 0;
   };
 
 private:
   MachineFunction *MF;
   Delegate *TheDelegate = nullptr;
 
   /// True if subregister liveness is tracked.
   const bool TracksSubRegLiveness;
 
   /// VRegInfo - Information we keep for each virtual register.
   ///
   /// Each element in this list contains the register class of the vreg and the
   /// start of the use/def list for the register.
private:
   /// vector is the frozen set of reserved registers when register allocation
   /// started.
   BitVector ReservedRegs;
 
   using VRegToTypeMap = IndexedMap<LLT, VirtReg2IndexFunctor>;
   /// Map generic virtual registers to their low-level type.
   VRegToTypeMap VRegToType;
 
+  /// Target-specific flags associated with a virtual register.
+  IndexedMap<uint8_t, VirtReg2IndexFunctor> VRegFlags;
+
   /// Keep track of the physical registers that are live in to the function.
   /// Live in values are typically arguments in registers.  LiveIn values are
   /// allowed to have virtual registers associated with them, stored in the
   /// second element.
   std::vector<std::pair<MCRegister, Register>> LiveIns;
 
 public:
   explicit MachineRegisterInfo(MachineFunction *MF);
   MachineRegisterInfo(const MachineRegisterInfo &) = delete;
   MachineRegisterInfo &operator=(const MachineRegisterInfo &) = delete;
 
   const TargetRegisterInfo *getTargetRegisterInfo() const {
     return MF->getSubtarget().getRegisterInfo();
   }
 
   void resetDelegate(Delegate *delegate) {
     // Ensure another delegate does not take over unless the current
     // delegate first unattaches itself. If we ever need to multicast
     // notifications, we will need to change to using a list.
     assert(TheDelegate == delegate &&
            "Only the current delegate can perform reset!");
     TheDelegate = nullptr;
   }
 
   void setDelegate(Delegate *delegate) {
     assert(delegate && !TheDelegate &&
            "Attempted to set delegate to null, or to change it without "
            "first resetting it!");
 
     TheDelegate = delegate;
   }
 
   //===--------------------------------------------------------------------===//
   // Function State
   //===--------------------------------------------------------------------===//
public:
 
   /// Set the low-level type of \p VReg to \p Ty.
   void setType(Register VReg, LLT Ty);
 
   /// Create and return a new generic virtual register with low-level
   /// type \p Ty.
   Register createGenericVirtualRegister(LLT Ty, StringRef Name = "");
 
+  /// Set \p Flag to \p Reg
+  void setFlag(Register Reg, uint8_t Flag);
+
+  /// Returns true if \p Flag is set for \p Reg
+  bool checkFlag(Register Reg, uint8_t Flag) const;
+
   /// Remove all types associated to virtual registers (after instruction
   /// selection and constraining of all generic virtual registers).
   void clearVirtRegTypes();
 
   /// Creates a new virtual register that has no register class, register bank
   /// or size assigned yet. This is only allowed to be used
   /// temporarily while constructing machine instructions. Most operations are
   /// undefined on an incomplete register until one of setRegClass(),

llvm/lib/CodeGen/MachineRegisterInfo.cpp

 void MachineRegisterInfo::Delegate::anchor() {}
 
 MachineRegisterInfo::MachineRegisterInfo(MachineFunction *MF)
     : MF(MF), TracksSubRegLiveness(MF->getSubtarget().enableSubRegLiveness() &&
                                    EnableSubRegLiveness) {
   unsigned NumRegs = getTargetRegisterInfo()->getNumRegs();
   VRegInfo.reserve(256);
   RegAllocHints.reserve(256);
+  VRegFlags.reserve(256);
   UsedPhysRegMask.resize(NumRegs);
   PhysRegUseDefLists.reset(new MachineOperand*[NumRegs]());
 }
 
 /// setRegClass - Set the register class of the specified virtual register.
 ///
 void
 MachineRegisterInfo::setRegClass(Register Reg, const TargetRegisterClass *RC) {
MachineRegisterInfo::recomputeRegClass(Register Reg) {
   setRegClass(Reg, NewRC);
   return true;
 }
 
 Register MachineRegisterInfo::createIncompleteVirtualRegister(StringRef Name) {
   Register Reg = Register::index2VirtReg(getNumVirtRegs());
   VRegInfo.grow(Reg);
   RegAllocHints.grow(Reg);
+  VRegFlags.grow(Reg);
   insertVRegByName(Name, Reg);
   return Reg;
 }
 
 /// createVirtualRegister - Create and return a new virtual register in the
 /// function with the specified register class.
 ///
 Register
 MachineRegisterInfo::createVirtualRegister(const TargetRegisterClass *RegClass,
                                            StringRef Name) {
   assert(RegClass && "Cannot create register without RegClass!");
   assert(RegClass->isAllocatable() &&
          "Virtual register RegClass must be allocatable.");
 
   // New virtual register number.
   Register Reg = createIncompleteVirtualRegister(Name);
   VRegInfo[Reg].first = RegClass;
   if (TheDelegate)
     TheDelegate->MRI_NoteNewVirtualRegister(Reg);
   return Reg;
 }
 
 Register MachineRegisterInfo::cloneVirtualRegister(Register VReg,
                                                    StringRef Name) {
   Register Reg = createIncompleteVirtualRegister(Name);
   VRegInfo[Reg].first = VRegInfo[VReg].first;
+  VRegFlags[Reg] = VRegFlags[VReg];
   setType(Reg, getType(VReg));
   if (TheDelegate)
     TheDelegate->MRI_NoteNewVirtualRegister(Reg);
   return Reg;
 }
 
 void MachineRegisterInfo::setType(Register VReg, LLT Ty) {
   VRegToType.grow(VReg);
   VRegToType[VReg] = Ty;
 }
 
+void MachineRegisterInfo::setFlag(Register Reg, uint8_t Flag) {
+  assert(Register::isVirtualRegister(Reg));
+  if (VRegFlags.inBounds(Reg))
+    VRegFlags[Reg] |= (uint8_t)1 << Flag;
+}
+
+bool MachineRegisterInfo::checkFlag(Register Reg, uint8_t Flag) const {
+  if (Register::isVirtualRegister(Reg))
+    return VRegFlags.inBounds(Reg) && VRegFlags[Reg] & ((uint8_t)1 << Flag);
+
+  return false;
+}
+
 Register
 MachineRegisterInfo::createGenericVirtualRegister(LLT Ty, StringRef Name) {
   // New virtual register number.
   Register Reg = createIncompleteVirtualRegister(Name);
   // FIXME: Should we use a dummy register class?
   VRegInfo[Reg].first = static_cast<RegisterBank *>(nullptr);
   setType(Reg, Ty);
   if (TheDelegate)