[llvm] [CodeGen] Add listener support to the rematerializer (NFC) (PR #184338)

[Lucas Ramirez via llvm-commits]

https://github.com/lucas-rami updated https://github.com/llvm/llvm-project/pull/184338

>From 0991e231ce602acedb4927fcd34ddf31e5782069 Mon Sep 17 00:00:00 2001
From: Lucas Ramirez <lucas.rami at proton.me>
Date: Thu, 26 Feb 2026 14:05:54 +0000
Subject: [PATCH 1/2] [CodeGen] Add listener support to the rematerializer
 (NFC)

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

This has no user at the moment, but the plan is to have listeners start
being responsible for secondary/optional functionnalities 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).
---
 llvm/include/llvm/CodeGen/Rematerializer.h | 43 ++++++++++++++++++++++
 llvm/lib/CodeGen/Rematerializer.cpp        |  8 +++-
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/CodeGen/Rematerializer.h b/llvm/include/llvm/CodeGen/Rematerializer.h
index 44b50b1e597c7..a04dedb36a6b7 100644
--- a/llvm/include/llvm/CodeGen/Rematerializer.h
+++ b/llvm/include/llvm/CodeGen/Rematerializer.h
@@ -186,6 +186,37 @@ 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 {
+  private:
+    /// Called before register \p RegIdx is rematerialized to register \p
+    /// RematRegIdx. At this point the rematerialization does not exist.
+    virtual void beforeRegRematerialized(const Rematerializer &Remater,
+                                         RegisterIdx RegIdx,
+                                         RegisterIdx RematRegIdx) {}
+
+    /// Called just after register \p NewRegIdx is created (following a
+    /// rematerializations). At this point the rematerialization exists but does
+    /// not yet have any user.
+    virtual void newRegCreated(const Rematerializer &Remater,
+                               RegisterIdx NewRegIdx) {}
+
+    /// Called juste before register \p RegIdx is deleted from the MIR. At this
+    /// point the register still exists and is guaranteed to have 0 users.
+    virtual void beforeRegDeleted(const Rematerializer &Remater,
+                                  RegisterIdx RegIdx) {}
+
+    friend Rematerializer;
+
+  public:
+    using RegisterIdx = Rematerializer::RegisterIdx;
+
+    virtual ~Listener() = default;
+  };
+
   /// Error value for register indices.
   static constexpr unsigned NoReg = ~0;
 
@@ -207,6 +238,11 @@ class Rematerializer {
   /// rematerializable register in regions.
   bool analyze(bool SupportRollback);
 
+  /// Adds a listener to the rematerializer.
+  void addListener(Listener *Listen) { Listeners.push_back(Listen); }
+  /// Removes all listeners from the rematerializers.
+  void clearListeners(Listener *Listen) { Listeners.clear(); }
+
   inline const Reg &getReg(RegisterIdx RegIdx) const {
     assert(RegIdx < Regs.size() && "out of bounds");
     return Regs[RegIdx];
@@ -386,6 +422,13 @@ class Rematerializer {
   LiveIntervals &LIS;
   const TargetInstrInfo &TII;
   const TargetRegisterInfo &TRI;
+  SmallVector<Listener *, 2> Listeners;
+
+  template <typename Callback, typename... ArgsTy>
+  void notifyListeners(Callback Cb, ArgsTy &&...Args) const {
+    for (Listener *Listen : Listeners)
+      (Listen->*Cb)(*this, std::forward<ArgsTy>(Args)...);
+  }
 
   /// 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..80633d35c22c5 100644
--- a/llvm/lib/CodeGen/Rematerializer.cpp
+++ b/llvm/lib/CodeGen/Rematerializer.cpp
@@ -401,8 +401,11 @@ void Rematerializer::deleteRegIfUnused(RegisterIdx RootIdx) {
 }
 
 void Rematerializer::deleteReg(RegisterIdx RegIdx) {
+  notifyListeners(&Listener::beforeRegDeleted, 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
   // boundary since we don't ever consider them rematerializable.
   MachineBasicBlock::iterator &RegionBegin = Regions[DeleteReg.DefRegion].first;
@@ -574,9 +577,10 @@ RegisterIdx Rematerializer::getDefRegIdx(const MachineInstr &MI) const {
 RegisterIdx Rematerializer::rematerializeReg(
     RegisterIdx RegIdx, MachineBasicBlock::iterator InsertPos,
     SmallVectorImpl<Reg::Dependency> &&Dependencies) {
-  unsigned UseRegion = MIRegion.at(&*InsertPos);
   RegisterIdx NewRegIdx = Regs.size();
+  notifyListeners(&Listener::beforeRegRematerialized, RegIdx, NewRegIdx);
 
+  unsigned UseRegion = MIRegion.at(&*InsertPos);
   Reg &NewReg = Regs.emplace_back();
   Reg &FromReg = Regs[RegIdx];
   NewReg.Mask = FromReg.Mask;
@@ -627,6 +631,8 @@ RegisterIdx Rematerializer::rematerializeReg(
     LISUpdates.insert(NewDep.RegIdx);
   }
 
+  notifyListeners(&Listener::newRegCreated, NewRegIdx);
+
   LLVM_DEBUG({
     dbgs() << "** Rematerialized " << printID(RegIdx) << " as "
            << printRematReg(NewRegIdx) << '\n';

>From ef56bdef624010969127d0e76ea61c86f0685589 Mon Sep 17 00:00:00 2001
From: Lucas Ramirez <lucas.rami at proton.me>
Date: Tue, 3 Mar 2026 13:35:21 +0000
Subject: [PATCH 2/2] Format

---
 llvm/lib/CodeGen/Rematerializer.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/Rematerializer.cpp b/llvm/lib/CodeGen/Rematerializer.cpp
index 80633d35c22c5..739c59d53c189 100644
--- a/llvm/lib/CodeGen/Rematerializer.cpp
+++ b/llvm/lib/CodeGen/Rematerializer.cpp
@@ -405,7 +405,7 @@ void Rematerializer::deleteReg(RegisterIdx 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
   // boundary since we don't ever consider them rematerializable.
   MachineBasicBlock::iterator &RegionBegin = Regions[DeleteReg.DefRegion].first;