[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.

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;

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 } : () -> ()
+}