[ModuleInliner][BugFix] Correctly update NLA after recursive inline (#…

…4942)

This fixes a bug with updating the NLA, when modules are recursively inlined
 to a single parent. The NLA was being updated with respect to a temporary and
 stale parent, the retop must be called with the original parent.
Also handle renamed instances correctly, by correctly tracking the NLAs that
 are contextually valid at any inlined instance.
Fixes #4920, #4915

lib/Dialect/FIRRTL/Transforms/ModuleInliner.cpp

@@ -554,7 +554,7 @@ class Inliner {
   /// the target, and does not trigger inlining on the target itself.
   void inlineInto(StringRef prefix, OpBuilder &b, IRMapping &mapper,
                   BackedgeBuilder &beb, SmallVectorImpl<Backedge> &edges,
-                  FModuleOp target,
+                  FModuleOp target, FModuleOp inlineToParent,
                   DenseMap<Attribute, Attribute> &symbolRenames,
                   ModuleNamespace &moduleNamespace);
 
@@ -676,6 +676,12 @@ bool Inliner::renameInstance(
     StringRef prefix, InstanceOp oldInst, InstanceOp newInst,
     ModuleNamespace &moduleNamespace,
     const DenseMap<Attribute, Attribute> &symbolRenames) {
+  // Add this instance to the activeHierpaths. This ensures that NLAs that this
+  // instance participates in will be updated correctly.
+  auto parentActivePaths = activeHierpaths;
+  if (auto instSym = getInnerSymName(oldInst))
+    setActiveHierPaths(oldInst->getParentOfType<FModuleOp>().getNameAttr(),
+                       instSym);
   // List of HierPathOps that are valid based on the InstanceOp being inlined
   // and the InstanceOp which is being replaced after inlining. That is the set
   // of HierPathOps that is common between these two.
@@ -742,6 +748,7 @@ bool Inliner::renameInstance(
         nlaList[en.index()] = newSym.cast<StringAttr>();
     }
   }
+  activeHierpaths = std::move(parentActivePaths);
   return symbolChanged;
 }
 
