From 3a241f73a022c4c81f59f6aaa6735d4c32a3181a Mon Sep 17 00:00:00 2001
From: Will Dietz <will.dietz@sifive.com>
Date: Fri, 8 Sep 2023 17:42:54 -0500
Subject: [PATCH] [FIRRTL][IMDCE] Don't try to remove main module, don't warn
 if empty. (#6065)

---
 lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp |  3 ++-
 test/Dialect/FIRRTL/imdce.mlir                   | 16 ++++++++++++++--
 test/firtool/firtool-errors.mlir                 |  1 -
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp b/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp
index 691b80a7837d..4fe159c20c29 100644
--- a/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp
+++ b/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp
@@ -451,7 +451,8 @@ void IMDeadCodeElimPass::runOnOperation() {
       op.erase();
 
   for (auto module : modules)
-    eraseEmptyModule(module);
+    if (module != circuit.getMainModule())
+      eraseEmptyModule(module);
 
   // Clean up data structures.
   executableBlocks.clear();
diff --git a/test/Dialect/FIRRTL/imdce.mlir b/test/Dialect/FIRRTL/imdce.mlir
index b5ea6002816b..7602d9de06d1 100644
--- a/test/Dialect/FIRRTL/imdce.mlir
+++ b/test/Dialect/FIRRTL/imdce.mlir
@@ -93,7 +93,6 @@ firrtl.circuit "top"  {
 
   // CHECK-LABEL: firrtl.module @top(in %clock: !firrtl.clock, in %input: !firrtl.uint<1>) {
   // CHECK-NEXT:  }
-  // expected-warning @+1 {{module `top` is empty but cannot be removed because the module is public}}
   firrtl.module @top(in %clock: !firrtl.clock, in %input: !firrtl.uint<1>) {
     %tile_input, %tile_output = firrtl.instance tile  @Child1(in input: !firrtl.uint<1>, out output: !firrtl.uint<1>)
     firrtl.strictconnect %tile_input, %input : !firrtl.uint<1>
@@ -251,7 +250,6 @@ firrtl.circuit "RefPorts" {
 
 firrtl.circuit "MemoryInDeadCycle" {
   // CHECK-LABEL: firrtl.module public @MemoryInDeadCycle
-  // expected-warning @+1{{module `MemoryInDeadCycle` is empty but cannot be removed because the module is public}}
   firrtl.module public @MemoryInDeadCycle(in %clock: !firrtl.clock, in %addr: !firrtl.uint<4>) {
 
     // CHECK-NOT: firrtl.mem
@@ -477,3 +475,17 @@ firrtl.circuit "AnnoAlive" {
      firrtl.wire {annotations = [{class = "circt.test"}]} : !firrtl.uint
   }
 }
+
+// -----
+// Test warning about not being able to remove dead public modules.
+
+// CHECK-LABEL: "DeadPublic"
+firrtl.circuit "DeadPublic" {
+  // CHECK: module @PublicDeadChild
+  // expected-warning @below {{module `PublicDeadChild` is empty but cannot be removed because the module is public}}
+  firrtl.module @PublicDeadChild() {}
+  // CHECK: module @DeadPublic
+  firrtl.module @DeadPublic() {
+     firrtl.instance pdc @PublicDeadChild()
+  }
+}
diff --git a/test/firtool/firtool-errors.mlir b/test/firtool/firtool-errors.mlir
index 6840fbb7518c..e4682ead413e 100644
--- a/test/firtool/firtool-errors.mlir
+++ b/test/firtool/firtool-errors.mlir
@@ -1,7 +1,6 @@
 // RUN: firtool %s -format=mlir -verilog -verify-diagnostics | FileCheck --allow-empty %s
 
 firrtl.circuit "top" {
-  // expected-warning @+3 {{module `top` is empty but cannot be removed because the module is public}}
   // expected-error @+2 {{'firrtl.module' op contains a 'chisel3.util.experimental.ForceNameAnnotation' that is not a non-local annotation}}
   // expected-note @+1 {{the erroneous annotation is '{class = "chisel3.util.experimental.ForceNameAnnotation"}'}}
   firrtl.module @top() attributes {annotations = [{class = "chisel3.util.experimental.ForceNameAnnotation"}]} {