[FIRRTL] Add force dedup reduction

Add a reduction to the FIRRTL dialect that takes MustDedup annotations
and forcefully replaces all modules in the group with the first module.
This is brittle, but works fairly well in practice.

Author: fabianschuiki

include/circt/Reduce/Reduction.h

@@ -104,19 +104,38 @@ struct Reduction {
   }
 };
 
+/// A reduction pattern for a specific operation.
+///
+/// Only matches on operations of type `OpTy`, and calls corresponding match and
+/// rewrite functions with the operation cast to this type, for convenience.
 template <typename OpTy>
 struct OpReduction : public Reduction {
-  uint64_t match(Operation *op) override {
+  void matches(Operation *op,
+               llvm::function_ref<void(uint64_t, uint64_t)> addMatch) override {
     if (auto concreteOp = dyn_cast<OpTy>(op))
-      return match(concreteOp);
-    return 0;
+      matches(concreteOp, addMatch);
   }
-  LogicalResult rewrite(Operation *op) override {
-    return rewrite(cast<OpTy>(op));
+  LogicalResult rewriteMatches(Operation *op,
+                               ArrayRef<uint64_t> matches) override {
+    return rewriteMatches(cast<OpTy>(op), matches);
   }
 
   virtual uint64_t match(OpTy op) { return 1; }
-  virtual LogicalResult rewrite(OpTy op) = 0;
+  virtual void matches(OpTy op,
+                       llvm::function_ref<void(uint64_t, uint64_t)> addMatch) {
+    addMatch(match(op), 0);
+  }
+  virtual LogicalResult rewrite(OpTy op) { return failure(); }
+  virtual LogicalResult rewriteMatches(OpTy op, ArrayRef<uint64_t> matches) {
+    assert(matches.size() == 1 && matches[0] == 0);
+    return rewrite(op);
+  }
+
+private:
+  /// Hide the base class match/rewrite functions to prevent compiler warnings
+  /// about the `OpTy`-specific ones hiding the base class functions.
+  using Reduction::match;
+  using Reduction::rewrite;
 };
 
 /// A reduction pattern that applies an `mlir::Pass`.

lib/Dialect/FIRRTL/FIRRTLReductions.cpp

@@ -7,7 +7,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "circt/Dialect/FIRRTL/FIRRTLReductions.h"
+#include "circt/Dialect/FIRRTL/AnnotationDetails.h"
 #include "circt/Dialect/FIRRTL/CHIRRTLDialect.h"
+#include "circt/Dialect/FIRRTL/FIRRTLAnnotations.h"
 #include "circt/Dialect/FIRRTL/FIRRTLInstanceGraph.h"
 #include "circt/Dialect/FIRRTL/FIRRTLOps.h"
 #include "circt/Dialect/FIRRTL/FIRRTLUtils.h"
@@ -25,6 +27,7 @@
 
 using namespace mlir;
 using namespace circt;
+using namespace firrtl;
 
 //===----------------------------------------------------------------------===//
 // Utilities
@@ -179,6 +182,8 @@ struct NLARemover {
 // Reduction patterns
 //===----------------------------------------------------------------------===//
 
+namespace {
+
 /// A sample reduction pattern that maps `firrtl.module` to `firrtl.extmodule`.
 struct FIRRTLModuleExternalizer : public OpReduction<firrtl::FModuleOp> {
   void beforeReduction(mlir::ModuleOp op) override {
@@ -1118,6 +1123,182 @@ struct ModuleNameSanitizer : OpReduction<firrtl::CircuitOp> {
   bool isOneShot() const override { return true; }
 };
 
+/// A reduction pattern that handles MustDedup annotations by replacing all
+/// module names in a dedup group with a single module name. This helps reduce
+/// the IR by consolidating module references that are required to be identical.
+///
+/// The pattern works by:
+/// 1. Finding all MustDeduplicateAnnotation annotations on the circuit
+/// 2. For each dedup group, using the first module as the canonical name
+/// 3. Replacing all instance references to other modules in the group with
+///    references to the canonical module
+/// 4. Removing the non-canonical modules from the circuit
+/// 5. Removing the processed MustDedup annotation
+///
+/// This reduction is particularly useful for reducing large circuits where
+/// multiple modules are known to be identical but haven't been deduplicated
+/// yet.
+struct ForceDedup : public OpReduction<CircuitOp> {
+  void beforeReduction(mlir::ModuleOp op) override {
+    symbols.clear();
+    nlaRemover.clear();
+  }
+  void afterReduction(mlir::ModuleOp op) override { nlaRemover.remove(op); }
+
+  /// Collect all MustDedup annotations and create matches for each dedup group.
+  void matches(CircuitOp circuitOp,
+               llvm::function_ref<void(uint64_t, uint64_t)> addMatch) override {
+    auto annotations = AnnotationSet(circuitOp);
+    for (auto [annoIdx, anno] : llvm::enumerate(annotations)) {
+      if (!anno.isClass(mustDedupAnnoClass))
+        continue;
+
+      auto modulesAttr = anno.getMember<ArrayAttr>("modules");
+      if (!modulesAttr)
+        continue;
+
+      // Each dedup group gets its own match with benefit proportional to group
+      // size.
+      uint64_t benefit = modulesAttr.size();
+      addMatch(benefit, annoIdx);
+    }
+  }
+
+  LogicalResult rewriteMatches(CircuitOp circuitOp,
+                               ArrayRef<uint64_t> matches) override {
+    auto *context = circuitOp->getContext();
+    NLATable nlaTable(circuitOp);
+    hw::InnerSymbolTableCollection innerSymTables;
+    auto annotations = AnnotationSet(circuitOp);
+    SmallVector<Annotation> newAnnotations;
+
+    for (auto [annoIdx, anno] : llvm::enumerate(annotations)) {
+      // Check if this annotation was selected.
+      if (!llvm::is_contained(matches, annoIdx)) {
+        newAnnotations.push_back(anno);
+        continue;
+      }
+      auto modulesAttr = anno.getMember<ArrayAttr>("modules");
+      assert(anno.isClass(mustDedupAnnoClass) && modulesAttr &&
+             modulesAttr.size() >= 2);
+
+      // Extract module names from the dedup group.
+      SmallVector<StringAttr> moduleNames;
+      for (auto moduleRef : modulesAttr.getAsRange<StringAttr>()) {
+        // Parse "~CircuitName|ModuleName" format.
+        auto refStr = moduleRef.getValue();
+        auto pipePos = refStr.find('|');
+        if (pipePos != StringRef::npos && pipePos + 1 < refStr.size()) {
+          auto moduleName = refStr.substr(pipePos + 1);
+          moduleNames.push_back(StringAttr::get(context, moduleName));
+        }
+      }
+
+      // Simply drop the annotation if there's only one module.
+      if (moduleNames.size() < 2)
+        continue;
+
+      // Replace all instances and references to other modules with the
+      // first module.
+      replaceModuleReferences(circuitOp, moduleNames, nlaTable, innerSymTables);
+      nlaRemover.markNLAsInAnnotation(anno.getAttr());
+    }
+    if (newAnnotations.size() == annotations.size())
+      return failure();
+
+    // Update circuit annotations.
+    AnnotationSet newAnnoSet(newAnnotations, context);
+    newAnnoSet.applyToOperation(circuitOp);
+    return success();
+  }
+
+  std::string getName() const override { return "firrtl-force-dedup"; }
+  bool acceptSizeIncrease() const override { return true; }
+
+private:
+  /// Replace all references to modules in the dedup group with the canonical
+  /// module name
+  void replaceModuleReferences(CircuitOp circuitOp,
+                               ArrayRef<StringAttr> moduleNames,
+                               NLATable &nlaTable,
+                               hw::InnerSymbolTableCollection &innerSymTables) {
+    auto *tableOp = SymbolTable::getNearestSymbolTable(circuitOp);
+    auto &symbolTable = symbols.getSymbolTable(tableOp);
+    auto *context = circuitOp->getContext();
+    auto innerRefs = hw::InnerRefNamespace{symbolTable, innerSymTables};
+
+    // Collect the modules.
+    FModuleLike canonicalModule;
+    SmallVector<FModuleLike> modulesToReplace;
+    for (auto name : moduleNames) {
+      if (auto mod = symbolTable.lookup<FModuleLike>(name)) {
+        if (!canonicalModule)
+          canonicalModule = mod;
+        else
+          modulesToReplace.push_back(mod);
+      }
+    }
+    if (modulesToReplace.empty())
+      return;
+
+    // Replace all instance references.
+    auto canonicalName = canonicalModule.getModuleNameAttr();
+    auto canonicalRef = FlatSymbolRefAttr::get(canonicalName);
+    circuitOp.walk([&](InstanceOp instOp) {
+      auto moduleName = instOp.getModuleNameAttr().getAttr();
+      if (llvm::is_contained(moduleNames, moduleName) &&
+          moduleName != canonicalName) {
+        instOp.setModuleNameAttr(canonicalRef);
+        instOp.setPortNamesAttr(canonicalModule.getPortNamesAttr());
+      }
+    });
+
+    // Update NLAs to reference the canonical module instead of modules being
+    // removed using NLATable for better performance.
+    for (auto oldMod : modulesToReplace) {
+      SmallVector<hw::HierPathOp> nlaOps(
+          nlaTable.lookup(oldMod.getModuleNameAttr()));
+      for (auto nlaOp : nlaOps) {
+        nlaTable.erase(nlaOp);
+        StringAttr oldModName = oldMod.getModuleNameAttr();
+        StringAttr newModName = canonicalName;
+        SmallVector<Attribute, 4> newPath;
+        for (auto nameRef : nlaOp.getNamepath()) {
+          if (auto ref = dyn_cast<hw::InnerRefAttr>(nameRef)) {
+            if (ref.getModule() == oldModName) {
+              auto oldInst = innerRefs.lookupOp<FInstanceLike>(ref);
+              ref = hw::InnerRefAttr::get(newModName, ref.getName());
+              auto newInst = innerRefs.lookupOp<FInstanceLike>(ref);
+              if (oldInst && newInst) {
+                oldModName = oldInst.getReferencedModuleNameAttr();
+                newModName = newInst.getReferencedModuleNameAttr();
+              }
+            }
+            newPath.push_back(ref);
+          } else if (cast<FlatSymbolRefAttr>(nameRef).getAttr() == oldModName) {
+            newPath.push_back(FlatSymbolRefAttr::get(newModName));
+          } else {
+            newPath.push_back(nameRef);
+          }
+        }
+        nlaOp.setNamepathAttr(ArrayAttr::get(context, newPath));
+        nlaTable.addNLA(nlaOp);
+      }
+    }
+
+    // Mark NLAs in modules to be removed.
+    for (auto module : modulesToReplace) {
+      nlaRemover.markNLAsInOperation(module);
+      module->erase();
+    }
+  }
+
+  ::detail::SymbolCache symbols;
+  NLARemover nlaRemover;
+};
+
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // Reduction Registration
 //===----------------------------------------------------------------------===//
@@ -1129,6 +1310,7 @@ void firrtl::FIRRTLReducePatternDialectInterface::populateReducePatterns(
   // prioritized). For example, things that can knock out entire modules while
   // being cheap should be tried first (and thus have higher benefit), before
   // trying to tweak operands of individual arithmetic ops.
+  patterns.add<ForceDedup, 31>();
   patterns.add<PassReduction, 30>(
       getContext(),
       firrtl::createDropName({/*preserveMode=*/PreserveValues::None}), false,

lib/Reduce/GenericReductions.cpp

@@ -82,8 +82,9 @@ static std::unique_ptr<Pass> createSimpleCanonicalizerPass() {
 
 void circt::populateGenericReducePatterns(MLIRContext *context,
                                           ReducePatternSet &patterns) {
-  patterns.add<UnusedSymbolPruner, 40>();
-  patterns.add<PassReduction, 3>(context, createCSEPass());
-  patterns.add<PassReduction, 2>(context, createSimpleCanonicalizerPass());
+  patterns.add<PassReduction, 103>(context, createSymbolDCEPass());
+  patterns.add<PassReduction, 102>(context, createCSEPass());
+  patterns.add<PassReduction, 101>(context, createSimpleCanonicalizerPass());
+  patterns.add<UnusedSymbolPruner, 100>();
   patterns.add<OperationPruner, 1>();
 }

test/Dialect/FIRRTL/Reduction/force-dedup.mlir

@@ -0,0 +1,110 @@
+// UNSUPPORTED: system-windows
+//   See https://github.com/llvm/circt/issues/4129
+// RUN: circt-reduce %s --test /usr/bin/env --test-arg true --include firrtl-force-dedup | FileCheck %s
+
+// Test that the MustDedup reducer can consolidate modules in a dedup group
+// by replacing all module names with a single canonical module name.
+
+firrtl.circuit "MustDedup" attributes {annotations = [{
+    class = "firrtl.transforms.MustDeduplicateAnnotation",
+    modules = ["~MustDedup|Simple0", "~MustDedup|Simple1", "~MustDedup|Simple2"]
+  }]} {
+
+  // CHECK: hw.hierpath private @nla [@MustDedup::@simple1, @Simple0]
+  hw.hierpath private @nla [@MustDedup::@simple1, @Simple1]
+
+  // CHECK: firrtl.module private @Simple0
+  firrtl.module private @Simple0() {
+    %w = firrtl.wire : !firrtl.uint<1>
+  }
+
+  // CHECK-NOT: firrtl.module private @Simple1
+  firrtl.module private @Simple1() {
+    %w = firrtl.wire : !firrtl.uint<1>
+  }
+
+  // CHECK-NOT: firrtl.module private @Simple2
+  firrtl.module private @Simple2() {
+    %w = firrtl.wire : !firrtl.uint<1>
+  }
+
+  // CHECK: firrtl.module @MustDedup
+  firrtl.module @MustDedup() {
+    // CHECK: firrtl.instance simple0 @Simple0
+    firrtl.instance simple0 @Simple0()
+    // CHECK: firrtl.instance simple1 sym @simple1 {annotations = [{circt.nonlocal = @nla, class = "test"}]} @Simple0()
+    firrtl.instance simple1 sym @simple1 {annotations = [{circt.nonlocal = @nla, class = "test"}]} @Simple1()
+    // CHECK: firrtl.instance simple2 @Simple0
+    firrtl.instance simple2 @Simple2()
+  }
+}
+
+// Test with multiple NLAs referencing different modules in the dedup group
+firrtl.circuit "MultiNLA" attributes {annotations = [{
+    class = "firrtl.transforms.MustDeduplicateAnnotation",
+    modules = ["~MultiNLA|ModA", "~MultiNLA|ModB"]
+  }]} {
+
+  // CHECK: hw.hierpath private @nla1 [@MultiNLA::@instA, @ModA]
+  hw.hierpath private @nla1 [@MultiNLA::@instA, @ModA]
+  // CHECK: hw.hierpath private @nla2 [@MultiNLA::@instB, @ModA]
+  hw.hierpath private @nla2 [@MultiNLA::@instB, @ModB]
+
+  // CHECK: firrtl.module private @ModA
+  firrtl.module private @ModA() {
+    %w = firrtl.wire : !firrtl.uint<1>
+  }
+
+  // CHECK-NOT: firrtl.module private @ModB
+  firrtl.module private @ModB() {
+    %w = firrtl.wire : !firrtl.uint<1>
+  }
+
+  // CHECK: firrtl.module @MultiNLA
+  firrtl.module @MultiNLA() {
+    // CHECK: firrtl.instance instA sym @instA {annotations = [{circt.nonlocal = @nla1, class = "test1"}]} @ModA()
+    firrtl.instance instA sym @instA {annotations = [{circt.nonlocal = @nla1, class = "test1"}]} @ModA()
+    // CHECK: firrtl.instance instB sym @instB {annotations = [{circt.nonlocal = @nla2, class = "test2"}]} @ModA()
+    firrtl.instance instB sym @instB {annotations = [{circt.nonlocal = @nla2, class = "test2"}]} @ModB()
+  }
+}
+
+// Test with multiple NLAs referencing different nested modules in the dedup group
+firrtl.circuit "MultiSubNLA" attributes {annotations = [{
+    class = "firrtl.transforms.MustDeduplicateAnnotation",
+    modules = ["~MultiSubNLA|ModA", "~MultiSubNLA|ModB"]
+  }]} {
+
+  // CHECK: hw.hierpath private @nla1 [@MultiSubNLA::@instA, @ModA::@sub, @SubModA::@wire]
+  hw.hierpath private @nla1 [@MultiSubNLA::@instA, @ModA::@sub, @SubModA::@wire]
+  // CHECK: hw.hierpath private @nla2 [@MultiSubNLA::@instB, @ModA::@sub, @SubModA::@wire]
+  hw.hierpath private @nla2 [@MultiSubNLA::@instB, @ModB::@sub, @SubModB::@wire]
+
+  // CHECK: firrtl.module private @ModA
+  firrtl.module private @ModA() {
+    firrtl.instance sub sym @sub @SubModA()
+  }
+
+  // CHECK-NOT: firrtl.module private @ModB
+  firrtl.module private @ModB() {
+    firrtl.instance sub sym @sub @SubModB()
+  }
+
+  // CHECK: firrtl.module private @SubModA
+  firrtl.module private @SubModA() {
+    %w = firrtl.wire sym @wire {annotations = [{circt.nonlocal = @nla1, class = "test1"}]} : !firrtl.uint<1>
+  }
+
+  // CHECK: firrtl.module private @SubModB
+  firrtl.module private @SubModB() {
+    %w = firrtl.wire sym @wire {annotations = [{circt.nonlocal = @nla2, class = "test2"}]} : !firrtl.uint<1>
+  }
+
+  // CHECK: firrtl.module @MultiSubNLA
+  firrtl.module @MultiSubNLA() {
+    // CHECK: firrtl.instance instA sym @instA @ModA()
+    firrtl.instance instA sym @instA @ModA()
+    // CHECK: firrtl.instance instB sym @instB @ModA()
+    firrtl.instance instB sym @instB @ModB()
+  }
+}