From bcc1eeab9363049396419b4f23124dbf38c23772 Mon Sep 17 00:00:00 2001
From: Prithayan Barua <prithayan@gmail.com>
Date: Mon, 11 Dec 2023 16:59:59 -0800
Subject: [PATCH] [ModuleInliner] Donot retop if flattening the module

---
 .../FIRRTL/Transforms/ModuleInliner.cpp       |  2 +-
 test/Dialect/FIRRTL/inliner.mlir              | 29 +++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/lib/Dialect/FIRRTL/Transforms/ModuleInliner.cpp b/lib/Dialect/FIRRTL/Transforms/ModuleInliner.cpp
index e7c760ae4ae5..bb16ebaa5246 100644
--- a/lib/Dialect/FIRRTL/Transforms/ModuleInliner.cpp
+++ b/lib/Dialect/FIRRTL/Transforms/ModuleInliner.cpp
@@ -1150,7 +1150,7 @@ void Inliner::inlineInstances(FModuleOp module) {
     // participate in any HierPathOp. But the reTop might add a symbol to it, if
     // a HierPathOp is added to this Op.
     DenseMap<Attribute, Attribute> symbolRenames;
-    if (!rootMap[target.getNameAttr()].empty()) {
+    if (!rootMap[target.getNameAttr()].empty() && !shouldFlatten(target)) {
       for (auto sym : rootMap[target.getNameAttr()]) {
         auto &mnla = nlaMap[sym];
         sym = mnla.reTop(module);
diff --git a/test/Dialect/FIRRTL/inliner.mlir b/test/Dialect/FIRRTL/inliner.mlir
index 5abaf45edabf..a84fc438ddc1 100644
--- a/test/Dialect/FIRRTL/inliner.mlir
+++ b/test/Dialect/FIRRTL/inliner.mlir
@@ -1417,3 +1417,32 @@ firrtl.circuit "PropertyUTurn" {
   }
   firrtl.extmodule @Consume(in in : !firrtl.string)
 }
+
+// -----
+
+// Test that inlining and flattening compose with nla well.
+firrtl.circuit "compose_nla" {
+  hw.hierpath private @nla1 [@test1::@sym, @test2::@sym, @test3]
+// CHECK-NOT:  hw.hierpath private @nla1
+firrtl.module @compose_nla() {
+// CHECK-LABEL: firrtl.module @compose_nla() {
+  firrtl.instance test1 @test1()
+  firrtl.instance test2 @test2()
+  firrtl.instance test3 @test3()
+}
+firrtl.module private @test1() attributes {annotations =
+        [{class = "firrtl.transforms.FlattenAnnotation"},
+         {class = "firrtl.passes.InlineAnnotation"}]} {
+  %test_wire = firrtl.wire : !firrtl.uint<2>
+  firrtl.instance test2 sym @sym @test2()
+  firrtl.instance test3 @test3()
+}
+firrtl.module private @test2() attributes {annotations =
+        [{class = "firrtl.passes.InlineAnnotation"}]} {
+  %test_wire = firrtl.wire : !firrtl.uint<2>
+  firrtl.instance test3 sym @sym @test3()
+}
+firrtl.module private @test3() {
+  %test_wire = firrtl.wire : !firrtl.uint<2>
+}
+}