From c93b41da95ca728c28963a1ca05325ed1b8676f0 Mon Sep 17 00:00:00 2001
From: Schuyler Eldridge <schuyler.eldridge@sifive.com>
Date: Fri, 10 Jan 2025 18:13:07 -0500
Subject: [PATCH] [FIRRTL] Reuse XMRRefOps in LowerXMR

Change the `LowerXMR` pass to reuse previously created `XMRRefOp`s within
a module.  This is done for later, easier processing of these operations
in `LowerToHW`.  While these operations could be made to be CSE'd (by
giving them the `Pure` trait), it is advantageous to instead not create
duplicate operations as this avoids needing to run schedule a mostly
useless CSE.

This is primarily done to allow for Verilator-compatible lowering of
`force_initial` and `release_initial` in `LowerToHW` where it would be
useful to know when a force and release are referring to the same
reference.  (Yes, full alias analysis is necessary to do this correctly
and relying on the CSE of `LowerXMR` is going to be buggy.)

Signed-off-by: Schuyler Eldridge <schuyler.eldridge@sifive.com>
---
 lib/Dialect/FIRRTL/Transforms/LowerXMR.cpp | 55 +++++++++++++++++++++-
 test/Dialect/FIRRTL/lowerXMR.mlir          | 13 ++---
 2 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/lib/Dialect/FIRRTL/Transforms/LowerXMR.cpp b/lib/Dialect/FIRRTL/Transforms/LowerXMR.cpp
index af671caed165..ca89f47bce49 100644
--- a/lib/Dialect/FIRRTL/Transforms/LowerXMR.cpp
+++ b/lib/Dialect/FIRRTL/Transforms/LowerXMR.cpp
@@ -76,6 +76,50 @@ struct XMRNode {
   os << ", next=" << node.next << ")";
   return os;
 }
