IMDeadCodeElim.cpp

Go to the documentation of this file.

//===- IMDeadCodeElim.cpp - Intermodule Dead Code Elimination ---*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
 
#include "circt/Dialect/FIRRTL/FIRRTLInstanceGraph.h"
#include "circt/Dialect/FIRRTL/FIRRTLOps.h"
#include "circt/Dialect/FIRRTL/Passes.h"
#include "circt/Dialect/HW/HWOps.h"
#include "circt/Dialect/HW/InnerSymbolTable.h"
#include "circt/Support/Debug.h"
#include "mlir/IR/Iterators.h"
#include "mlir/IR/Threading.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "llvm/ADT/BitVector.h"
#include "llvm/ADT/DenseMapInfoVariant.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/Support/Debug.h"
 
#define DEBUG_TYPE "firrtl-imdeadcodeelim"
 
namespace circt {
namespace firrtl {
#define GEN_PASS_DEF_IMDEADCODEELIM
#include "circt/Dialect/FIRRTL/Passes.h.inc"
} // namespace firrtl
} // namespace circt
 
using namespace circt;
using namespace firrtl;
 
// Return true if this op has side-effects except for alloc and read.
static bool hasUnknownSideEffect(Operation *op) {
  return !(mlir::isMemoryEffectFree(op) ||
           mlir::hasSingleEffect<mlir::MemoryEffects::Allocate>(op) ||
           mlir::hasSingleEffect<mlir::MemoryEffects::Read>(op));
}
 
/// Return true if this is a wire or a register or a node.
static bool isDeclaration(Operation *op) {
  return isa<WireOp, RegResetOp, RegOp, NodeOp, MemOp>(op);
}
 
/// Return true if this is a wire or register we're allowed to delete.
static bool isDeletableDeclaration(Operation *op) {
  if (auto name = dyn_cast<FNamableOp>(op))
    if (!name.hasDroppableName())
      return false;
  return !hasDontTouch(op) && AnnotationSet(op).empty();
}
 
namespace {
struct IMDeadCodeElimPass
    : public circt::firrtl::impl::IMDeadCodeElimBase<IMDeadCodeElimPass> {
  using Base::Base;
 
  void runOnOperation() override;
 
  void rewriteModuleSignature(FModuleOp module);
  void rewriteModuleBody(FModuleOp module);
  void eraseEmptyModule(FModuleOp module);
  void forwardConstantOutputPort(FModuleOp module);
 
  /// Return true if the value is known alive.
  bool isKnownAlive(Value value) const {
    assert(value && "null should not be used");
    return liveElements.count(value);
  }
 
  /// Return true if the value is assumed dead.
  bool isAssumedDead(Value value) const { return !isKnownAlive(value); }
  bool isAssumedDead(Operation *op) const {
    return llvm::none_of(op->getResults(),
                         [&](Value value) { return isKnownAlive(value); });
  }
 
  /// Return true if the block is alive.
  bool isBlockExecutable(Block *block) const {
    return executableBlocks.count(block);
  }
 
  void visitUser(Operation *op);
  void visitValue(Value value);
 
  void visitConnect(FConnectLike connect);
  void visitSubelement(Operation *op);
  void markBlockExecutable(Block *block);
  void markBlockUndeletable(Operation *op) {
    markAlive(op->getParentOfType<FModuleOp>());
  }
 
  void markDeclaration(Operation *op);
  void markFInstanceLikeOp(FInstanceLike instanceLike);
  void markUnknownSideEffectOp(Operation *op);
  void visitFInstanceLikeOp(FInstanceLike instanceLike);
  void visitHierPathOp(hw::HierPathOp hierpath);
  void visitModuleOp(FModuleOp module);
 
private:
  /// The set of blocks that are known to execute, or are intrinsically alive.
  DenseSet<Block *> executableBlocks;
 
  InstanceGraph *instanceGraph;
 
  // The type with which we associate liveness.
  using ElementType =
      std::variant<Value, FModuleOp, FInstanceLike, hw::HierPathOp>;
 
  void markAlive(ElementType element) {
    if (!liveElements.insert(element).second)
      return;
    worklist.push_back(element);
  }
 
  /// A worklist of values whose liveness recently changed, indicating
  /// the users need to be reprocessed.
  SmallVector<ElementType, 64> worklist;
  llvm::DenseSet<ElementType> liveElements;
 
  /// A map from instances to hierpaths whose last path is the associated
  /// instance.
  DenseMap<InstanceOp, SmallVector<hw::HierPathOp>> instanceToHierPaths;
 
  /// Hierpath to its users (=non-local annotation targets).
  DenseMap<hw::HierPathOp, SetVector<ElementType>> hierPathToElements;
 
  /// A cache for a (inner)symbol lookp.
  circt::hw::InnerRefNamespace *innerRefNamespace;
  mlir::SymbolTable *symbolTable;
};
} // namespace
 
