From f05c89b32ec52066cbb24f231eb797c4c25df205 Mon Sep 17 00:00:00 2001
From: Fabian Schuiki <fabian@schuiki.ch>
Date: Mon, 28 Aug 2023 11:52:50 -0700
Subject: [PATCH] [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.
---
 .../FIRRTL/Transforms/ExtractInstances.cpp     | 15 +++++++--------
 test/Dialect/FIRRTL/extract-instances.mlir     | 18 ++++++++++++++++++
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/lib/Dialect/FIRRTL/Transforms/ExtractInstances.cpp b/lib/Dialect/FIRRTL/Transforms/ExtractInstances.cpp
index 543ec6d6df3d..939dacc216ce 100644
--- a/lib/Dialect/FIRRTL/Transforms/ExtractInstances.cpp
+++ b/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.
diff --git a/test/Dialect/FIRRTL/extract-instances.mlir b/test/Dialect/FIRRTL/extract-instances.mlir
index 56f776a9c6df..85d2d71e125f 100644
--- a/test/Dialect/FIRRTL/extract-instances.mlir
+++ b/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()
+  }
+}