[FIRRTL][Dedup] Fix cloning annotations for more context (#4285)

When we merge two modules, sometimes we have to add more context to a non-local annotation. When this happens, we create a new hierarchical path operation and update all annotations to use the new path. In order for us to find all annotations which need to be updated, we store a mapping, `targetMap`, from each HierPathOp to the annotations that reference them. This fixes a bug where we were not keeping the `targetMap` updated when we copy a non-local annotation when deduping operations. This is a regression from 46e9b21, where we stopped updating the target map. See 46e9b21#diff-d0c16a1337e2a6200ae0f4c763fadb7ddaa3f8b389130aa04cd1e61ff2fba7eaL896 for the exact line that was removed.
llvm · Nov 10, 2022 · 79a0a6a · 79a0a6a
1 parent 7f0c65b
commit 79a0a6a
Show file tree

Hide file tree

Showing 2 changed files with 101 additions and 1 deletion.
diff --git a/lib/Dialect/FIRRTL/Transforms/Dedup.cpp b/lib/Dialect/FIRRTL/Transforms/Dedup.cpp
@@ -947,8 +947,9 @@ struct Deduper {
       }
       // If the annotation is already non-local, we add it as is.  It is already
       // added to the target map.
-      if (anno.getMember("circt.nonlocal")) {
+      if (auto nla = anno.getMember<FlatSymbolRefAttr>("circt.nonlocal")) {
         newAnnotations.push_back(anno);
+        targetMap[nla.getAttr()].insert(to);
         continue;
       }
       // Otherwise make the annotation non-local and add it to the set.

diff --git a/test/Dialect/FIRRTL/dedup.mlir b/test/Dialect/FIRRTL/dedup.mlir
@@ -286,6 +286,105 @@ firrtl.circuit "Context" {
 }
 
 
+// When an annotation is already non-local, and is copied over to another
+// module, and in further dedups force us to add more context to the
+// hierarchical path, the target of the annotation should be updated to use the
+// new NLA.
+// CHECK-LABEL: firrtl.circuit "Context"
+firrtl.circuit "Context" {
+
+  // CHECK-NOT: firrtl.hierpath private @nla0
+  firrtl.hierpath private @nla0 [@Context0::@leaf0, @ContextLeaf0::@w0]
+  // CHECK-NOT: firrtl.hierpath private @nla1
+  firrtl.hierpath private @nla1 [@Context1::@leaf1, @ContextLeaf1::@w1]
+
+  // CHECK: firrtl.hierpath private [[NLA0:@.+]] [@Context::@context1, @Context0::@leaf0, @ContextLeaf0::@w0]
+  // CHECK: firrtl.hierpath private [[NLA1:@.+]] [@Context::@context0, @Context0::@leaf0, @ContextLeaf0::@w0]
+
+  // CHECK: firrtl.module @ContextLeaf0()
+  firrtl.module @ContextLeaf0() {
+    // CHECK: %w0 = firrtl.wire sym @w0  {annotations = [
+    // CHECK-SAME: {circt.nonlocal = [[NLA1]], class = "fake0"}
+    // CHECK-SAME: {circt.nonlocal = [[NLA0]], class = "fake1"}]}
+    %w0 = firrtl.wire sym @w0 {annotations = [
+      {circt.nonlocal = @nla0, class = "fake0"}]}: !firrtl.uint<3>
+  }
+
+  firrtl.module @ContextLeaf1() {
+    %w1 = firrtl.wire sym @w1 {annotations = [
+      {circt.nonlocal = @nla1, class = "fake1"}]}: !firrtl.uint<3>
+  }
+
+  firrtl.module @Context0() {
+    firrtl.instance leaf0 sym @leaf0 @ContextLeaf0()
+  }
+
+  firrtl.module @Context1() {
+    firrtl.instance leaf1 sym @leaf1 @ContextLeaf1()
+  }
+
+  firrtl.module @Context() {
+    firrtl.instance context0 @Context0()
+    firrtl.instance context1 @Context1()
+  }
+}
+
+
+// This is a larger version of the above test using 3 modules.
+// CHECK-LABEL: firrtl.circuit "DuplicateNLAs"
+firrtl.circuit "DuplicateNLAs" {
+  // CHECK-NOT: firrtl.hierpath private @annos_nla_1 [@Mid_1::@core, @Core_1]
+  // CHECK-NOT: firrtl.hierpath private @annos_nla_2 [@Mid_2::@core, @Core_2]
+  // CHECK-NOT: firrtl.hierpath private @annos_nla_3 [@Mid_3::@core, @Core_3]
+  firrtl.hierpath private @annos_nla_1 [@Mid_1::@core, @Core_1]
+  firrtl.hierpath private @annos_nla_2 [@Mid_2::@core, @Core_2]
+  firrtl.hierpath private @annos_nla_3 [@Mid_3::@core, @Core_3]
+
+  // CHECK: firrtl.hierpath private [[NLA0:@.+]] [@DuplicateNLAs::@core_3, @Mid_1::@core, @Core_1]
+  // CHECK: firrtl.hierpath private [[NLA1:@.+]] [@DuplicateNLAs::@core_2, @Mid_1::@core, @Core_1]
+  // CHECK: firrtl.hierpath private [[NLA2:@.+]] [@DuplicateNLAs::@core_1, @Mid_1::@core, @Core_1]
+
+  firrtl.module @DuplicateNLAs() {
+    firrtl.instance core_1 sym @core_1 @Mid_1()
+    firrtl.instance core_2 sym @core_2 @Mid_2()
+    firrtl.instance core_3 sym @core_3 @Mid_3()
+  }
+
+  firrtl.module private @Mid_1() {
+    firrtl.instance core sym @core @Core_1()
+  }
+
+  firrtl.module private @Mid_2() {
+    firrtl.instance core sym @core @Core_2()
+  }
+
+  firrtl.module private @Mid_3() {
+    firrtl.instance core sym @core @Core_3()
+  }
+
+  // CHECK: firrtl.module private @Core_1() attributes {annotations = [
+  // CHECK-SAME: {circt.nonlocal = [[NLA2]], class = "SomeAnno1"}
+  // CHECK-SAME: {circt.nonlocal = [[NLA1]], class = "SomeAnno2"}
+  // CHECK-SAME: {circt.nonlocal = [[NLA0]], class = "SomeAnno3"}
+  firrtl.module private @Core_1() attributes {
+    annotations = [
+      {circt.nonlocal = @annos_nla_1, class = "SomeAnno1"}
+    ]
+  } { }
+
+  firrtl.module private @Core_2() attributes {
+    annotations = [
+      {circt.nonlocal = @annos_nla_2, class = "SomeAnno2"}
+    ]
+  } { }
+
+  firrtl.module private @Core_3() attributes {
+    annotations = [
+      {circt.nonlocal = @annos_nla_3, class = "SomeAnno3"}
+    ]
+  } { }
+}
+
 // External modules should dedup and fixup any NLAs.
 // CHECK: firrtl.circuit "ExtModuleTest"
 firrtl.circuit "ExtModuleTest" {