void IMDeadCodeElimPass::visitFInstanceLikeOp(FInstanceLike instanceLike) {
  auto instance = dyn_cast<InstanceOp>(*instanceLike);
 
  // This is already marked as alive in the initialization.
  if (!instance)
    return;
 
  markBlockUndeletable(instance);
 
  auto module = instance.getReferencedModule<FModuleOp>(*instanceGraph);
 
  if (!module)
    return;
 
  // NOTE: Don't call `markAlive(module)` here as liveness of instance doesn't
  // imply the global liveness of the module.
 
  // Propgate liveness through hierpath.
  for (auto hierPath : instanceToHierPaths[instance])
    markAlive(hierPath);
 
  // Input ports get alive only when the instance is alive.
  for (auto &blockArg : module.getBody().getArguments()) {
    auto portNo = blockArg.getArgNumber();
    if (module.getPortDirection(portNo) == Direction::In &&
        isKnownAlive(module.getArgument(portNo)))
      markAlive(instance.getResult(portNo));
  }
}
 
void IMDeadCodeElimPass::visitModuleOp(FModuleOp module) {
  // If the module needs to be alive, so are its instances and instance choices.
  for (auto *use : instanceGraph->lookup(module)->uses()) {
    if (auto instance = use->getInstance<FInstanceLike>())
      markAlive(instance);
  }
}
 
void IMDeadCodeElimPass::visitHierPathOp(hw::HierPathOp hierPathOp) {
  // If the hierpath is alive, mark all instances on the path alive.
  for (auto path : hierPathOp.getNamepathAttr())
    if (auto innerRef = dyn_cast<hw::InnerRefAttr>(path)) {
      auto *op = innerRefNamespace->lookupOp(innerRef);
      if (auto instance = dyn_cast_or_null<InstanceOp>(op))
        markAlive(instance);
    }
 
  for (auto elem : hierPathToElements[hierPathOp])
    markAlive(elem);
}
 
void IMDeadCodeElimPass::markDeclaration(Operation *op) {
  assert(isDeclaration(op) && "only a declaration is expected");
  if (!isDeletableDeclaration(op)) {
    for (auto result : op->getResults())
      markAlive(result);
    markBlockUndeletable(op);
  }
}
 
void IMDeadCodeElimPass::markUnknownSideEffectOp(Operation *op) {
  // For operations with side effects, pessimistically mark results and
  // operands as alive.
  for (auto result : op->getResults())
    markAlive(result);
  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) {
  LLVM_DEBUG(llvm::dbgs() << "Visit: " << *op << "\n");
  if (auto connectOp = dyn_cast<FConnectLike>(op))
    return visitConnect(connectOp);
  if (isa<SubfieldOp, SubindexOp, SubaccessOp, ObjectSubfieldOp>(op))
    return visitSubelement(op);
}
 
