[FIRRTL] Don't force non-local trackers in Dedup.

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
@@ -1251,6 +1252,16 @@ struct Deduper {
         targetMap[nla.getAttr()].insert(to);
         continue;
       }
+      // 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")) {
+        auto [it, inserted] =
+            localTrackerIds.insert(anno.getMember<DistinctAttr>("id"));
+        if (inserted)
+          newAnnotations.push_back(anno);
+        continue;
+      }
+
       // Otherwise make the annotation non-local and add it to the set.
       makeAnnotationNonLocal(toModule.getModuleNameAttr(), to, fromModule, anno,
                              newAnnotations);
@@ -1265,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())

test/Dialect/FIRRTL/dedup.mlir

@@ -773,3 +773,34 @@ firrtl.circuit "NoDedupPublic" {
   firrtl.module @DUTLike() {}
   firrtl.module @NotDUT() {}
 }
+
+// CHECK-LABEL: "AllowLocalTrackers"
+// If the user requested local trackers, don't make them non-local.
+firrtl.circuit "AllowLocalTrackers" {
+  firrtl.module @AllowLocalTrackers() {
+    firrtl.instance foo0 @Foo0()
+    firrtl.instance foo1 @Foo1()
+    firrtl.instance foo2 @Foo2()
+  }
+
+  // CHECK-NOT: hw.hierpath
+
+  // CHECK: firrtl.module private @Foo0
+  // 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>>
+  }
+
+  // CHECK-NOT: @Foo1
+  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>>
+  }
+}