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

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>();

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>
   }
 }