From 2f23f5c0d53c619101541a348cd817174ac27783 Mon Sep 17 00:00:00 2001
From: Christudasan Devadasan <Christudasan.Devadasan@amd.com>
Date: Fri, 30 Sep 2022 13:20:58 +0530
Subject: [PATCH] [CodeGen] Use delegate to notify targets when virtual
 registers are created

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 support multicast.

Reviewed By: qcolombet

Differential Revision: https://reviews.llvm.org/D134950
---
 llvm/include/llvm/CodeGen/LiveRangeEdit.h       |  2 +-
 llvm/include/llvm/CodeGen/MachineRegisterInfo.h | 34 +++++++++++++++++--------
 llvm/lib/CodeGen/MachineRegisterInfo.cpp        | 10 +++-----
 3 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/LiveRangeEdit.h b/llvm/include/llvm/CodeGen/LiveRangeEdit.h
index aaaf8ea..fedcff0 100644
--- a/llvm/include/llvm/CodeGen/LiveRangeEdit.h
+++ b/llvm/include/llvm/CodeGen/LiveRangeEdit.h
@@ -133,7 +133,7 @@ public:
       : Parent(parent), NewRegs(newRegs), MRI(MF.getRegInfo()), LIS(lis),
         VRM(vrm), TII(*MF.getSubtarget().getInstrInfo()), TheDelegate(delegate),
         FirstNew(newRegs.size()), DeadRemats(deadRemats) {
-    MRI.setDelegate(this);
+    MRI.addDelegate(this);
   }
 
   ~LiveRangeEdit() override { MRI.resetDelegate(this); }
diff --git a/llvm/include/llvm/CodeGen/MachineRegisterInfo.h b/llvm/include/llvm/CodeGen/MachineRegisterInfo.h
index 75a444f..5722172 100644
--- a/llvm/include/llvm/CodeGen/MachineRegisterInfo.h
+++ b/llvm/include/llvm/CodeGen/MachineRegisterInfo.h
@@ -17,6 +17,7 @@
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/IndexedMap.h"
 #include "llvm/ADT/PointerUnion.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/iterator_range.h"
@@ -56,11 +57,15 @@ public:
     virtual ~Delegate() = default;
 
     virtual void MRI_NoteNewVirtualRegister(Register Reg) = 0;
+    virtual void MRI_NotecloneVirtualRegister(Register NewReg,
+                                              Register SrcReg) {
+      MRI_NoteNewVirtualRegister(NewReg);
+    }
   };
 
 private:
   MachineFunction *MF;
-  Delegate *TheDelegate = nullptr;
+  SmallPtrSet<Delegate *, 1> TheDelegates;
 
   /// True if subregister liveness is tracked.
   const bool TracksSubRegLiveness;
@@ -154,19 +159,28 @@ public:
 
   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;
+    // delegate first unattaches itself.
+    assert(TheDelegates.count(delegate) &&
+           "Only an existing delegate can perform reset!");
+    TheDelegates.erase(delegate);
   }
 
-  void setDelegate(Delegate *delegate) {
-    assert(delegate && !TheDelegate &&
-           "Attempted to set delegate to null, or to change it without "
+  void addDelegate(Delegate *delegate) {
+    assert(delegate && !TheDelegates.count(delegate) &&
+           "Attempted to add null delegate, or to change it without "
            "first resetting it!");
 
-    TheDelegate = delegate;
+    TheDelegates.insert(delegate);
+  }
+
+  void noteNewVirtualRegister(Register Reg) {
+    for (auto *TheDelegate : TheDelegates)
+      TheDelegate->MRI_NoteNewVirtualRegister(Reg);
+  }
+
+  void noteCloneVirtualRegister(Register NewReg, Register SrcReg) {
+    for (auto *TheDelegate : TheDelegates)
+      TheDelegate->MRI_NotecloneVirtualRegister(NewReg, SrcReg);
   }
 
   //===--------------------------------------------------------------------===//
diff --git a/llvm/lib/CodeGen/MachineRegisterInfo.cpp b/llvm/lib/CodeGen/MachineRegisterInfo.cpp
index e322250..5a7d432 100644
--- a/llvm/lib/CodeGen/MachineRegisterInfo.cpp
+++ b/llvm/lib/CodeGen/MachineRegisterInfo.cpp
@@ -48,6 +48,7 @@ MachineRegisterInfo::MachineRegisterInfo(MachineFunction *MF)
   RegAllocHints.reserve(256);
   UsedPhysRegMask.resize(NumRegs);
   PhysRegUseDefLists.reset(new MachineOperand*[NumRegs]());
+  TheDelegates.clear();
 }
 
 /// setRegClass - Set the register class of the specified virtual register.
@@ -162,8 +163,7 @@ MachineRegisterInfo::createVirtualRegister(const TargetRegisterClass *RegClass,
   // New virtual register number.
   Register Reg = createIncompleteVirtualRegister(Name);
   VRegInfo[Reg].first = RegClass;
-  if (TheDelegate)
-    TheDelegate->MRI_NoteNewVirtualRegister(Reg);
+  noteNewVirtualRegister(Reg);
   return Reg;
 }
 
@@ -172,8 +172,7 @@ Register MachineRegisterInfo::cloneVirtualRegister(Register VReg,
   Register Reg = createIncompleteVirtualRegister(Name);
   VRegInfo[Reg].first = VRegInfo[VReg].first;
   setType(Reg, getType(VReg));
-  if (TheDelegate)
-    TheDelegate->MRI_NoteNewVirtualRegister(Reg);
+  noteCloneVirtualRegister(Reg, VReg);
   return Reg;
 }
 
@@ -189,8 +188,7 @@ MachineRegisterInfo::createGenericVirtualRegister(LLT Ty, StringRef Name) {
   // FIXME: Should we use a dummy register class?
   VRegInfo[Reg].first = static_cast<RegisterBank *>(nullptr);
   setType(Reg, Ty);
-  if (TheDelegate)
-    TheDelegate->MRI_NoteNewVirtualRegister(Reg);
+  noteNewVirtualRegister(Reg);
   return Reg;
 }
 
-- 
2.7.4