[ExtractInstances] Fix module prefix not being applied to NLA (#5971)

When instances are grouped after extraction, the group module would
inherit the appropriate prefix from the DUT, but any hierarchical paths
pointing to it would not include the prefix. Fix this by using the same
name for the module and the hierarchical path.

Fixes #5961.

lib/Dialect/FIRRTL/Transforms/ExtractInstances.cpp

@@ -866,7 +866,8 @@ void ExtractInstancesPass::groupInstances() {
     OpBuilder builder(parentOp);
 
     // Uniquify the wrapper name.
-    wrapperName = circuitNamespace.newName(wrapperName);
+    auto wrapperModuleName = builder.getStringAttr(
+        circuitNamespace.newName(dutPrefix + wrapperName));
     auto wrapperInstName =
         builder.getStringAttr(getModuleNamespace(parent).newName(wrapperName));
 
@@ -918,13 +919,12 @@ void ExtractInstancesPass::groupInstances() {
 
         // The relevant part of the NLA is of the form `Top::bb`, which we want
         // to expand to `Top::wrapperInst` and `Wrapper::bb`.
-        auto wrapperNameAttr = builder.getStringAttr(wrapperName);
         auto ref1 = InnerRefAttr::get(parent.moduleNameAttr(), wrapperInstName);
         Attribute ref2;
         if (auto innerRef = nlaPath[nlaIdx].dyn_cast<InnerRefAttr>())
-          ref2 = InnerRefAttr::get(wrapperNameAttr, innerRef.getName());
+          ref2 = InnerRefAttr::get(wrapperModuleName, innerRef.getName());
         else
-          ref2 = FlatSymbolRefAttr::get(wrapperNameAttr);
+          ref2 = FlatSymbolRefAttr::get(wrapperModuleName);
         LLVM_DEBUG(llvm::dbgs() << "    - Expanding " << nlaPath[nlaIdx]
                                 << " to (" << ref1 << ", " << ref2 << ")\n");
         nlaPath[nlaIdx] = ref1;
@@ -935,14 +935,13 @@ void ExtractInstancesPass::groupInstances() {
         nla.namepathAttr(builder.getArrayAttr(nlaPath));
         LLVM_DEBUG(llvm::dbgs() << "    - Modified to " << nla << "\n");
         // Add the NLA to the wrapper module.
-        nlaTable.addNLAtoModule(nla, wrapperNameAttr);
+        nlaTable.addNLAtoModule(nla, wrapperModuleName);
       }
     }
 
     // Create the wrapper module.
-    auto wrapper = builder.create<FModuleOp>(
-        builder.getUnknownLoc(), builder.getStringAttr(dutPrefix + wrapperName),
-        ports);
+    auto wrapper = builder.create<FModuleOp>(builder.getUnknownLoc(),
+                                             wrapperModuleName, ports);
 
     // Instantiate the wrapper module in the parent and replace uses of the
     // extracted instances' ports with the corresponding wrapper module ports.

test/Dialect/FIRRTL/extract-instances.mlir

@@ -504,3 +504,21 @@ firrtl.circuit "InstSymConflict" {
   // CHECK-SAME: #hw.innerNameRef<@InstSymConflict::@bb>
   // CHECK-SAME: ]
 }
+
+// Module prefixing should not break extraction.
+// https://github.com/llvm/circt/issues/5961
+// CHECK-LABEL: firrtl.circuit "Plop_Foo"
+firrtl.circuit "Plop_Foo" attributes {annotations = [{class = "sifive.enterprise.firrtl.ExtractClockGatesFileAnnotation", filename = "ckgates.txt", group = "ClockGates"}]} {
+  // CHECK: firrtl.hierpath @nla_1 [@Plop_Foo::@ClockGates, @Plop_ClockGates::@ckg, @EICG_wrapper]
+  firrtl.hierpath @nla_1 [@Plop_Foo::@core, @Plop_Bar::@ckg, @EICG_wrapper]
+  firrtl.extmodule private @EICG_wrapper() attributes {defname = "EICG_wrapper"}
+  firrtl.module private @Plop_Bar() {
+    firrtl.instance ckg sym @ckg @EICG_wrapper()
+  }
+  // CHECK: firrtl.module @Plop_ClockGates()
+  // CHECK-LABEL: firrtl.module @Plop_Foo()
+  firrtl.module @Plop_Foo() attributes {annotations = [{class = "sifive.enterprise.firrtl.MarkDUTAnnotation", prefix = "Plop_"}]} {
+    // CHECK: firrtl.instance ClockGates sym @ClockGates @Plop_ClockGates()
+    firrtl.instance core sym @core @Plop_Bar()
+  }
+}