+
+/// Track information about operations being created in a module.  This is used
+/// to generate more compact code and reuse operations where possible.
+class ModuleState {
+
+public:
+  ModuleState(FModuleOp &moduleOp) : body(moduleOp.getBodyBlock()) {}
+
+  /// Return the existing XMRRefOp for this type, symbol, and suffix for this
+  /// module.  Otherwise, create a new one.  The first XMRRefOp will be created
+  /// at the beginning of the module.  Subsequent XMRRefOps will be created
+  /// immediately following the first one.
+  Value getOrCreateXMRRefOp(Type type, FlatSymbolRefAttr symbol,
+                            StringAttr suffix, ImplicitLocOpBuilder &builder) {
+    // Return the saved XMRRefOp.
+    auto it = xmrRefCache.find({type, symbol, suffix});
+    if (it != xmrRefCache.end())
+      return it->getSecond();
+
+    // Create a new XMRRefOp.
+    OpBuilder::InsertionGuard guard(builder);
+    if (xmrRefPoint.isSet())
+      builder.restoreInsertionPoint(xmrRefPoint);
+    else
+      builder.setInsertionPointToStart(body);
+
+    Value xmr = builder.create<XMRRefOp>(type, symbol, suffix);
+    xmrRefCache.insert({{type, symbol, suffix}, xmr});
+
+    xmrRefPoint = builder.saveInsertionPoint();
+    return xmr;
+  };
+
+private:
+  /// The module's body.  This is used to set the insertion point for the first
+  /// created operation.
+  Block *body;
+
+  /// Map used to know if we created this XMRRefOp before.
+  DenseMap<std::tuple<Type, SymbolRefAttr, StringAttr>, Value> xmrRefCache;
+
+  /// The saved insertion point for XMRRefOps.
+  OpBuilder::InsertPoint xmrRefPoint;
+};
 } // end anonymous namespace
 
 class LowerXMRPass : public circt::firrtl::impl::LowerXMRBase<LowerXMRPass> {
@@ -263,6 +307,8 @@ class LowerXMRPass : public circt::firrtl::impl::LowerXMRBase<LowerXMRPass> {
       LLVM_DEBUG(llvm::dbgs()
                  << "Traversing module:" << module.getModuleNameAttr() << "\n");
 
+      moduleStates.insert({module, ModuleState(module)});
+
       if (module.isPublic())
         publicModules.push_back(module);
 
@@ -488,7 +534,10 @@ class LowerXMRPass : public circt::firrtl::impl::LowerXMRBase<LowerXMRPass> {
               if (failed(resolveReference(op.getDest(), builder, ref, str)))
                 return failure();
 
-              Value xmr = builder.create<XMRRefOp>(destType, ref, str);
+              Value xmr =
+                  moduleStates.find(op->template getParentOfType<FModuleOp>())
+                      ->getSecond()
+                      .getOrCreateXMRRefOp(destType, ref, str, builder);
               op.getDestMutable().assign(xmr);
               return success();
             })
@@ -772,6 +821,7 @@ class LowerXMRPass : public circt::firrtl::impl::LowerXMRBase<LowerXMRPass> {
     refPortsToRemoveMap.clear();
     dataflowAt.clear();
     refSendPathList.clear();
+    moduleStates.clear();
   }
 
   bool isZeroWidth(FIRRTLBaseType t) { return t.getBitWidthOrSentinel() == 0; }
@@ -859,6 +909,9 @@ class LowerXMRPass : public circt::firrtl::impl::LowerXMRBase<LowerXMRPass> {
 
   /// The insertion point where the pass inserts HierPathOps.
   OpBuilder::InsertPoint pathInsertPoint = {};
+
+  /// Per-module helpers for creating operations within modules.
+  DenseMap<FModuleOp, ModuleState> moduleStates;
 };
 
 std::unique_ptr<mlir::Pass> circt::firrtl::createLowerXMRPass() {
diff --git a/test/Dialect/FIRRTL/lowerXMR.mlir b/test/Dialect/FIRRTL/lowerXMR.mlir
index c8969b4764a8..1e55aeb1c69e 100644
--- a/test/Dialect/FIRRTL/lowerXMR.mlir
+++ b/test/Dialect/FIRRTL/lowerXMR.mlir
@@ -437,19 +437,16 @@ firrtl.circuit "ForceRelease" {
   }
   // CHECK-LABEL: firrtl.module @ForceRelease
   firrtl.module @ForceRelease(in %c: !firrtl.uint<1>, in %clock: !firrtl.clock, in %x: !firrtl.uint<4>) {
+      // CHECK-NEXT: %[[REF:.+]] = firrtl.xmr.ref @[[XMRPATH]] : !firrtl.rwprobe<uint<4>>
       // CHECK-NEXT: firrtl.instance r sym @[[INST_SYM]] @RefMe()
       %r_p = firrtl.instance r @RefMe(out p: !firrtl.rwprobe<uint<4>>)
-      // CHECK-NEXT: %[[REF1:.+]] = firrtl.xmr.ref @[[XMRPATH]] : !firrtl.rwprobe<uint<4>>
-      // CHECK-NEXT: firrtl.ref.force %clock, %c, %[[REF1]], %x : !firrtl.clock, !firrtl.uint<1>, !firrtl.rwprobe<uint<4>>, !firrtl.uint<4>
+      // CHECK-NEXT: firrtl.ref.force %clock, %c, %[[REF]], %x : !firrtl.clock, !firrtl.uint<1>, !firrtl.rwprobe<uint<4>>, !firrtl.uint<4>
       firrtl.ref.force %clock, %c, %r_p, %x : !firrtl.clock, !firrtl.uint<1>, !firrtl.rwprobe<uint<4>>, !firrtl.uint<4>
-      // CHECK-NEXT: %[[REF2:.+]] = firrtl.xmr.ref @[[XMRPATH]] : !firrtl.rwprobe<uint<4>>
-      // CHECK-NEXT: firrtl.ref.force_initial %c, %[[REF2]], %x : !firrtl.uint<1>, !firrtl.rwprobe<uint<4>>, !firrtl.uint<4>
+      // CHECK-NEXT: firrtl.ref.force_initial %c, %[[REF]], %x : !firrtl.uint<1>, !firrtl.rwprobe<uint<4>>, !firrtl.uint<4>
       firrtl.ref.force_initial %c, %r_p, %x : !firrtl.uint<1>, !firrtl.rwprobe<uint<4>>, !firrtl.uint<4>
-      // CHECK-NEXT: %[[REF3:.+]] = firrtl.xmr.ref @[[XMRPATH]] : !firrtl.rwprobe<uint<4>>
-      // CHECK-NEXT: firrtl.ref.release %clock, %c, %[[REF3]] : !firrtl.clock, !firrtl.uint<1>, !firrtl.rwprobe<uint<4>>
+      // CHECK-NEXT: firrtl.ref.release %clock, %c, %[[REF]] : !firrtl.clock, !firrtl.uint<1>, !firrtl.rwprobe<uint<4>>
       firrtl.ref.release %clock, %c, %r_p : !firrtl.clock, !firrtl.uint<1>, !firrtl.rwprobe<uint<4>>
-      // CHECK-NEXT: %[[REF4:.+]] = firrtl.xmr.ref @[[XMRPATH]] : !firrtl.rwprobe<uint<4>>
-      // CHECK-NEXT: firrtl.ref.release_initial %c, %[[REF4]] : !firrtl.uint<1>, !firrtl.rwprobe<uint<4>>
+      // CHECK-NEXT: firrtl.ref.release_initial %c, %[[REF]] : !firrtl.uint<1>, !firrtl.rwprobe<uint<4>>
       firrtl.ref.release_initial %c, %r_p : !firrtl.uint<1>, !firrtl.rwprobe<uint<4>>
     }
   }