From e4a2d66a4c2484c18e90c0ce1113f9b20306ef53 Mon Sep 17 00:00:00 2001
From: Fabian Schuiki <fabian@schuiki.ch>
Date: Wed, 22 Jan 2025 20:12:28 -0800
Subject: [PATCH] [FIRRTL][IMDCE] Don't remove modules with symbol uses (#8114)

Mark any modules alive that appear as symbol references outside of
hierarchical path and instance operations. Prevent modules from being
deleted if they are explicitly marked as alive. This fixes an issue in
IMDCE where references to modules in `firrtl.formal` or other unknown
operations would become invalidated due to the pass deleting the modules
despite these references.

The pattern of walking operations' attribute dictionary an looking for
symbol references is inspired by what MLIR's `SymbolTable` does
internally in `walkSymbolRefs` when looking for symbol uses.
---
 .../FIRRTL/Transforms/IMDeadCodeElim.cpp      | 53 +++++++++++++------
 test/Dialect/FIRRTL/imdce.mlir                | 25 ++++++++-
 2 files changed, 62 insertions(+), 16 deletions(-)

diff --git a/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp b/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp
index 13f8c8a2363e..ceae7de7627c 100644
--- a/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp
+++ b/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp
@@ -355,21 +355,40 @@ void IMDeadCodeElimPass::runOnOperation() {
       return;
     }
 
-    // If there is an unknown use of inner sym or hierpath, just mark all of
-    // them alive.
-    for (NamedAttribute namedAttr : op->getAttrs()) {
-      namedAttr.getValue().walk([&](Attribute subAttr) {
-        if (auto innerRef = dyn_cast<hw::InnerRefAttr>(subAttr))
-          if (auto instance = dyn_cast_or_null<firrtl::InstanceOp>(
-                  innerRefNamespace->lookupOp(innerRef)))
-            markAlive(instance);
-
-        if (auto flatSymbolRefAttr = dyn_cast<FlatSymbolRefAttr>(subAttr))
-          if (auto hierPath = symbolTable->template lookup<hw::HierPathOp>(
-                  flatSymbolRefAttr.getAttr()))
-            markAlive(hierPath);
-      });
-    }
+    // If there is an unknown symbol or inner symbol use, mark all of them
+    // alive.
+    op->getAttrDictionary().walk([&](Attribute attr) {
+      if (auto innerRef = dyn_cast<hw::InnerRefAttr>(attr)) {
+        // Mark instances alive that are targeted by an inner ref.
+        if (auto instance = dyn_cast_or_null<firrtl::InstanceOp>(
+                innerRefNamespace->lookupOp(innerRef)))
+          markAlive(instance);
+        return;
+      }
+
+      if (auto symbolRef = dyn_cast<FlatSymbolRefAttr>(attr)) {
+        auto *symbol = symbolTable->lookup(symbolRef.getAttr());
+        if (!symbol)
+          return;
+
+        // Mark referenced hierarchical paths alive.
+        if (auto hierPath = dyn_cast<hw::HierPathOp>(symbol))
+          markAlive(hierPath);
+
+        // Mark modules referenced by unknown ops alive.
+        if (auto module = dyn_cast<FModuleOp>(symbol)) {
+          if (!isa<firrtl::InstanceOp>(op)) {
+            LLVM_DEBUG(llvm::dbgs()
+                       << "Unknown use of " << module.getModuleNameAttr()
+                       << " in " << op->getName() << "\n");
+            markAlive(module);
+            markBlockExecutable(module.getBodyBlock());
+          }
+        }
+
+        return;
+      }
+    });
   });
 
   // Create a vector of modules in the post order of instance graph.
@@ -801,6 +820,10 @@ void IMDeadCodeElimPass::eraseEmptyModule(FModuleOp module) {
     return;
   }
 
+  // We cannot delete alive modules.
+  if (liveElements.contains(module))
+    return;
+
   instanceGraph->erase(instanceGraphNode);
   module.erase();
   ++numErasedModules;
diff --git a/test/Dialect/FIRRTL/imdce.mlir b/test/Dialect/FIRRTL/imdce.mlir
index 3fa028744318..c3dca150e2cf 100644
--- a/test/Dialect/FIRRTL/imdce.mlir
+++ b/test/Dialect/FIRRTL/imdce.mlir
@@ -1,4 +1,4 @@
-// RUN: circt-opt -pass-pipeline='builtin.module(firrtl-imdeadcodeelim)' --split-input-file -verify-diagnostics %s | FileCheck %s
+// RUN: circt-opt -pass-pipeline='builtin.module(firrtl-imdeadcodeelim)' --split-input-file --verify-diagnostics --allow-unregistered-dialect %s | FileCheck %s
 firrtl.circuit "top" {
   // In `dead_module`, %source is connected to %dest through several dead operations such as
   // node, wire, reg or rgereset. %dest is also dead at any instantiation, so check that
@@ -583,3 +583,26 @@ firrtl.circuit "Foo" {
     firrtl.matchingconnect %b, %bar_b : !firrtl.uint<1>
   }
 }
+
+// -----
+
+// Modules referenced from unknown ops cannot be removed, even if they are empty
+// and all their instances have been deleted.
+// CHECK-LABEL: firrtl.circuit "FormalMarkerIsUse"
+firrtl.circuit "FormalMarkerIsUse" {
+  // expected-warning @below {{module `FormalMarkerIsUse` is empty but cannot be removed because the module is public}}
+  firrtl.module @FormalMarkerIsUse() {
+    // CHECK-NOT: firrtl.instance foo @Foo
+    // CHECK-NOT: firrtl.instance bar @Bar
+    firrtl.instance foo @Foo()
+    firrtl.instance bar @Bar()
+  }
+  // CHECK: firrtl.module private @Foo
+  // CHECK: firrtl.module private @Bar
+  // CHECK: firrtl.module private @Uninstantiated
+  firrtl.module private @Foo() {}
+  firrtl.module private @Bar() {}
+  firrtl.module private @Uninstantiated() {}
+  firrtl.formal @Test, @Foo {}
+  "some_unknown_dialect.op"() { magic = @Bar, other = @Uninstantiated } : () -> ()
+}