[FIRRTL] Make IMDCE work for ops w/ regions/blocks

Fix a bug in FIRRTL's IMDCE Pass where it would not visit the blocks of
operations.  This can result in crashes for blocks which contain users of
FIRRTL modules, e.g., instances inside layerblocks of sv.ifdef.

This conceptually is two changes: (1) when marking a block live, the block
needs to be recursively walked and (2) when erasing ops, a recursive walk
is needed.

Signed-off-by: Schuyler Eldridge <[email protected]>

lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp

@@ -14,6 +14,7 @@
 #include "circt/Dialect/HW/InnerSymbolTable.h"
 #include "circt/Support/Debug.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/IR/Iterators.h"
 #include "mlir/IR/Threading.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Pass/Pass.h"
@@ -240,15 +241,16 @@ void IMDeadCodeElimPass::markBlockExecutable(Block *block) {
   if (!executableBlocks.insert(block).second)
     return; // Already executable.
 
-  auto fmodule = cast<FModuleOp>(block->getParentOp());
-  if (fmodule.isPublic())
+  auto fmodule = dyn_cast<FModuleOp>(block->getParentOp());
+  if (fmodule && fmodule.isPublic())
     markAlive(fmodule);
 
   // Mark ports with don't touch as alive.
   for (auto blockArg : block->getArguments())
     if (hasDontTouch(blockArg)) {
       markAlive(blockArg);
-      markAlive(fmodule);
+      if (fmodule)
+        markAlive(fmodule);
     }
 
   for (auto &op : *block) {
@@ -261,8 +263,14 @@ void IMDeadCodeElimPass::markBlockExecutable(Block *block) {
     else if (isa<FConnectLike>(op))
       // Skip connect op.
       continue;
-    else if (hasUnknownSideEffect(&op))
+    else if (hasUnknownSideEffect(&op)) {
       markUnknownSideEffectOp(&op);
+      // Recursively mark any blocks contained within these operations as
+      // executable.
+      for (auto &region : op.getRegions())
+        for (auto &block : region.getBlocks())
+          markBlockExecutable(&block);
+    }
 
     // TODO: Handle attach etc.
   }
@@ -561,33 +569,35 @@ void IMDeadCodeElimPass::rewriteModuleBody(FModuleOp module) {
       std::bind(removeDeadNonLocalAnnotations, -1, std::placeholders::_1));
 
   // Walk the IR bottom-up when deleting operations.
-  for (auto &op : llvm::make_early_inc_range(llvm::reverse(*body))) {
-    // Connects to values that we found to be dead can be dropped.
-    if (auto connect = dyn_cast<FConnectLike>(op)) {
-      if (isAssumedDead(connect.getDest())) {
-        LLVM_DEBUG(llvm::dbgs() << "DEAD: " << connect << "\n";);
-        connect.erase();
-        ++numErasedOps;
-      }
-      continue;
-    }
-
-    // Delete dead wires, regs, nodes and alloc/read ops.
-    if ((isDeclaration(&op) || !hasUnknownSideEffect(&op)) &&
-        isAssumedDead(&op)) {
-      LLVM_DEBUG(llvm::dbgs() << "DEAD: " << op << "\n";);
-      assert(op.use_empty() && "users should be already removed");
-      op.erase();
-      ++numErasedOps;
-      continue;
-    }
-
-    // Remove non-sideeffect op using `isOpTriviallyDead`.
-    if (mlir::isOpTriviallyDead(&op)) {
-      op.erase();
-      ++numErasedOps;
-    }
-  }
+  module.walk<mlir::WalkOrder::PostOrder, mlir::ReverseIterator>(
+      [&](Operation *op) {
+        // Connects to values that we found to be dead can be dropped.
+        LLVM_DEBUG(llvm::dbgs() << "Visit: " << *op << "\n");
+        if (auto connect = dyn_cast<FConnectLike>(op)) {
+          if (isAssumedDead(connect.getDest())) {
+            LLVM_DEBUG(llvm::dbgs() << "DEAD: " << connect << "\n";);
+            connect.erase();
+            ++numErasedOps;
+          }
+          return;
+        }
+
+        // Delete dead wires, regs, nodes and alloc/read ops.
+        if ((isDeclaration(op) || !hasUnknownSideEffect(op)) &&
+            isAssumedDead(op)) {
+          LLVM_DEBUG(llvm::dbgs() << "DEAD: " << *op << "\n";);
+          assert(op->use_empty() && "users should be already removed");
+          op->erase();
+          ++numErasedOps;
+          return;
+        }
+
+        // Remove non-sideeffect op using `isOpTriviallyDead`.
+        if (mlir::isOpTriviallyDead(op)) {
+          op->erase();
+          ++numErasedOps;
+        }
+      });
 }
 
 void IMDeadCodeElimPass::rewriteModuleSignature(FModuleOp module) {

test/Dialect/FIRRTL/imdce.mlir

@@ -617,3 +617,58 @@ firrtl.circuit "OMIRRemoval" {
       ]} : !firrtl.uint<4>
   }
 }
+
+// -----
+
+// Test that an operation with a nested block user will be removed (and not
+// crash).  This should work for both FIRRTL operations and non-FIRRTL
+// operations.
+//
+// CHECK-LAEBL: "Foo"
+firrtl.circuit "Foo" {
+  firrtl.layer @A bind {}
+  sv.macro.decl @B["B"]
+  // CHECK-NOT: @Bar
+  firrtl.module private @Bar() {}
+  firrtl.module @Foo() {
+    // CHECK-LABEL: firrtl.layerblock @A
+    firrtl.layerblock @A {
+      // CHECK-NOT: firrtl.instance
+      firrtl.instance bar @Bar()
+      // CHECK-LABEL: sv.ifdef @B
+      sv.ifdef @B {
+        // CHECK-NOT: firrtl.instance
+        firrtl.instance bar2 @Bar()
+      }
+    }
+  }
+}
+
+// -----
+
+// CHECK-LABEL: "Foo"
+firrtl.circuit "Foo" {
+  firrtl.layer @A bind {}
+  // CHECK-LABEL: @Bar
+  // CHECK-NOT:     out %probe
+  firrtl.module private @Bar(
+    in %a: !firrtl.uint<1>,
+    out %b: !firrtl.uint<1>,
+    out %probe: !firrtl.probe<uint<1>, @A>
+  ) {
+    // CHECK-NOT: firrtl.layerblock
+    firrtl.layerblock @A {
+      %1 = firrtl.not %a : (!firrtl.uint<1>) -> !firrtl.uint<1>
+      %a_not = firrtl.node %1 : !firrtl.uint<1>
+      %2 = firrtl.ref.send %a_not : !firrtl.uint<1>
+      %3 = firrtl.ref.cast %2 : (!firrtl.probe<uint<1>>) -> !firrtl.probe<uint<1>, @A>
+      firrtl.ref.define %probe, %3 : !firrtl.probe<uint<1>, @A>
+    }
+    firrtl.matchingconnect %b, %a : !firrtl.uint<1>
+  }
+  firrtl.module @Foo(in %a: !firrtl.uint<1>, out %b: !firrtl.uint<1>) {
+    %bar_a, %bar_b, %bar_probe = firrtl.instance bar @Bar(in a: !firrtl.uint<1>, out b: !firrtl.uint<1>, out probe: !firrtl.probe<uint<1>, @A>)
+    firrtl.matchingconnect %bar_a, %a : !firrtl.uint<1>
+    firrtl.matchingconnect %b, %bar_b : !firrtl.uint<1>
+  }
+}