void IMDeadCodeElimPass::markFInstanceLikeOp(FInstanceLike instanceLike) {
  if (auto instance = dyn_cast<InstanceOp>(*instanceLike)) {
    // Get the module being referenced.
    Operation *op = instance.getReferencedModule(*instanceGraph);
 
    // If this is an extmodule, just remember that any inputs and inouts are
    // alive.
    if (!isa<FModuleOp>(op)) {
      auto module = dyn_cast<FModuleLike>(op);
      for (auto resultNo : llvm::seq(0u, instance.getNumResults())) {
        // If this is an output to the extmodule, we can ignore it.
        if (module.getPortDirection(resultNo) == Direction::Out)
          continue;
 
        // Otherwise this is an input from it or an inout, mark it as alive.
        markAlive(instance.getResult(resultNo));
      }
      markAlive(instance);
 
      return;
    }
 
    // Otherwise this is a defined module.
    auto fModule = cast<FModuleOp>(op);
    markBlockExecutable(fModule.getBodyBlock());
    return;
  }
 
  // Conservatively mark everything.
 
  markAlive(instanceLike);
  markBlockUndeletable(instanceLike);
 
  // Mark all results (ports) as alive.
  for (auto result : instanceLike->getResults())
    markAlive(result);
 
  // Mark all possible target modules as executable.
  for (auto moduleName : instanceLike.getReferencedModuleNamesAttr()) {
    auto *node = instanceGraph->lookup(cast<StringAttr>(moduleName));
    if (!node)
      continue;
 
    if (auto fModule = dyn_cast<FModuleOp>(*node->getModule())) {
      markBlockExecutable(fModule.getBodyBlock());
      for (auto result : fModule.getBodyBlock()->getArguments())
        markAlive(result);
    }
  }
}
 
