[llvm] 8a06085 - [CodeGen] Add listener support to the rematerializer (NFC) (#184338)

[via llvm-commits]

Author: Lucas Ramirez
Date: 2026-04-01T13:35:37+02:00
New Revision: 8a06085c610a49739beb0af716a30714f073540c

URL: https://github.com/llvm/llvm-project/commit/8a06085c610a49739beb0af716a30714f073540c
DIFF: https://github.com/llvm/llvm-project/commit/8a06085c610a49739beb0af716a30714f073540c.diff

LOG: [CodeGen] Add listener support to the rematerializer (NFC) (#184338)

This change adds support for adding listeners to the target-independent
rematerializer; listeners can catch certain rematerialization-related
events to implement some additional functionality on top of what the
rematerializer already performs.

This is NFC and has no user at the moment, but the plan is to have
listeners start being responsible for secondary/optional functionalities
that are at the moment integrated with the rematerializer itself. Two
examples of that are:

1. rollback support (currently optional), and
2. region tracking (currently mandatory, but not fundamentally necessary
to the rematerializer).

Added: 
    

Modified: 
    llvm/include/llvm/CodeGen/Rematerializer.h
    llvm/lib/CodeGen/Rematerializer.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/CodeGen/Rematerializer.h b/llvm/include/llvm/CodeGen/Rematerializer.h
index 44b50b1e597c7..4619223834c36 100644
--- a/llvm/include/llvm/CodeGen/Rematerializer.h
+++ b/llvm/include/llvm/CodeGen/Rematerializer.h
@@ -186,6 +186,31 @@ class Rematerializer {
     friend Rematerializer;
   };
 
+  /// Rematerializer listener. Defines overridable hooks that allow to catch
+  /// specific events inside the rematerializer. All hooks do nothing by
+  /// default. Listeners can be added or removed at any time during the
+  /// rematerializer's lifetime.
+  class Listener {
+  public:
+    using RegisterIdx = Rematerializer::RegisterIdx;
+
+    /// Called just after register \p NewRegIdx is created (following a
+    /// rematerialization). At this point the rematerialization exists in the \p
+    /// Remater state and the MIR but does not yet have any user.
+    virtual void rematerializerNoteRegCreated(const Rematerializer &Remater,
+                                              RegisterIdx NewRegIdx) {}
+
+    /// Called juste before register \p RegIdx is deleted from the MIR. At this
+    /// point the register still exists in the MIR but no longer has any user.
+    virtual void rematerializerNoteRegDeleted(const Rematerializer &Remater,
+                                              RegisterIdx RegIdx) {}
+
+    virtual ~Listener() = default;
+
+  private:
+    virtual void anchor();
+  };
+
   /// Error value for register indices.
   static constexpr unsigned NoReg = ~0;
 
@@ -207,6 +232,22 @@ class Rematerializer {
   /// rematerializable register in regions.
   bool analyze(bool SupportRollback);
 
+  /// Adds a new listener to the rematerializer.
+  void addListener(Listener *Listen) {
+    assert(Listen && "null listener");
+    assert(!Listeners.contains(Listen) && "duplicate listener");
+    Listeners.insert(Listen);
+  }
+
+  /// Removes a listener from the rematerializer.
+  void removeListener(Listener *Listen) {
+    assert(Listeners.contains(Listen) && "unknown listener");
+    Listeners.erase(Listen);
+  }
+
+  /// Removes all listeners from the rematerializer.
+  void clearListeners() { Listeners.clear(); }
+
   inline const Reg &getReg(RegisterIdx RegIdx) const {
     assert(RegIdx < Regs.size() && "out of bounds");
     return Regs[RegIdx];
@@ -386,6 +427,17 @@ class Rematerializer {
   LiveIntervals &LIS;
   const TargetInstrInfo &TII;
   const TargetRegisterInfo &TRI;
+  SmallPtrSet<Listener *, 1> Listeners;
+
+  void noteRegCreated(RegisterIdx RegIdx) const {
+    for (Listener *Listen : Listeners)
+      Listen->rematerializerNoteRegCreated(*this, RegIdx);
+  }
+
+  void noteRegDeleted(RegisterIdx RegIdx) const {
+    for (Listener *Listen : Listeners)
+      Listen->rematerializerNoteRegDeleted(*this, RegIdx);
+  }
 
   /// Rematerializable registers identified since the rematerializer's creation,
   /// both dead and alive, originals and rematerializations. No register is ever

diff  --git a/llvm/lib/CodeGen/Rematerializer.cpp b/llvm/lib/CodeGen/Rematerializer.cpp
index 1c28c8afc52eb..5e248c552f3a7 100644
--- a/llvm/lib/CodeGen/Rematerializer.cpp
+++ b/llvm/lib/CodeGen/Rematerializer.cpp
@@ -30,6 +30,9 @@
 using namespace llvm;
 using RegisterIdx = Rematerializer::RegisterIdx;
 
+// Pin the vtable to this file.
+void Rematerializer::Listener::anchor() {}
+
 /// Checks whether the value in \p LI at \p UseIdx is identical to \p OVNI (this
 /// implies it is also live there). When \p LI has sub-ranges, checks that
 /// all sub-ranges intersecting with \p Mask are also live at \p UseIdx.
@@ -401,6 +404,8 @@ void Rematerializer::deleteRegIfUnused(RegisterIdx RootIdx) {
 }
 
 void Rematerializer::deleteReg(RegisterIdx RegIdx) {
+  noteRegDeleted(RegIdx);
+
   Reg &DeleteReg = Regs[RegIdx];
   assert(DeleteReg.DefMI && "register was already deleted");
   // It is not possible for the deleted instruction to be the upper region
@@ -627,6 +632,8 @@ RegisterIdx Rematerializer::rematerializeReg(
     LISUpdates.insert(NewDep.RegIdx);
   }
 
+  noteRegCreated(NewRegIdx);
+
   LLVM_DEBUG({
     dbgs() << "** Rematerialized " << printID(RegIdx) << " as "
            << printRematReg(NewRegIdx) << '\n';