From 8a3165bcb907247afc7a66f97243a14a1130fe42 Mon Sep 17 00:00:00 2001
From: Fabian Schuiki <fabian@schuiki.ch>
Date: Tue, 28 Jan 2025 18:39:49 -0800
Subject: [PATCH] [FIRRTL][IMDCE] Support operations with blocks

Avoid a crash in IMDCE when encountering a block argument that is not
defined by a module, but some other operation that has nested blocks,
such as a contract. Additionally, when handling liveness of generic
operations, conservatively mark any nested blocks executable.
---
 .../FIRRTL/Transforms/IMDeadCodeElim.cpp      | 33 +++++++++++--------
 test/Dialect/FIRRTL/imdce.mlir                |  9 ++---
 2 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp b/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp
index ceae7de7627c..0b27c499b1ae 100644
--- a/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp
+++ b/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp
@@ -492,19 +492,21 @@ void IMDeadCodeElimPass::visitValue(Value value) {
 
   // Requiring an input port propagates the liveness to each instance.
   if (auto blockArg = dyn_cast<BlockArgument>(value)) {
-    auto module = cast<FModuleOp>(blockArg.getParentBlock()->getParentOp());
-    auto portDirection = module.getPortDirection(blockArg.getArgNumber());
-    // If the port is input, it's necessary to mark corresponding input ports of
-    // instances as alive. We don't have to propagate the liveness of output
-    // ports.
-    if (portDirection == Direction::In) {
-      for (auto *instRec : instanceGraph->lookup(module)->uses()) {
-        auto instance = cast<InstanceOp>(instRec->getInstance());
-        if (liveElements.contains(instance))
-          markAlive(instance.getResult(blockArg.getArgNumber()));
+    if (auto module =
+            dyn_cast<FModuleOp>(blockArg.getParentBlock()->getParentOp())) {
+      auto portDirection = module.getPortDirection(blockArg.getArgNumber());
+      // If the port is input, it's necessary to mark corresponding input ports
+      // of instances as alive. We don't have to propagate the liveness of
+      // output ports.
+      if (portDirection == Direction::In) {
+        for (auto *instRec : instanceGraph->lookup(module)->uses()) {
+          auto instance = cast<InstanceOp>(instRec->getInstance());
+          if (liveElements.contains(instance))
+            markAlive(instance.getResult(blockArg.getArgNumber()));
+        }
       }
+      return;
     }
-    return;
   }
 
   // Marking an instance port as alive propagates to the corresponding port of
@@ -533,10 +535,15 @@ void IMDeadCodeElimPass::visitValue(Value value) {
     return;
   }
 
-  // If op is defined by an operation, mark its operands as alive.
-  if (auto op = value.getDefiningOp())
+  // If the value is defined by an operation, mark its operands alive and any
+  // nested blocks executable.
+  if (auto op = value.getDefiningOp()) {
     for (auto operand : op->getOperands())
       markAlive(operand);
+    for (auto &region : op->getRegions())
+      for (auto &block : region)
+        markBlockExecutable(&block);
+  }
 
   // If either result of a forceable declaration is alive, they both are.
   if (auto fop = value.getDefiningOp<Forceable>();
diff --git a/test/Dialect/FIRRTL/imdce.mlir b/test/Dialect/FIRRTL/imdce.mlir
index 02811d5efb3c..e9d6abd2ef83 100644
--- a/test/Dialect/FIRRTL/imdce.mlir
+++ b/test/Dialect/FIRRTL/imdce.mlir
@@ -613,10 +613,11 @@ firrtl.circuit "FormalMarkerIsUse" {
 // crash the `firrtl::hasDontTouch` query.
 // CHECK-LABEL: firrtl.circuit "NonModuleBlockArgsMustNotCrash"
 firrtl.circuit "NonModuleBlockArgsMustNotCrash" {
-  firrtl.module @NonModuleBlockArgsMustNotCrash(in %a: !firrtl.uint<42>, out %b: !firrtl.uint<42>) {
-    %0 = firrtl.contract %a : !firrtl.uint<42> {
-    ^bb0(%arg0: !firrtl.uint<42>):
+  firrtl.module @NonModuleBlockArgsMustNotCrash(in %a: !firrtl.uint<1>, out %b: !firrtl.uint<1>) {
+    %0 = firrtl.contract %a : !firrtl.uint<1> {
+    ^bb0(%arg0: !firrtl.uint<1>):
+      firrtl.int.verif.assert %arg0 : !firrtl.uint<1>
     }
-    firrtl.matchingconnect %b, %0 : !firrtl.uint<42>
+    firrtl.matchingconnect %b, %0 : !firrtl.uint<1>
   }
 }