void IMDeadCodeElimPass::markBlockExecutable(Block *block) {
  if (!executableBlocks.insert(block).second)
    return; // Already executable.
 
  auto fmodule = dyn_cast<FModuleOp>(block->getParentOp());
  if (fmodule && (fmodule.isPublic() || removePortsOnly))
    markAlive(fmodule);
 
  // Mark ports with don't touch as alive.
  for (auto blockArg : block->getArguments())
    if (hasDontTouch(blockArg)) {
      markAlive(blockArg);
      if (fmodule)
        markAlive(fmodule);
    }
 
  for (auto &op : *block) {
    if (auto instance = dyn_cast<FInstanceLike>(op)) {
      markFInstanceLikeOp(instance);
      continue;
    }
 
    // 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);
      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. std::nullopt represents an
  // invalid value.
  SmallVector<std::pair<unsigned, std::optional<APSInt>>>
      constantPortIndicesAndValues;
  auto ports = module.getPorts();
  auto *instanceGraphNode = instanceGraph->lookup(module);
 
  for (const auto &e : llvm::enumerate(ports)) {
    unsigned index = e.index();
    auto port = e.value();
    auto arg = module.getArgument(index);
 
    // If the port has don't touch, don't propagate the constant value.
    if (!port.isOutput() || hasDontTouch(arg))
      continue;
 
    // Remember the index and constant value connected to an output port.
    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 or invalid port, abort.
  if (constantPortIndicesAndValues.empty())
    return;
 
  // Rewrite all uses.
  for (auto *use : instanceGraphNode->uses()) {
    auto instance = use->getInstance<InstanceOp>();
    // Only constprop for instance ops.
    if (!instance)
      continue;
    ImplicitLocOpBuilder builder(instance.getLoc(), instance);
    for (auto [index, constant] : constantPortIndicesAndValues) {
      auto result = instance.getResult(index);
      assert(ports[index].isOutput() && "must be an output port");
 
      // 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);
    }
  }
}
 
void IMDeadCodeElimPass::runOnOperation() {
  CIRCT_DEBUG_SCOPED_PASS_LOGGER(this);
 
  auto circuits = getOperation().getOps<CircuitOp>();
  if (circuits.empty())
    return;
 
  auto circuit = *circuits.begin();
 
  if (!llvm::hasSingleElement(circuits)) {
    mlir::emitError(circuit.getLoc(),
                    "cannot process multiple circuit operations")
            .attachNote((*std::next(circuits.begin())).getLoc())
        << "second circuit here";
    return signalPassFailure();
  }
 
  instanceGraph = &getChildAnalysis<InstanceGraph>(circuit);
  symbolTable = &getChildAnalysis<SymbolTable>(circuit);
  auto &istc = getChildAnalysis<hw::InnerSymbolTableCollection>(circuit);
 
  circt::hw::InnerRefNamespace theInnerRefNamespace{*symbolTable, istc};
  innerRefNamespace = &theInnerRefNamespace;
 
  // Walk attributes and find unknown uses of inner symbols or hierpaths.
  getOperation().walk([&](Operation *op) {
    if (isa<FModuleOp>(op)) // Port or module annotations are ok to ignore.
      return;
 
    if (auto hierPath = dyn_cast<hw::HierPathOp>(op)) {
      auto namePath = hierPath.getNamepath().getValue();
      // If the hierpath is public or ill-formed, the verifier should have
      // caught the error. Conservatively mark the symbol as alive.
      if (hierPath.isPublic() || namePath.size() <= 1 ||
          isa<hw::InnerRefAttr>(namePath.back()))
        return markAlive(hierPath);
 
      if (auto instance =
              dyn_cast_or_null<firrtl::InstanceOp>(innerRefNamespace->lookupOp(
                  cast<hw::InnerRefAttr>(namePath.drop_back().back()))))
        instanceToHierPaths[instance].push_back(hierPath);
      return;
    }
 
    // If there is an unknown symbol or inner symbol use, mark all of them
    // alive.
    op->getAttrDictionary().walk([&](Attribute attr) {
      if (auto innerRef = dyn_cast<hw::InnerRefAttr>(attr)) {
        // Mark instances alive that are targeted by an inner ref.
        if (auto instance = dyn_cast_or_null<firrtl::InstanceOp>(
                innerRefNamespace->lookupOp(innerRef)))
          markAlive(instance);
        return;
      }
 
      if (auto symbolRef = dyn_cast<FlatSymbolRefAttr>(attr)) {
        auto *symbol = symbolTable->lookup(symbolRef.getAttr());
        if (!symbol)
          return;
 
        // Mark referenced hierarchical paths alive.
        if (auto hierPath = dyn_cast<hw::HierPathOp>(symbol))
          markAlive(hierPath);
 
        // Mark modules referenced by unknown ops alive.
        if (auto module = dyn_cast<FModuleOp>(symbol)) {
          if (!isa<firrtl::InstanceOp>(op)) {
            LLVM_DEBUG(llvm::dbgs()
                       << "Unknown use of " << module.getModuleNameAttr()
                       << " in " << op->getName() << "\n");
            markAlive(module);
            markBlockExecutable(module.getBodyBlock());
          }
        }
 
        return;
      }
    });
  });
 
  // Create a vector of modules in the post order of instance graph.
  // FIXME: We copy the list of modules into a vector first to avoid iterator
  // invalidation while we mutate the instance graph. See issue 3387.
  SmallVector<FModuleOp, 0> modules(llvm::make_filter_range(
      llvm::map_range(
          llvm::post_order(instanceGraph),
          [](auto *node) { return dyn_cast<FModuleOp>(*node->getModule()); }),
      [](auto module) { return module; }));
 
  // Forward constant output ports to caller sides so that we can eliminate
  // constant outputs.
  for (auto module : modules)
    forwardConstantOutputPort(module);
 
  for (auto module : circuit.getBodyBlock()->getOps<FModuleOp>()) {
    bool isPublic = module.isPublic();
    if (isPublic || removePortsOnly) {
      markBlockExecutable(module.getBodyBlock());
      // 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
    // targets. `portId` is `-1` for module annotations.
    auto visitAnnotation = [&](int portId, Annotation anno) -> bool {
      auto hierPathSym = anno.getMember<FlatSymbolRefAttr>("circt.nonlocal");
      hw::HierPathOp hierPathOp;
      if (hierPathSym)
        hierPathOp =
            symbolTable->template lookup<hw::HierPathOp>(hierPathSym.getAttr());
 
      if (hierPathOp)
        markAlive(hierPathOp);
      if (portId >= 0)
        markAlive(module.getArgument(portId));
      markAlive(module);
      return false;
    };
 
    AnnotationSet::removePortAnnotations(module, visitAnnotation);
    AnnotationSet::removeAnnotations(
        module, std::bind(visitAnnotation, -1, std::placeholders::_1));
  }
 
  // If an element changed liveness then propagate liveness through it.
  while (!worklist.empty()) {
    auto v = worklist.pop_back_val();
    if (auto *value = std::get_if<Value>(&v))
      visitValue(*value);
    else if (auto *instance = std::get_if<FInstanceLike>(&v))
      visitFInstanceLikeOp(*instance);
    else if (auto *hierpath = std::get_if<hw::HierPathOp>(&v))
      visitHierPathOp(*hierpath);
    else if (auto *module = std::get_if<FModuleOp>(&v))
      visitModuleOp(*module);
  }
 
  // Rewrite module signatures or delete unreachable modules.
  for (auto module : llvm::make_early_inc_range(
           circuit.getBodyBlock()->getOps<FModuleOp>())) {
    if (isBlockExecutable(module.getBodyBlock()))
      rewriteModuleSignature(module);
    else {
      // If the module is unreachable from the toplevel, just delete it.
      // Note that post-order traversal on the instance graph never visit
      // unreachable modules so it's safe to erase the module even though
      // `modules` seems to be capturing module pointers.
      module.erase();
    }
  }
 
  // Rewrite module bodies parallelly.
  mlir::parallelForEach(circuit.getContext(),
                        circuit.getBodyBlock()->getOps<FModuleOp>(),
                        [&](auto op) { rewriteModuleBody(op); });
 
  // Clean up hierpaths.
  for (auto op : llvm::make_early_inc_range(
           circuit.getBodyBlock()->getOps<hw::HierPathOp>()))
    if (!liveElements.count(op))
      op.erase();
 
  if (!removePortsOnly)
    for (auto module : modules)
      eraseEmptyModule(module);
 
  // Clean up data structures.
  executableBlocks.clear();
  liveElements.clear();
  instanceToHierPaths.clear();
  hierPathToElements.clear();
}
 
void IMDeadCodeElimPass::visitValue(Value value) {
  assert(isKnownAlive(value) && "only alive values reach here");
 
  // Propagate liveness through users.
  for (Operation *user : value.getUsers())
    visitUser(user);
 
  if (auto blockArg = dyn_cast<BlockArgument>(value)) {
    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()) {
          // If this is not an instance, it must be marked overdefined already.
          if (auto instance = instRec->getInstance<InstanceOp>()) {
            if (liveElements.contains(instance))
              markAlive(instance.getResult(blockArg.getArgNumber()));
          }
        }
      }
 
      // Any domain ports associated with non-domain ports are also alive.
      if (!type_isa<DomainType>(blockArg.getType()))
        for (auto domain : cast<ArrayAttr>(
                 module.getDomainInfoAttrForPort(blockArg.getArgNumber())))
          markAlive(module.getArgument(
              cast<IntegerAttr>(domain).getValue().getZExtValue()));
      return;
    }
  }
 
  // Marking an instance port as alive propagates to the corresponding port of
  // the module.
  if (auto instance = value.getDefiningOp<InstanceOp>()) {
    auto instanceResult = cast<mlir::OpResult>(value);
    // Update the src, when it's an instance op.
    auto module = instance.getReferencedModule<FModuleOp>(*instanceGraph);
 
    // Propagate liveness only when a port is output.
    if (!module || module.getPortDirection(instanceResult.getResultNumber()) ==
                       Direction::In)
      return;
 
    markAlive(instance);
 
    BlockArgument modulePortVal =
        module.getArgument(instanceResult.getResultNumber());
    return markAlive(modulePortVal);
  }
 
  // If a port of a memory is alive, all other ports are.
  if (auto mem = value.getDefiningOp<MemOp>()) {
    for (auto port : mem->getResults())
      markAlive(port);
    return;
  }
 
  // 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>();
      fop && fop.isForceable() &&
      (fop.getData() == value || fop.getDataRef() == value)) {
    markAlive(fop.getData());
    markAlive(fop.getDataRef());
  }
}
 
