[FIRRTL] Update NLA during prefix module pass (llvm#2030)

This change updates the `NonLocalAnchor` symbol reference to modules that are
 renamed in prefix module pass. 
The `nla` op contains reference to module symbols, which might be left
 invalid after module renaming. For example, 
`firrtl.nla @nla_2 [@top, @aardvark, @Zebra] ["test", "test1", "Zebra"]` 

There can be multiple copies of any module created after the prefix module
 pass. The `firrtl.nla` renaming is context sensitive and depends on the
 instance op, which is anchored to the `nla`.

This renaming is also based on the assumption, that the modules in the
 `firrtl.nla` begin with the top level module. Otherwise, the `firrtl.nla`
 might be ambiguous if multiple clones of the first module is created.

lib/Dialect/FIRRTL/Transforms/PrefixModules.cpp

@@ -108,6 +108,9 @@ class PrefixModulesPass : public PrefixModulesBase<PrefixModulesPass> {
   /// Cached instance graph analysis.
   InstanceGraph *instanceGraph = nullptr;
 
+  /// Map of symbol name to NonLocalAnchor op.
+  llvm::StringMap<Operation *> nlaMap;
+
   /// Boolean keeping track of any name changes.
   bool anythingChanged = false;
 };
@@ -135,13 +138,40 @@ void PrefixModulesPass::renameModuleBody(std::string prefix, FModuleOp module) {
       // Skip this rename if the instance is an external module.
       if (!target)
         return;
-
       // Record that we must prefix the target module with the current prefix.
       recordPrefix(prefixMap, target.getName(), prefix);
 
       // Fixup this instance op to use the prefixed module name.  Note that the
       // referenced FModuleOp will be renamed later.
       auto newTarget = (prefix + getPrefix(target) + target.getName()).str();
+      AnnotationSet instAnnos(instanceOp);
+      // If the instance has NonLocalAnchor, then update its module name also.
+      // There can be multiple NonLocalAnchors attached to the instance op.
+
+      for (Annotation anno : instAnnos) {
+        if (anno.isClass("circt.nonlocal"))
+          if (auto nla = anno.getMember("circt.nonlocal")) {
+            auto nlaName = nla.cast<FlatSymbolRefAttr>().getValue();
+            auto f = nlaMap.find(nlaName);
+            if (f == nlaMap.end())
+              instanceOp.emitError("cannot find NonLocalAnchor :" + nlaName);
+            else {
+              auto nlaOp = dyn_cast<NonLocalAnchor>(f->second);
+              // Iterate over the modules of the NonLocalAnchor op, and update
+              // it.
+              SmallVector<Attribute, 4> newMods;
+              for (auto oldMod : nlaOp.modpath()) {
+                if (instanceOp.moduleNameAttr() ==
+                    oldMod.cast<FlatSymbolRefAttr>())
+                  newMods.push_back(FlatSymbolRefAttr::get(context, newTarget));
+                else
+                  newMods.push_back(oldMod.cast<FlatSymbolRefAttr>());
+              }
+              nlaOp->setAttr("modpath", ArrayAttr::get(context, newMods));
+            }
+          }
+      }
+
       instanceOp.moduleNameAttr(FlatSymbolRefAttr::get(context, newTarget));
     }
   });
