From 6cf1dd040bb4c69e499dbe5ac7cc9de526518df4 Mon Sep 17 00:00:00 2001
From: Mike Urbach <mike.urbach@sifive.com>
Date: Tue, 15 Oct 2024 14:28:51 -0700
Subject: [PATCH] Also dedupe local trackers so LowerClasses doesn't see
 duplicates.

---
 lib/Dialect/FIRRTL/Transforms/Dedup.cpp | 18 ++++++++++++------
 test/Dialect/FIRRTL/dedup.mlir          |  7 +++++++
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/lib/Dialect/FIRRTL/Transforms/Dedup.cpp b/lib/Dialect/FIRRTL/Transforms/Dedup.cpp
index 11cdf4c99d91..8bcf2ec20f6a 100644
--- a/lib/Dialect/FIRRTL/Transforms/Dedup.cpp
+++ b/lib/Dialect/FIRRTL/Transforms/Dedup.cpp
@@ -1233,7 +1233,8 @@ struct Deduper {
   void copyAnnotations(FModuleLike toModule, AnnoTarget to,
                        FModuleLike fromModule, AnnotationSet annos,
                        SmallVectorImpl<Annotation> &newAnnotations,
-                       SmallPtrSetImpl<Attribute> &dontTouches) {
+                       SmallPtrSetImpl<Attribute> &dontTouches,
+                       SmallPtrSetImpl<Attribute> &localTrackerIds) {
     for (auto anno : annos) {
       if (anno.isClass(dontTouchAnnoClass)) {
         // Remove the nonlocal field of the annotation if it has one, since this
@@ -1254,7 +1255,10 @@ struct Deduper {
       // If the annotation is a local tracker, don't make it non-local. This
       // allows trackers that refer to all instances when the user wants that.
       if (anno.isClass("circt.tracker")) {
-        newAnnotations.push_back(anno);
+        auto [it, inserted] =
+            localTrackerIds.insert(anno.getMember<DistinctAttr>("id"));
+        if (inserted)
+          newAnnotations.push_back(anno);
         continue;
       }
 
@@ -1272,16 +1276,18 @@ struct Deduper {
     // This is a list of all the annotations which will be added to `to`.
     SmallVector<Annotation> newAnnotations;
 
-    // We have special case handling of DontTouch to prevent it from being
-    // turned into a non-local annotation, and to remove duplicates.
+    // We have special case handling of DontTouch and local trackers to prevent
+    // them from being turned into a non-local annotation, and to remove
+    // duplicates.
     llvm::SmallPtrSet<Attribute, 4> dontTouches;
+    llvm::SmallPtrSet<Attribute, 4> localTrackerIds;
 
     // Iterate the annotations, transforming most annotations into non-local
     // ones.
     copyAnnotations(toModule, to, toModule, toAnnos, newAnnotations,
-                    dontTouches);
+                    dontTouches, localTrackerIds);
     copyAnnotations(toModule, to, fromModule, fromAnnos, newAnnotations,
-                    dontTouches);
+                    dontTouches, localTrackerIds);
 
     // Copy over all the new annotations.
     if (!newAnnotations.empty())
diff --git a/test/Dialect/FIRRTL/dedup.mlir b/test/Dialect/FIRRTL/dedup.mlir
index 525df99c3022..fcf368a953e6 100644
--- a/test/Dialect/FIRRTL/dedup.mlir
+++ b/test/Dialect/FIRRTL/dedup.mlir
@@ -780,6 +780,7 @@ firrtl.circuit "AllowLocalTrackers" {
   firrtl.module @AllowLocalTrackers() {
     firrtl.instance foo0 @Foo0()
     firrtl.instance foo1 @Foo1()
+    firrtl.instance foo2 @Foo2()
   }
 
   // CHECK-NOT: hw.hierpath
@@ -788,6 +789,7 @@ firrtl.circuit "AllowLocalTrackers" {
   // CHECK: firrtl.mem
   // CHECK-SAME: {class = "circt.tracker", id = distinct[0]<>}
   // CHECK-SAME: {class = "circt.tracker", id = distinct[1]<>}
+  // CHECK-NOT: {class = "circt.tracker", id = distinct[1]<>}
   firrtl.module private @Foo0() {
       %mem_RW0 = firrtl.mem Undefined {annotations = [{class = "circt.tracker", id = distinct[0]<>}], depth = 24 : i64, name = "mem", portNames = ["RW0"], readLatency = 1 : i32, writeLatency = 1 : i32} : !firrtl.bundle<addr: uint<5>, en: uint<1>, clk: clock, rdata flip: uint<72>, wmode: uint<1>, wdata: uint<72>, wmask: uint<1>>
   }
@@ -796,4 +798,9 @@ firrtl.circuit "AllowLocalTrackers" {
   firrtl.module private @Foo1() {
       %mem_RW0 = firrtl.mem Undefined {annotations = [{class = "circt.tracker", id = distinct[1]<>}], depth = 24 : i64, name = "mem", portNames = ["RW0"], readLatency = 1 : i32, writeLatency = 1 : i32} : !firrtl.bundle<addr: uint<5>, en: uint<1>, clk: clock, rdata flip: uint<72>, wmode: uint<1>, wdata: uint<72>, wmask: uint<1>>
   }
+
+  // CHECK-NOT: @Foo2
+  firrtl.module private @Foo2() {
+      %mem_RW0 = firrtl.mem Undefined {annotations = [{class = "circt.tracker", id = distinct[1]<>}], depth = 24 : i64, name = "mem", portNames = ["RW0"], readLatency = 1 : i32, writeLatency = 1 : i32} : !firrtl.bundle<addr: uint<5>, en: uint<1>, clk: clock, rdata flip: uint<72>, wmode: uint<1>, wdata: uint<72>, wmask: uint<1>>
+  }
 }