void IMDeadCodeElimPass::visitConnect(FConnectLike connect) {
  // If the dest is alive, mark the source value as alive.
  if (isKnownAlive(connect.getDest()))
    markAlive(connect.getSrc());
}
 
void IMDeadCodeElimPass::visitSubelement(Operation *op) {
  if (isKnownAlive(op->getOperand(0)))
    markAlive(op->getResult(0));
}
 
void IMDeadCodeElimPass::rewriteModuleBody(FModuleOp module) {
  assert(isBlockExecutable(module.getBodyBlock()) &&
         "unreachable modules must be already deleted");
 
  auto removeDeadNonLocalAnnotations = [&](int _, Annotation anno) -> bool {
    auto hierPathSym = anno.getMember<FlatSymbolRefAttr>("circt.nonlocal");
    if (!hierPathSym)
      return false;
    auto hierPathOp =
        symbolTable->template lookup<hw::HierPathOp>(hierPathSym.getAttr());
    return !liveElements.count(hierPathOp);
  };
 
  AnnotationSet::removePortAnnotations(module, removeDeadNonLocalAnnotations);
  AnnotationSet::removeAnnotations(
      module,
      std::bind(removeDeadNonLocalAnnotations, -1, std::placeholders::_1));
 
  // Walk the IR bottom-up when deleting operations.
  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) {
  assert(isBlockExecutable(module.getBodyBlock()) &&
         "unreachable modules must be already deleted");
  InstanceGraphNode *instanceGraphNode = instanceGraph->lookup(module);
  LLVM_DEBUG(llvm::dbgs() << "Prune ports of module: " << module.getName()
                          << "\n");
 
  auto replaceInstanceResultWithWire =
      [&](ImplicitLocOpBuilder &builder, unsigned index, InstanceOp instance) {
        auto result = instance.getResult(index);
        if (isAssumedDead(result)) {
          // If the result is dead, replace the result with an unrealized
          // conversion cast which works as a dummy placeholder.
          auto wire =
              mlir::UnrealizedConversionCastOp::create(
                  builder, ArrayRef<Type>{result.getType()}, ArrayRef<Value>{})
                  ->getResult(0);
          result.replaceAllUsesWith(wire);
          return;
        }
 
        Value wire = WireOp::create(builder, result.getType()).getResult();
        result.replaceAllUsesWith(wire);
        // If a module port is dead but its instance result is alive, the port
        // is used as a temporary wire so make sure that a replaced wire is
        // putted into `liveSet`.
        liveElements.erase(result);
        liveElements.insert(wire);
      };
 
  // First, delete dead instances.
  for (auto *use : llvm::make_early_inc_range(instanceGraphNode->uses())) {
    auto instanceLike = use->getInstance<FInstanceLike>();
    if (instanceLike && !liveElements.count(instanceLike)) {
      auto instance = cast<InstanceOp>(instanceLike);
      // Replace old instance results with dummy wires.
      ImplicitLocOpBuilder builder(instance.getLoc(), instance);
      for (auto index : llvm::seq(0u, instance.getNumResults()))
        replaceInstanceResultWithWire(builder, index, instance);
      // Make sure that we update the instance graph.
      use->erase();
      instance.erase();
    }
  }
 
  // Ports of public modules cannot be modified.
  if (module.isPublic())
    return;
 
  unsigned numOldPorts = module.getNumPorts();
  llvm::BitVector deadPortIndexes(numOldPorts);
 
  ImplicitLocOpBuilder builder(module.getLoc(), module.getContext());
  builder.setInsertionPointToStart(module.getBodyBlock());
  auto oldPorts = module.getPorts();
 
  for (auto index : llvm::seq(0u, numOldPorts)) {
    auto argument = module.getArgument(index);
    assert((!hasDontTouch(argument) || isKnownAlive(argument)) &&
           "If the port has don't touch, it should be known alive");
 
    // If the port has dontTouch, skip.
    if (hasDontTouch(argument))
      continue;
 
    if (isKnownAlive(argument)) {
 
      // If an output port is only used internally in the module, then we can
      // remove the port and replace it with a wire.
      if (module.getPortDirection(index) == Direction::In)
        continue;
 
      // Check if the output port is demanded by any instance.  If not, then it
      // is only demanded internally to the module.
      if (llvm::any_of(instanceGraph->lookup(module)->uses(),
                       [&](InstanceRecord *record) {
                         return isKnownAlive(
                             record->getInstance().getOperation()->getResult(
                                 index));
                       }))
        continue;
 
      // Ok, this port is used only within its defined module. So we can replace
      // the port with a wire.
      auto wire = WireOp::create(builder, argument.getType()).getResult();
 
      // Since `liveSet` contains the port, we have to erase it from the set.
      liveElements.erase(argument);
      liveElements.insert(wire);
      argument.replaceAllUsesWith(wire);
      deadPortIndexes.set(index);
      continue;
    }
 
    // Replace the port with a dummy wire. This wire should be erased within
    // `rewriteModuleBody`.
    Value wire =
        mlir::UnrealizedConversionCastOp::create(
            builder, ArrayRef<Type>{argument.getType()}, ArrayRef<Value>{})
            ->getResult(0);
 
    argument.replaceAllUsesWith(wire);
    assert(isAssumedDead(wire) && "dummy wire must be dead");
    deadPortIndexes.set(index);
  }
 
  // If there is nothing to remove, abort.
  if (deadPortIndexes.none())
    return;
 
  // Erase arguments of the old module from liveSet to prevent from creating
  // dangling pointers.
  for (auto arg : module.getArguments())
    liveElements.erase(arg);
 
  // Delete ports from the module.
  module.erasePorts(deadPortIndexes);
 
  // Add arguments of the new module to liveSet.
  for (auto arg : module.getArguments())
    liveElements.insert(arg);
 
  // Rewrite all uses.
  for (auto *use : llvm::make_early_inc_range(instanceGraphNode->uses())) {
    auto instance = use->getInstance<InstanceOp>();
    if (!instance) {
      assert(!use->getInstance<FInstanceLike>() &&
             "if this is not an instance, it must be marked overdefined");
      continue;
    }
 
    ImplicitLocOpBuilder builder(instance.getLoc(), instance);
    // Replace old instance results with dummy wires.
    for (auto index : deadPortIndexes.set_bits())
      replaceInstanceResultWithWire(builder, index, instance);
 
    // Since we will rewrite instance op, it is necessary to remove old
    // instance results from liveSet.
    for (auto oldResult : instance.getResults())
      liveElements.erase(oldResult);
 
    // Create a new instance op without dead ports.
    auto newInstance =
        instance.cloneWithErasedPortsAndReplaceUses(deadPortIndexes);
 
    // Mark new results as alive.
    for (auto newResult : newInstance->getResults())
      liveElements.insert(newResult);
 
    instanceGraph->replaceInstance(instance, newInstance);
    if (liveElements.contains(instance)) {
      liveElements.erase(instance);
      liveElements.insert(newInstance);
    }
    // Remove old one.
    instance.erase();
  }
 
  numRemovedPorts += deadPortIndexes.count();
}
 
