[FIRRTL] Replace remove-unused-ports with port-only IMDCE mode

include/circt/Dialect/FIRRTL/Passes.td

@@ -124,22 +124,6 @@ def RegisterOptimizer : Pass<"firrtl-register-optimizer", "firrtl::FModuleOp"> {
   }];
 }
 
-def RemoveUnusedPorts : Pass<"firrtl-remove-unused-ports", "firrtl::CircuitOp"> {
-  let summary = "Remove unused ports";
-  let description = [{
-    This pass removes unused ports without annotations or symbols. Implementation
-    wise, this pass iterates over the instance graph in a topological order from
-    leaves to the top so that we can remove unused ports optimally.
-  }];
-  let options = [
-    Option<"ignoreDontTouch", "ignore-dont-touch", "bool", "false",
-           "remove unused ports even if they have a symbol or annotation">
-  ];
-  let statistics = [
-    Statistic<"numRemovedPorts", "num-removed-ports", "Number of ports erased">,
-  ];
-}
-
 def IMDeadCodeElim : Pass<"firrtl-imdeadcodeelim", "mlir::ModuleOp"> {
   let summary = "Intermodule dead code elimination";
   let description = [{
@@ -148,6 +132,10 @@ def IMDeadCodeElim : Pass<"firrtl-imdeadcodeelim", "mlir::ModuleOp"> {
     of public modules or a value with a symbol. We first populate alive values
     into a set, and then propagate the liveness by looking at their dataflow.
   }];
+  let options = [
+    Option<"removePortsOnly", "remove-ports-only", "bool", "false",
+           "only remove dead module ports and preserve other operations">
+  ];
   let statistics = [
     Statistic<"numErasedOps", "num-erased-ops", "Number of operations erased">,
     Statistic<"numErasedModules", "num-erased-modules", "Number of modules erased">,

lib/Dialect/FIRRTL/FIRRTLReductions.cpp

@@ -2663,6 +2663,16 @@ struct ExtmoduleConventionRemover : public OpReduction<FExtModuleOp> {
   bool isOneShot() const override { return true; }
 };
 
+struct IMDCEPortReduction : public PassReduction {
+  IMDCEPortReduction(MLIRContext *context)
+      : PassReduction(
+            context, firrtl::createIMDeadCodeElim({/*removePortsOnly=*/true})) {
+  }
+  std::string getName() const override {
+    return "firrtl-imdeadcodeelim-remove-ports";
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // Reduction Registration
 //===----------------------------------------------------------------------===//
@@ -2702,9 +2712,7 @@ void firrtl::FIRRTLReducePatternDialectInterface::populateReducePatterns(
                                   true, true);
   patterns.add<PassReduction, 19>(getContext(), firrtl::createInliner());
   patterns.add<PassReduction, 18>(getContext(), firrtl::createIMConstProp());
-  patterns.add<PassReduction, 17>(
-      getContext(),
-      firrtl::createRemoveUnusedPorts({/*ignoreDontTouch=*/true}));
+  patterns.add<IMDCEPortReduction, 17>(getContext());
   patterns.add<NodeSymbolRemover, 16>();
   patterns.add<PassReduction, 15>(getContext(), firrtl::createIMDeadCodeElim());
   patterns.add<ConnectForwarder, 14>();

lib/Dialect/FIRRTL/Transforms/CMakeLists.txt

@@ -56,7 +56,6 @@ add_circt_dialect_library(CIRCTFIRRTLTransforms
   ProbesToSignals.cpp
   RandomizeRegisterInit.cpp
   RegisterOptimizer.cpp
-  RemoveUnusedPorts.cpp
   ResolvePaths.cpp
   ResolveTraces.cpp
   SFCCompat.cpp

lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp

@@ -56,6 +56,8 @@ static bool isDeletableDeclaration(Operation *op) {
 namespace {
 struct IMDeadCodeElimPass
     : public circt::firrtl::impl::IMDeadCodeElimBase<IMDeadCodeElimPass> {
+  using Base::Base;
+
   void runOnOperation() override;
 
   void rewriteModuleSignature(FModuleOp module);
@@ -200,6 +202,12 @@ void IMDeadCodeElimPass::markUnknownSideEffectOp(Operation *op) {
   for (auto operand : op->getOperands())
     markAlive(operand);
   markBlockUndeletable(op);
+
+  // Recursively mark any blocks contained within these operations as
+  // executable.
+  for (auto &region : op->getRegions())
+    for (auto &block : region.getBlocks())
+      markBlockExecutable(&block);
 }
 
 void IMDeadCodeElimPass::visitUser(Operation *op) {
@@ -266,7 +274,7 @@ void IMDeadCodeElimPass::markBlockExecutable(Block *block) {
     return; // Already executable.
 
   auto fmodule = dyn_cast<FModuleOp>(block->getParentOp());
-  if (fmodule && fmodule.isPublic())
+  if (fmodule && (fmodule.isPublic() || removePortsOnly))
     markAlive(fmodule);
 
   // Mark ports with don't touch as alive.
@@ -278,29 +286,36 @@ void IMDeadCodeElimPass::markBlockExecutable(Block *block) {
     }
 
   for (auto &op : *block) {
-    if (isDeclaration(&op))
-      markDeclaration(&op);
-    else if (auto instance = dyn_cast<FInstanceLike>(op))
+    if (auto instance = dyn_cast<FInstanceLike>(op)) {
       markFInstanceLikeOp(instance);
-    else if (isa<FConnectLike>(op))
-      // Skip connect op.
       continue;
-    else if (hasUnknownSideEffect(&op)) {
+    }
+
+    // Skip connects in both modes.
+    if (isa<FConnectLike>(op))
+      continue;
+
+    if (removePortsOnly) {
+      // In port-only mode, all non-instance, non-connect ops are alive.
       markUnknownSideEffectOp(&op);
-      // Recursively mark any blocks contained within these operations as
-      // executable.
-      for (auto &region : op.getRegions())
-        for (auto &block : region.getBlocks())
-          markBlockExecutable(&block);
+      continue;
     }
 
+    // Full IMDCE mode: handle declarations and side effects.
+    if (isDeclaration(&op))
+      markDeclaration(&op);
+    else if (hasUnknownSideEffect(&op))
+      markUnknownSideEffectOp(&op);
+
     // TODO: Handle attach etc.
   }
 }
 
 void IMDeadCodeElimPass::forwardConstantOutputPort(FModuleOp module) {
-  // This tracks constant values of output ports.
-  SmallVector<std::pair<unsigned, APSInt>> constantPortIndicesAndValues;
+  // This tracks constant values of output ports. std::nullopt represents an
+  // invalid value.
+  SmallVector<std::pair<unsigned, std::optional<APSInt>>>
+      constantPortIndicesAndValues;
   auto ports = module.getPorts();
   auto *instanceGraphNode = instanceGraph->lookup(module);
 
@@ -314,12 +329,15 @@ void IMDeadCodeElimPass::forwardConstantOutputPort(FModuleOp module) {
       continue;
 
     // Remember the index and constant value connected to an output port.
-    if (auto connect = getSingleConnectUserOf(arg))
+    if (auto connect = getSingleConnectUserOf(arg)) {
       if (auto constant = connect.getSrc().getDefiningOp<ConstantOp>())
         constantPortIndicesAndValues.push_back({index, constant.getValue()});
+      else if (connect.getSrc().getDefiningOp<InvalidValueOp>())
+        constantPortIndicesAndValues.push_back({index, std::nullopt});
+    }
   }
 
-  // If there is no constant port, abort.
+  // If there is no constant or invalid port, abort.
   if (constantPortIndicesAndValues.empty())
     return;
 
@@ -334,8 +352,13 @@ void IMDeadCodeElimPass::forwardConstantOutputPort(FModuleOp module) {
       auto result = instance.getResult(index);
       assert(ports[index].isOutput() && "must be an output port");
 
-      // Replace the port with the constant.
-      result.replaceAllUsesWith(ConstantOp::create(builder, constant));
+      // Replace the port with the constant or invalid value.
+      Value replacement;
+      if (constant)
+        replacement = ConstantOp::create(builder, *constant);
+      else
+        replacement = InvalidValueOp::create(builder, result.getType());
+      result.replaceAllUsesWith(replacement);
     }
   }
 }
@@ -435,11 +458,13 @@ void IMDeadCodeElimPass::runOnOperation() {
     forwardConstantOutputPort(module);
 
   for (auto module : circuit.getBodyBlock()->getOps<FModuleOp>()) {
-    // Mark the ports of public modules as alive.
-    if (module.isPublic()) {
+    bool isPublic = module.isPublic();
+    if (isPublic || removePortsOnly) {
       markBlockExecutable(module.getBodyBlock());
-      for (auto port : module.getBodyBlock()->getArguments())
-        markAlive(port);
+      // Mark the ports of public modules as alive.
+      if (isPublic)
+        for (auto port : module.getBodyBlock()->getArguments())
+          markAlive(port);
     }
 
     // Walk annotations and populate a map from hierpath to attached annotation
@@ -502,8 +527,9 @@ void IMDeadCodeElimPass::runOnOperation() {
     if (!liveElements.count(op))
       op.erase();
 
-  for (auto module : modules)
-    eraseEmptyModule(module);
+  if (!removePortsOnly)
+    for (auto module : modules)
+      eraseEmptyModule(module);
 
   // Clean up data structures.
   executableBlocks.clear();