@@ -231,11 +261,29 @@ void PrefixModulesPass::runOnOperation() {
   instanceGraph = &getAnalysis<InstanceGraph>();
   auto circuitOp = getOperation();
 
+  // Record all the NLA ops in the circt.
+  for (auto nla : circuitOp.body().getOps<NonLocalAnchor>())
+    nlaMap[nla.sym_name()] = nla;
+
   // If the main module is prefixed, we have to update the CircuitOp.
   auto mainModule = instanceGraph->getTopLevelModule();
   auto prefix = getPrefix(mainModule);
-  if (!prefix.empty())
-    circuitOp.nameAttr(StringAttr::get(context, prefix + circuitOp.name()));
+  if (!prefix.empty()) {
+    auto newMainModuleName = ((prefix + circuitOp.name()).str());
+    circuitOp.nameAttr(StringAttr::get(context, newMainModuleName));
+    // Now update all the NLAs that have the top level module symbol.
+    for (auto &n : nlaMap) {
+      auto nla = cast<NonLocalAnchor>(n.second);
+      auto oldMods = nla.modpath();
+      if (oldMods.empty())
+        continue;
+      SmallVector<Attribute, 4> newMods(oldMods.begin(), oldMods.end());
+      if (nla.modpath()[0].cast<FlatSymbolRefAttr>().getValue().equals(
+              mainModule.moduleName()))
+        newMods[0] = FlatSymbolRefAttr::get(context, newMainModuleName);
+      nla->setAttr("modpath", ArrayAttr::get(context, newMods));
+    }
+  }
 
   // Walk all Modules in a top-down order.  For each module, look at the list of
   // required prefixes to be applied.

test/Dialect/FIRRTL/prefix-modules.mlir

@@ -187,3 +187,49 @@ firrtl.circuit "GCTInterfacePrefix"
     firrtl.instance dut @DUT()
   }
 }
+
+// CHECK: firrtl.circuit "T_NLATop"
+firrtl.circuit "NLATop" {
+
+  firrtl.nla @nla [@NLATop, @Aardvark, @Zebra] ["test", "test", "Zebra"]
+  firrtl.nla @nla_1 [@NLATop, @Aardvark, @Zebra] ["test", "test_1", "Zebra"]
+  // CHECK: firrtl.nla @nla [@T_NLATop, @T_Aardvark, @T_A_Z_Zebra] ["test", "test", "Zebra"]
+  // CHECK: firrtl.nla @nla_1 [@T_NLATop, @T_Aardvark, @T_A_Z_Zebra] ["test", "test_1", "Zebra"]
+  // CHECK: firrtl.module @T_NLATop
+  firrtl.module @NLATop()
+    attributes {annotations = [{
+      class = "sifive.enterprise.firrtl.NestedPrefixModulesAnnotation",
+      prefix = "T_",
+      inclusive = true
+    }]} {
+
+    // CHECK:  firrtl.instance test  {annotations = [{circt.nonlocal = @nla, class = "circt.nonlocal"}, {circt.nonlocal = @nla_1, class = "circt.nonlocal"}]} @T_Aardvark()
+    firrtl.instance test  {annotations = [{circt.nonlocal = @nla, class = "circt.nonlocal"}, {circt.nonlocal = @nla_1, class = "circt.nonlocal"} ]}@Aardvark()
+
+    // CHECK: firrtl.instance test2 @T_Z_Zebra()
+    firrtl.instance test2 @Zebra()
+  }
+
+  // CHECK: firrtl.module @T_Aardvark
+  firrtl.module @Aardvark()
+    attributes {annotations = [{
+      class = "sifive.enterprise.firrtl.NestedPrefixModulesAnnotation",
+      prefix = "A_",
+      inclusive = false
+    }]} {
+
+    // CHECK:  firrtl.instance test  {annotations = [{circt.nonlocal = @nla, class = "circt.nonlocal"}]} @T_A_Z_Zebra()
+    firrtl.instance test {annotations = [{circt.nonlocal = @nla, class = "circt.nonlocal"}]}@Zebra()
+    firrtl.instance test1 {annotations = [{circt.nonlocal = @nla_1, class = "circt.nonlocal"}]}@Zebra()
+  }
+
+  // CHECK: firrtl.module @T_Z_Zebra
+  // CHECK: firrtl.module @T_A_Z_Zebra
+  firrtl.module @Zebra() 
+    attributes {annotations = [{
+      class = "sifive.enterprise.firrtl.NestedPrefixModulesAnnotation",
+      prefix = "Z_",
+      inclusive = true
+    }]} {
+  }
+}