void IMDeadCodeElimPass::eraseEmptyModule(FModuleOp module) {
  // If the module is not empty, just skip.
  if (!module.getBodyBlock()->empty())
    return;
 
  // We cannot delete public modules so generate a warning.
  if (module.isPublic()) {
    mlir::emitWarning(module.getLoc())
        << "module `" << module.getName()
        << "` is empty but cannot be removed because the module is public";
    return;
  }
 
  if (!module.getAnnotations().empty()) {
    module.emitWarning() << "module `" << module.getName()
                         << "` is empty but cannot be removed "
                            "because the module has annotations "
                         << module.getAnnotations();
    return;
  }
 
  if (!module.getBodyBlock()->args_empty()) {
    auto diag = module.emitWarning()
                << "module `" << module.getName()
                << "` is empty but cannot be removed because the "
                   "module has ports ";
    llvm::interleaveComma(module.getPortNames(), diag);
    diag << " are referenced by name or dontTouched";
    return;
  }
 
  // Ok, the module is empty. Delete instances unless they have symbols.
  LLVM_DEBUG(llvm::dbgs() << "Erase " << module.getName() << "\n");
 
  InstanceGraphNode *instanceGraphNode =
      instanceGraph->lookup(module.getModuleNameAttr());
 
  SmallVector<Location> instancesWithSymbols;
  for (auto *use : llvm::make_early_inc_range(instanceGraphNode->uses())) {
    auto instance = use->getInstance<InstanceOp>();
    if (!instance)
      continue;
    if (instance.getInnerSym()) {
      instancesWithSymbols.push_back(instance.getLoc());
      continue;
    }
    use->erase();
    instance.erase();
  }
 
  // If there is an instance with a symbol, we don't delete the module itself.
  if (!instancesWithSymbols.empty()) {
    auto diag = module.emitWarning()
                << "module `" << module.getName()
                << "` is empty but cannot be removed because an instance is "
                   "referenced by name";
    diag.attachNote(FusedLoc::get(&getContext(), instancesWithSymbols))
        << "these are instances with symbols";
    return;
  }
 
  // We cannot delete alive modules.
  if (liveElements.contains(module))
    return;
 
  instanceGraph->erase(instanceGraphNode);
  module.erase();
  ++numErasedModules;
}