[FIRRTL] Allow local targets to be multiply-instantiated.

The existing path support was built up based on the assumption that
every target is unique. That is true for FIRRTL produced by standard
Chisel code, which elaborates unique modules for each instance. We
definitely don't want to limit ourselves to this world, and we should
support targeting things that are multiply instantiated when it is not
ambiguous what we refer to.

This patch relaxes the single-instantiation constraints for local
targets, which refer to a module or something inside a module,
regardless of how many times or at what paths that particular module
was instantiated.

This required a couple changes through the pipeline:

ResolvePaths already had an early exit for the local path case, but
this needed to come before the single-instantiation check.

PrefixModules has support for cleaning up dead non-local annotations,
but we also need to support cleaning up cloned local trackers now.

LowerClasses needed a couple small changes to not enforce the
single-instantiation check in the local path case, and to build a
hierpath that just has a single element.

lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp

@@ -275,7 +275,8 @@ struct PathTracker {
   // and `owningModule`.
   FailureOr<bool> getOrComputeNeedsAltBasePath(Location loc,
                                                StringAttr moduleName,
-                                               FModuleOp owningModule);
+                                               FModuleOp owningModule,
+                                               bool isNonLocal);
   FModuleLike module;
 
   // Local data structures.
@@ -364,7 +365,8 @@ LogicalResult PathTracker::runOnModule() {
 
 FailureOr<bool>
 PathTracker::getOrComputeNeedsAltBasePath(Location loc, StringAttr moduleName,
-                                          FModuleOp owningModule) {
+                                          FModuleOp owningModule,
+                                          bool isNonLocal) {
 
   auto it = needsAltBasePathCache.find({moduleName, owningModule});
   if (it != needsAltBasePathCache.end())
@@ -385,7 +387,7 @@ PathTracker::getOrComputeNeedsAltBasePath(Location loc, StringAttr moduleName,
     }
     // If there is more than one instance of this module, then the path
     // operation is ambiguous, which is an error.
-    if (!node->hasOneUse()) {
+    if (isNonLocal && !node->hasOneUse()) {
       auto diag = mlir::emitError(loc)
                   << "unable to uniquely resolve target due "
                      "to multiple instantiation";
@@ -517,8 +519,8 @@ PathTracker::processPathTrackers(const AnnoTarget &target) {
     }
 
     // Check if we need an alternative base path.
-    auto needsAltBasePath =
-        getOrComputeNeedsAltBasePath(op->getLoc(), moduleName, owningModule);
+    auto needsAltBasePath = getOrComputeNeedsAltBasePath(
+        op->getLoc(), moduleName, owningModule, hierPathOp);
     if (failed(needsAltBasePath)) {
       error = true;
       return false;
@@ -528,6 +530,10 @@ PathTracker::processPathTrackers(const AnnoTarget &target) {
     // to the start of the annotation's NLA.
     InstanceGraphNode *node = instanceGraph.lookup(moduleName);
     while (true) {
+      // If it's not a non-local target, we don't have to append anything.
+      if (!hierPathOp)
+        break;
+
       // If we get to the owning module or the top, we're done.
       if (node->getModule() == owningModule || node->noUses())
         break;

lib/Dialect/FIRRTL/Transforms/PrefixModules.cpp

@@ -100,9 +100,10 @@ static StringRef getPrefix(Operation *module) {
 namespace {
 class PrefixModulesPass
     : public circt::firrtl::impl::PrefixModulesBase<PrefixModulesPass> {
-  void removeDeadAnnotations(StringAttr moduleName, Operation *op);
-  void renameModuleBody(std::string prefix, StringRef oldName,
-                        FModuleOp module);
+  void removeDeadAnnotations(StringAttr moduleName, Operation *op,
+                             bool removeTrackers = false);
+  void renameModuleBody(std::string prefix, StringRef oldName, FModuleOp module,
+                        bool removeTrackers = false);
   void renameModule(FModuleOp module);
   void renameExtModule(FExtModuleOp extModule);
   void renameMemModule(FMemModuleOp memModule);
@@ -135,16 +136,25 @@ class PrefixModulesPass
 /// function will remove all non-local annotations from the clone with a path
 /// that doesn't match.
 void PrefixModulesPass::removeDeadAnnotations(StringAttr moduleName,
-                                              Operation *op) {
+                                              Operation *op,
+                                              bool removeTrackers) {
   // A predicate to check if an annotation can be removed. If there is a
   // reference to a NLA, the NLA should either contain this module in its path,
   // if its an InstanceOp. Else, it must exist at the leaf of the NLA. Otherwise
   // the NLA reference can be removed, since its a spurious annotation, result
   // of cloning the original module.
   auto canRemoveAnno = [&](Annotation anno, Operation *op) -> bool {
     auto nla = anno.getMember("circt.nonlocal");
-    if (!nla)
-      return false;
+    if (!nla) {
+      // For clones, removeTrackers will be true, and if this is a local
+      // tracker, we need to remove it to avoid duplicating trackers.
+      if (removeTrackers && anno.isClass("circt.tracker")) {
+        op->emitRemark("removing tracker");
+        return true;
+      } else {
+        return false;
+      }
+    }
     auto nlaName = cast<FlatSymbolRefAttr>(nla).getAttr();
     auto nlaOp = nlaTable->getNLA(nlaName);
     if (!nlaOp) {
@@ -169,7 +179,8 @@ void PrefixModulesPass::removeDeadAnnotations(StringAttr moduleName,
 /// Applies the prefix to the module.  This will update the required prefixes of
 /// any referenced module in the prefix map.
 void PrefixModulesPass::renameModuleBody(std::string prefix, StringRef oldName,
-                                         FModuleOp module) {
+                                         FModuleOp module,
+                                         bool removeTrackers) {
   auto *context = module.getContext();
 
   // If we are renaming the body of this module, we need to mark that we have
@@ -182,7 +193,7 @@ void PrefixModulesPass::renameModuleBody(std::string prefix, StringRef oldName,
   // Remove spurious NLA references from the module ports and the module itself.
   // Some of the NLA references become invalid after a module is cloned, based
   // on the instance.
-  removeDeadAnnotations(thisMod, module);
+  removeDeadAnnotations(thisMod, module, removeTrackers);
 
   mlir::AttrTypeReplacer replacer;
   replacer.addReplacement(
@@ -309,7 +320,8 @@ void PrefixModulesPass::renameModule(FModuleOp module) {
     nlaTable->addModule(moduleClone);
     fixNLAsRootedAt(oldModName, newModNameAttr);
     // Each call to this function could invalidate the `prefixes` reference.
-    renameModuleBody((outerPrefix + innerPrefix).str(), oldName, moduleClone);
+    renameModuleBody((outerPrefix + innerPrefix).str(), oldName, moduleClone,
+                     /*removeTrackers=*/true);
   }
 
   auto prefixFull = (firstPrefix + innerPrefix).str();

lib/Dialect/FIRRTL/Transforms/ResolvePaths.cpp

@@ -41,6 +41,12 @@ struct PathResolver {
                                 const AnnoPathValue &target,
                                 FlatSymbolRefAttr &result) {
 
+    // If the path is empty then this is a local reference and we should not
+    // construct a HierPathOp.
+    if (target.instances.empty()) {
+      return success();
+    }
+
     // We want to root this path at the top level module, or in the case of an
     // unreachable module, we settle for as high as we can get.
     auto module = target.ref.getModule();
@@ -69,12 +75,6 @@ struct PathResolver {
       node = (*node->usesBegin())->getParent();
     }
 
-    // If the path is empty then this is a local reference and we should not
-    // construct a HierPathOp.
-    if (target.instances.empty()) {
-      return success();
-    }
-
     // Transform the instances into a list of FlatSymbolRefs.
     SmallVector<Attribute> insts;
     insts.reserve(target.instances.size());

test/Dialect/FIRRTL/lower-classes.mlir

@@ -158,6 +158,7 @@ firrtl.circuit "PathModule" {
   // CHECK: hw.hierpath private [[WIRE_PATH:@.+]] [@PathModule::[[WIRE_SYM:@.+]]]
   // CHECK: hw.hierpath private [[VECTOR_PATH:@.+]] [@PathModule::[[VECTOR_SYM:@.+]]]
   // CHECK: hw.hierpath private [[INST_PATH:@.+]] [@PathModule::@child]
+  // CHECK: hw.hierpath private [[LOCAL_PATH:@.+]] [@Child]
   // CHECK: hw.hierpath private [[MODULE_PATH:@.+]] [@PathModule::@child, @Child::[[NONLOCAL_SYM:@.+]]]
 
   // CHECK: firrtl.module @PathModule(in %in: !firrtl.uint<1> sym [[PORT_SYM]]) {
@@ -209,7 +210,7 @@ firrtl.circuit "PathModule" {
     // CHECK: om.path_create member_instance %basepath [[INST_PATH]]
     // CHECK: om.path_create member_instance %basepath [[INST_PATH]]
     // CHECK: om.path_create instance %basepath [[INST_PATH]]
-    // CHECK: om.path_create reference %basepath [[NONLOCAL_PATH]]
+    // CHECK: om.path_create reference %basepath [[LOCAL_PATH]]
     %instance_member_instance = firrtl.path member_instance distinct[4]<>
     %instance_member_reference = firrtl.path member_reference distinct[4]<>
     %instance = firrtl.path instance distinct[4]<>
@@ -471,3 +472,25 @@ firrtl.circuit "PathTargetReplaced" {
   firrtl.module private @WillBeReplaced(out %output: !firrtl.integer) {
   }
 }
+
+// CHECK-LABEL: firrtl.circuit "LocalPath"
+firrtl.circuit "LocalPath" {
+  // CHECK: hw.hierpath private [[MODULE_NLA:@.+]] [@Child]
+  // CHECK: hw.hierpath private [[WIRE_NLA:@.+]] [@Child::[[WIRE_SYM:@.+]]]
+
+  firrtl.module @Child() attributes {annotations = [{class = "circt.tracker", id = distinct[0]<>}]} {
+    // CHECK: firrtl.wire sym [[WIRE_SYM]]
+    %wire = firrtl.wire {annotations = [{class = "circt.tracker", id = distinct[1]<>}]} : !firrtl.uint<8>
+  }
+
+  firrtl.module @LocalPath() {
+    // CHECK: om.path_create instance %basepath [[MODULE_NLA]]
+    %0 = firrtl.path instance distinct[0]<>
+
+    // CHECK: om.path_create reference %basepath [[WIRE_NLA]]
+    %1 = firrtl.path reference distinct[1]<>
+
+    firrtl.instance child0 @Child()
+    firrtl.instance child1 @Child()
+  }
+}

test/Dialect/FIRRTL/prefix-modules.mlir

@@ -427,3 +427,46 @@ firrtl.circuit "RewriteInnerNameRefs" {
     %wire = firrtl.wire sym @wire : !firrtl.uint<1>
   }
 }
+
+firrtl.circuit "LocalTrackerRemoval" {
+  firrtl.module @LocalTrackerRemoval() {
+    firrtl.instance child @Child()
+  }
+
+  firrtl.module @Child() attributes {
+    annotations = [
+      {
+        class = "sifive.enterprise.firrtl.NestedPrefixModulesAnnotation",
+        prefix = "Prefix1_",
+        inclusive = true
+      }
+    ]
+  } {
+    firrtl.instance grandchild @GrandChild()
+    firrtl.instance target @Target()
+  }
+
+  firrtl.module @GrandChild() attributes {
+    annotations = [
+      {
+        class = "sifive.enterprise.firrtl.NestedPrefixModulesAnnotation",
+        prefix = "Prefix2_",
+        inclusive = true
+      }
+    ]
+  } {
+    firrtl.instance target @Target()
+  }
+
+  // CHECK: firrtl.module @Prefix1_Target() attributes {annotations = [{class = "circt.tracker", id = distinct[0]<>}]}
+  // CHECK: firrtl.module @Prefix1_Prefix2_Target() {
+  firrtl.module @Target() attributes {
+    annotations = [
+      {
+        class = "circt.tracker",
+	id = distinct[1]<>
+      }
+    ]
+  } {
+  }
+}

test/Dialect/FIRRTL/resolve-paths-errors.mlir

@@ -75,13 +75,16 @@ firrtl.module @VectorTarget(in %a : !firrtl.vector<uint<1>, 1>) {
 firrtl.circuit "AmbiguousPath" {
 firrtl.module @AmbiguousPath() {
     // expected-error @below {{unable to uniquely resolve target due to multiple instantiation}}
-    %0 = firrtl.unresolved_path "OMReferenceTarget:~AmbiguousPath|Child"
+    %0 = firrtl.unresolved_path "OMReferenceTarget:~AmbiguousPath|Child/grandchild:GrandChild"
     // expected-note @below {{instance here}}
     firrtl.instance child0 @Child()
     // expected-note @below {{instance here}}
     firrtl.instance child1 @Child()
 }
-firrtl.module @Child() {}
+firrtl.module @Child() {
+    firrtl.instance grandchild @GrandChild()
+}
+firrtl.module @GrandChild() {}
 }
 
 // -----
@@ -102,18 +105,17 @@ firrtl.class @Test(){
 
 firrtl.circuit "UpwardPath" {
 firrtl.module @UpwardPath() {
-  %wire = firrtl.wire : !firrtl.uint<8>
   firrtl.instance child @Child()
 }
 
 firrtl.module @Child() {
   %om = firrtl.object @OM()
-
+  %wire = firrtl.wire : !firrtl.uint<8>
 }
 
 firrtl.class @OM(){
   // expected-error @below {{unable to resolve path relative to owning module "Child"}}
-  %0 = firrtl.unresolved_path "OMReferenceTarget:~UpwardPath|UpwardPath>wire"
+  %0 = firrtl.unresolved_path "OMReferenceTarget:~UpwardPath|UpwardPath/child:Child>wire"
 }
 }
 

test/Dialect/FIRRTL/resolve-paths.mlir

@@ -91,6 +91,30 @@ firrtl.module @Child() {
 
 // -----
 
+// CHECK-LABEL: firrtl.circuit "LocalPath"
+firrtl.circuit "LocalPath" {
+// CHECK: firrtl.module @LocalPath()
+firrtl.module @LocalPath() {
+    // CHECK: %0 = firrtl.path instance distinct[0]<>
+    %0 = firrtl.unresolved_path "OMInstanceTarget:~LocalPath|Child"
+
+    // CHECK: %1 = firrtl.path reference distinct[1]<>
+    %1 = firrtl.unresolved_path "OMReferenceTarget:~LocalPath|Child>wire"
+
+    // CHECK: firrtl.instance child0 @Child()
+    // CHECK: firrtl.instance child1 @Child()
+    firrtl.instance child0 @Child()
+    firrtl.instance child1 @Child()
+}
+// CHECK: firrtl.module @Child() attributes {annotations = [{class = "circt.tracker", id = distinct[0]<>}]}
+firrtl.module @Child() {
+    // CHECK: %wire = firrtl.wire {annotations = [{class = "circt.tracker", id = distinct[1]<>}]} : !firrtl.uint<8>
+    %wire = firrtl.wire : !firrtl.uint<8>
+}
+}
+
+// -----
+
 // CHECK-LABEL: firrtl.circuit "PathMinimization"
 firrtl.circuit "PathMinimization" {
 // CHECK: firrtl.module @PathMinimization()