@@ -1009,7 +1016,7 @@ void Inliner::flattenInstances(FModuleOp module) {
 // NOLINTNEXTLINE(misc-no-recursion)
 void Inliner::inlineInto(StringRef prefix, OpBuilder &b, IRMapping &mapper,
                          BackedgeBuilder &beb, SmallVectorImpl<Backedge> &edges,
-                         FModuleOp target,
+                         FModuleOp target, FModuleOp inlineToParent,
                          DenseMap<Attribute, Attribute> &symbolRenames,
                          ModuleNamespace &moduleNamespace) {
   auto moduleName = target.getNameAttr();
@@ -1066,7 +1073,9 @@ void Inliner::inlineInto(StringRef prefix, OpBuilder &b, IRMapping &mapper,
     if (!rootMap[childModule.getNameAttr()].empty()) {
       for (auto sym : rootMap[childModule.getNameAttr()]) {
         auto &mnla = nlaMap[sym];
-        sym = mnla.reTop(target);
+        // Retop to the new parent, which is the topmost module (and not
+        // immediate parent) in case of recursive inlining.
+        sym = mnla.reTop(inlineToParent);
         StringAttr instSym = getInnerSymName(instance);
         if (!instSym) {
           instSym = StringAttr::get(
@@ -1099,7 +1108,7 @@ void Inliner::inlineInto(StringRef prefix, OpBuilder &b, IRMapping &mapper,
                   moduleNamespace);
     } else {
       inlineInto(nestedPrefix, b, mapper, beb, edges, childModule,
-                 symbolRenames, moduleNamespace);
+                 inlineToParent, symbolRenames, moduleNamespace);
     }
     currentPath.pop_back();
     activeHierpaths = parentActivePaths;
@@ -1195,8 +1204,10 @@ void Inliner::inlineInstances(FModuleOp parent) {
       flattenInto(nestedPrefix, b, mapper, beb, edges, target, {},
                   moduleNamespace);
     } else {
-      inlineInto(nestedPrefix, b, mapper, beb, edges, target, symbolRenames,
-                 moduleNamespace);
+      // Recursively inline all the child modules under `parent`, that are
+      // marked to be inlined.
+      inlineInto(nestedPrefix, b, mapper, beb, edges, target, parent,
+                 symbolRenames, moduleNamespace);
     }
     currentPath.pop_back();
     activeHierpaths = parentActivePaths;

test/Dialect/FIRRTL/inliner.mlir

@@ -1139,3 +1139,54 @@ firrtl.circuit "Bug4882Rename"  {
     // CHECK-NEXT: %bar2_x = firrtl.wire sym @x {annotations = [{class = "test0"}]}
   }
 }
+
+// -----
+
+// Issue #4920, the recursive inlining should consider the correct retop for NLAs.
+
+firrtl.circuit "DidNotContainSymbol" {
+  hw.hierpath private @path [@Bar1::@w, @Bar3]
+  firrtl.module private @Bar2() attributes {annotations = [{class = "firrtl.passes.InlineAnnotation"}]} {
+    firrtl.instance no sym @no @Bar1()
+  }
+  firrtl.module private @Bar1() attributes {annotations = [{class = "firrtl.passes.InlineAnnotation"}]} {
+    firrtl.instance bar3 sym @w @Bar3()
+  }
+  firrtl.module private @Bar3() attributes {annotations = [{class = "firrtl.passes.InlineAnnotation"}]} {
+    %w = firrtl.wire sym @w {annotations = [{circt.nonlocal = @path , class = "test0"}]} : !firrtl.uint<8>
+  }
+  firrtl.module @DidNotContainSymbol() {
+    firrtl.instance bar2 sym @w @Bar2()
+  }
+  // CHECK-LABEL: firrtl.module @DidNotContainSymbol() {
+  // CHECK-NEXT:     %bar2_no_bar3_w = firrtl.wire sym @w_0 {annotations = [{class = "test0"}]} : !firrtl.uint<8>
+  // CHECK-NEXT:  }
+}
+
+// -----
+
+// Issue #4915, the NLAs should be updated with renamed extern module instance.
+
+firrtl.circuit "SimTop" {
+  hw.hierpath private @nla_61 [@Rob::@difftest_3, @DifftestLoadEvent]
+  // CHECK: hw.hierpath private @nla_61 [@SimTop::@difftest_3_0, @DifftestLoadEvent]
+  hw.hierpath private @nla_60 [@Rob::@difftest_2, @DifftestLoadEvent]
+  // CHECK: hw.hierpath private @nla_60 [@SimTop::@difftest_2, @DifftestLoadEvent]
+  firrtl.extmodule private @DifftestIntWriteback()
+  firrtl.extmodule private @DifftestLoadEvent() attributes {annotations = [{circt.nonlocal = @nla_60, class = "B"}, {circt.nonlocal = @nla_61, class = "B"}]}
+	// CHECK: firrtl.extmodule private @DifftestLoadEvent() attributes {annotations = [{circt.nonlocal = @nla_60, class = "B"}, {circt.nonlocal = @nla_61, class = "B"}]}
+  firrtl.module private @Rob() attributes {annotations = [{class = "firrtl.passes.InlineAnnotation"}]} {
+    firrtl.instance difftest_2 sym @difftest_2 @DifftestLoadEvent()
+    firrtl.instance difftest_3 sym @difftest_3 @DifftestLoadEvent()
+  }
+  firrtl.module private @CtrlBlock() attributes {annotations = [{class = "firrtl.passes.InlineAnnotation"}]} {
+    firrtl.instance rob @Rob()
+  }
+  firrtl.module @SimTop() {
+    firrtl.instance difftest_3 sym @difftest_3 @DifftestIntWriteback()
+    firrtl.instance ctrlBlock @CtrlBlock()
+    // CHECK:  firrtl.instance difftest_3 sym @difftest_3 @DifftestIntWriteback()
+    // CHECK:  firrtl.instance ctrlBlock_rob_difftest_2 sym @difftest_2 @DifftestLoadEvent()
+    // CHECK:  firrtl.instance ctrlBlock_rob_difftest_3 sym @difftest_3_0 @DifftestLoadEvent()
+  }
+}