llvm · seldridge · Jan 13, 2025 · Jan 10, 2025
@@ -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() {

@@ -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>>
     }
   }