From 34ba163c00dadf4ef3c71ba4a209892cb938e5d8 Mon Sep 17 00:00:00 2001
From: Schuyler Eldridge <schuyler.eldridge@sifive.com>
Date: Fri, 22 Nov 2024 22:12:00 -0500
Subject: [PATCH] [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 <schuyler.eldridge@sifive.com>
---
 .../FIRRTL/Transforms/IMDeadCodeElim.cpp      | 70 +++++++++++--------
 test/Dialect/FIRRTL/imdce.mlir                | 25 +++++++
 2 files changed, 65 insertions(+), 30 deletions(-)

diff --git a/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp b/lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp
index 478ca2ddf9ab..ecdd0f4173d2 100644
--- 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
index ad8fe3f82b13..5f7594cb26af 100644
--- 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()
+      }
+    }
+  }
+}