[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]>
llvm · Nov 23, 2024 · 34ba163 · 34ba163
1 parent 9f22da8
commit 34ba163
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 30 deletions.
diff --git a/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp b/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) {
@@ -264,6 +266,12 @@ void IMDeadCodeElimPass::markBlockExecutable(Block *block) {
     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) {

diff --git a/test/Dialect/FIRRTL/imdce.mlir b/test/Dialect/FIRRTL/imdce.mlir
@@ -617,3 +617,28 @@ firrtl.circuit "OMIRRemoval" {
       ]} : !firrtl.uint<4>
   }
 }
+
+// -----
+
+// Test that an operation with a nested block user will be removed (and not
+// crash).
+//
+// 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()
+      }
+    }
+  }
+}