[HW] Add Passes: hw-expunge-module, hw-tree-shake

CircuitCoder · CircuitCoder · commit 10489b2dafdc · 2025-01-13T01:19:56.000+08:00
diff --git a/include/circt/Dialect/HW/HWPasses.h b/include/circt/Dialect/HW/HWPasses.h
@@ -33,6 +33,8 @@ std::unique_ptr<mlir::Pass> createFlattenIOPass(bool recursiveFlag = true,
 std::unique_ptr<mlir::Pass> createVerifyInnerRefNamespacePass();
 std::unique_ptr<mlir::Pass> createFlattenModulesPass();
 std::unique_ptr<mlir::Pass> createFooWiresPass();
+std::unique_ptr<mlir::Pass> createHWExpungeModulePass();
+std::unique_ptr<mlir::Pass> createHWTreeShakePass();
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
diff --git a/include/circt/Dialect/HW/Passes.td b/include/circt/Dialect/HW/Passes.td
@@ -75,6 +75,38 @@ def VerifyInnerRefNamespace : Pass<"hw-verify-irn"> {
   let constructor =  "circt::hw::createVerifyInnerRefNamespacePass()";
 }
 
+def HWExpungeModule : Pass<"hw-expunge-module", "mlir::ModuleOp"> {
+  let summary = "Remove module from the hierarchy, and recursively expose their ports to upper level.";
+  let description = [{
+    This pass removes a list of modules from the hierarchy on-by-one, recursively exposing their ports to upper level.
+    The newly generated ports are by default named as <instance_path>__<port_name>. During a naming conflict, an warning would be genreated,
+    and an random suffix would be added to the <instance_path> part.
+
+    For each given (transitive) parent module, the prefix can alternatively be specified by option instead of using the instance path.
+  }];
+  let constructor =  "circt::hw::createHWExpungeModulePass()";
+
+  let options = [
+    ListOption<"modules", "modules", "std::string",
+      "Comma separated list of module names to be removed from the hierarchy.">,
+    ListOption<"portPrefixes", "port-prefixes", "std::string",
+      "Specify the prefix for ports of a given parent module's expunged childen. Each specification is formatted as <module>:<instance-path>=<prefix>. Only affect the top-most level module of the instance path.">,
+  ];
+}
+
+def HWTreeShake : Pass<"hw-tree-shake", "mlir::ModuleOp"> {
+  let summary = "Remove unused modules.";
+  let description = [{
+    This pass removes all modules besides a specified list of modules and their transitive dependencies.
+  }];
+  let constructor =  "circt::hw::createHWTreeShakePass()";
+
+  let options = [
+    ListOption<"keep", "keep", "std::string",
+      "Comma separated list of module names to be kept.">,
+  ];
+}
+
 /**
  * Tutorial Pass, doesn't do anything interesting
  */
diff --git a/lib/Dialect/HW/Transforms/CMakeLists.txt b/lib/Dialect/HW/Transforms/CMakeLists.txt
@@ -6,6 +6,8 @@ add_circt_dialect_library(CIRCTHWTransforms
   VerifyInnerRefNamespace.cpp
   FlattenModules.cpp
   FooWires.cpp
+  HWExpungeModule.cpp
+  HWTreeShake.cpp
 
   DEPENDS
   CIRCTHWTransformsIncGen
diff --git a/lib/Dialect/HW/Transforms/HWExpungeModule.cpp b/lib/Dialect/HW/Transforms/HWExpungeModule.cpp
@@ -0,0 +1,330 @@
+#include "circt/Dialect/HW/HWOps.h"
+#include "circt/Dialect/HW/HWPasses.h"
+#include "circt/Dialect/HW/HWTypes.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/ImmutableList.h"
+#include "llvm/Support/Regex.h"
+
+namespace circt {
+namespace hw {
+#define GEN_PASS_DEF_HWEXPUNGEMODULE
+#include "circt/Dialect/HW/Passes.h.inc"
+} // namespace hw
+} // namespace circt
+
+namespace {
+struct HWExpungeModulePass
+    : circt::hw::impl::HWExpungeModuleBase<HWExpungeModulePass> {
+  void runOnOperation() override;
+};
+
+struct InstPathSeg {
+  llvm::StringRef seg;
+
+  InstPathSeg(llvm::StringRef seg) : seg(seg) {}
+  const llvm::StringRef &getSeg() const { return seg; }
+  operator llvm::StringRef() const { return seg; }
+
+  void Profile(llvm::FoldingSetNodeID &ID) const { ID.AddString(seg); }
+};
+using InstPath = llvm::ImmutableList<InstPathSeg>;
+std::string defaultPrefix(InstPath path) {
+  std::string accum;
+  while (!path.isEmpty()) {
+    accum += path.getHead().getSeg();
+    accum += "_";
+    path = path.getTail();
+  }
+  accum += "_";
+  return std::move(accum);
+}
+
+// The regex for port prefix specification
+// "^([@#a-zA-Z0-9_]+):([a-zA-Z0-9_]+)(\\.[a-zA-Z0-9_]+)*=([a-zA-Z0-9_]+)$"
+// Unfortunately, the LLVM Regex cannot capture repeating capture groups, so
+// manually parse the spec This parser may accept identifiers with invalid
+// characters
+
+std::variant<std::tuple<llvm::StringRef, InstPath, llvm::StringRef>,
+             std::string>
+parsePrefixSpec(llvm::StringRef in, InstPath::Factory &listFac) {
+  auto [l, r] = in.split("=");
+  if (r == "")
+    return "No '=' found in input";
+  auto [ll, lr] = l.split(":");
+  if (lr == "")
+    return "No ':' found before '='";
+  llvm::SmallVector<llvm::StringRef, 4> segs;
+  while (lr != "") {
+    auto [seg, rest] = lr.split(".");
+    segs.push_back(seg);
+    lr = rest;
+  }
+  InstPath path;
+  for (auto &seg : llvm::reverse(segs))
+    path = listFac.add(seg, path);
+  return std::make_tuple(ll, path, r);
+}
+} // namespace
+
+void HWExpungeModulePass::runOnOperation() {
+  auto root = getOperation();
+  llvm::DenseMap<mlir::StringRef, circt::hw::HWModuleLike> allModules;
+  root.walk(
+      [&](circt::hw::HWModuleLike mod) { allModules[mod.getName()] = mod; });
+  // Reverse-topo order generated by module instances.
+  // The topo sorted order does not change throught out the operation. It only
+  // gets weakened, but still valid.
+  llvm::SmallVector<circt::hw::HWModuleLike> revTopo;
+  llvm::DenseSet<mlir::StringRef> visited;
+  auto visit = [&allModules, &visited,
+                &revTopo](auto &self, circt::hw::HWModuleLike mod) -> void {
+    if (visited.contains(mod.getName()))
+      return;
+    visited.insert(mod.getName());
+
+    mod.walk([&](circt::hw::InstanceOp inst) {
+      auto instModName = inst.getModuleName();
+      auto instMod = allModules.lookup(instModName);
+      if (!instMod)
+        inst.emitError("Unknown module ") << instModName;
+      self(self, instMod);
+    });
+
+    // Reverse topo order, insert on stack pop
+    revTopo.push_back(mod);
+  };
+  for (auto &[_, mod] : allModules)
+    visit(visit, mod);
+  assert(revTopo.size() == allModules.size());
+
+  // Instance path.
+  InstPath::Factory pathFactory;
+
+  // Process port prefix specifications
+  // (Module name, Instance path) -> Prefix
+  llvm::DenseMap<std::pair<mlir::StringRef, InstPath>, mlir::StringRef>
+      designatedPrefixes;
+  bool containsFailure = false;
+  for (const auto &raw : portPrefixes) {
+    auto matched = parsePrefixSpec(raw, pathFactory);
+    if (std::holds_alternative<std::string>(matched)) {
+      llvm::errs() << "Invalid port prefix specification: " << raw << "\n";
+      llvm::errs() << "Error: " << std::get<std::string>(matched) << "\n";
+      containsFailure = true;
+      continue;
+    }
+
+    auto [module, path, prefix] =
+        std::get<std::tuple<llvm::StringRef, InstPath, llvm::StringRef>>(
+            matched);
+    if (!allModules.contains(module)) {
+      llvm::errs() << "Module not found in port prefix specification: "
+                   << module << "\n";
+      llvm::errs() << "From specification: " << raw << "\n";
+      containsFailure = true;
+      continue;
+    }
+
+    // Skip checking instance paths' existence. Non-existent paths are ignored
+    designatedPrefixes.insert({{module, path}, prefix});
+  }
+
+  if (containsFailure)
+    return signalPassFailure();
+
+  // Instance path * prefix name
+  using ReplacedDescendent = std::pair<InstPath, std::string>;
+  // This map holds the expunged descendents of a module
+  llvm::DenseMap<llvm::StringRef, llvm::SmallVector<ReplacedDescendent>>
+      expungedDescendents;
+  for (auto &expunging : this->modules) {
+    // Clear expungedDescendents
+    for (auto &it : expungedDescendents)
+      it.getSecond().clear();
+
+    auto expungingMod = allModules.lookup(expunging);
+    if (!expungingMod)
+      continue; // Ignored missing modules
+    auto expungingModTy = expungingMod.getHWModuleType();
+    auto expungingModPorts = expungingModTy.getPorts();
+
+    auto createPortsOn = [&expungingModPorts](circt::hw::HWModuleOp mod,
+                                              const std::string &prefix,
+                                              auto genOutput, auto emitInput) {
+      mlir::OpBuilder builder(mod);
+      // Create ports using *REVERSE* direction of their definitions
+      for (auto &port : expungingModPorts) {
+        auto defaultName = prefix + port.name.getValue();
+        auto finalName = defaultName;
+        if (port.dir == circt::hw::PortInfo::Input) {
+          auto val = genOutput(port);
+          assert(val.getType() == port.type);
+          mod.appendOutput(finalName, val);
+        } else if (port.dir == circt::hw::PortInfo::Output) {
+          auto [_, arg] = mod.appendInput(finalName, port.type);
+          emitInput(port, arg);
+        }
+      }
+    };
+
+    for (auto &processingRaw : revTopo) {
+      // Skip extmodule and intmodule because they cannot contain anything
+      if (!llvm::isa<circt::hw::HWModuleOp>(processingRaw))
+        continue;
+      circt::hw::HWModuleOp processing =
+          llvm::cast<circt::hw::HWModuleOp>(processingRaw);
+
+      std::optional<decltype(expungedDescendents.lookup("")) *>
+          outerExpDescHold = {};
+      auto getOuterExpDesc = [&]() -> decltype(**outerExpDescHold) {
+        if (!outerExpDescHold.has_value())
+          outerExpDescHold = {
+              &expungedDescendents.insert({processing.getName(), {}})
+                   .first->getSecond()};
+        return **outerExpDescHold;
+      };
+
+      mlir::OpBuilder outerBuilder(processing);
+
+      processing.walk([&](circt::hw::InstanceOp inst) {
+        auto instName = inst.getInstanceName();
+        auto instMod = allModules.lookup(inst.getModuleName());
+
+        if (instMod.getOutputNames().size() != inst.getResults().size() ||
+            instMod.getNumInputPorts() != inst.getInputs().size()) {
+          // Module have been modified during this pass, create new instances
+          assert(instMod.getNumOutputPorts() >= inst.getResults().size());
+          assert(instMod.getNumInputPorts() >= inst.getInputs().size());
+
+          auto instModInTypes = instMod.getInputTypes();
+
+          llvm::SmallVector<mlir::Value> newInputs;
+          newInputs.reserve(instMod.getNumInputPorts());
+
+          outerBuilder.setInsertionPointAfter(inst);
+
+          // Appended inputs are at the end of the input list
+          for (size_t i = 0; i < instMod.getNumInputPorts(); ++i) {
+            mlir::Value input;
+            if (i < inst.getNumInputPorts()) {
+              input = inst.getInputs()[i];
+              if (auto existingName = inst.getInputName(i))
+                assert(existingName == instMod.getInputName(i));
+            } else {
+              input =
+                  outerBuilder
+                      .create<mlir::UnrealizedConversionCastOp>(
+                          inst.getLoc(), instModInTypes[i], mlir::ValueRange{})
+                      .getResult(0);
+            }
+            newInputs.push_back(input);
+          }
+
+          auto newInst = outerBuilder.create<circt::hw::InstanceOp>(
+              inst.getLoc(), instMod, inst.getInstanceNameAttr(), newInputs,
+              inst.getParameters(),
+              inst.getInnerSym().value_or<circt::hw::InnerSymAttr>({}));
+
+          for (size_t i = 0; i < inst.getNumResults(); ++i)
+            assert(inst.getOutputName(i) == instMod.getOutputName(i));
+          inst.replaceAllUsesWith(
+              newInst.getResults().slice(0, inst.getNumResults()));
+          inst.erase();
+          inst = newInst;
+        }
+
+        llvm::DenseMap<llvm::StringRef, mlir::Value> instOMap;
+        llvm::DenseMap<llvm::StringRef, mlir::Value> instIMap;
+        assert(instMod.getOutputNames().size() == inst.getResults().size());
+        for (auto [oname, oval] :
+             llvm::zip(instMod.getOutputNames(), inst.getResults()))
+          instOMap[llvm::cast<mlir::StringAttr>(oname).getValue()] = oval;
+        assert(instMod.getInputNames().size() == inst.getInputs().size());
+        for (auto [iname, ival] :
+             llvm::zip(instMod.getInputNames(), inst.getInputs()))
+          instIMap[llvm::cast<mlir::StringAttr>(iname).getValue()] = ival;
+
+        // Get outer expunged descendent first because it may modify the map and
+        // invalidate iterators.
+        auto &outerExpDesc = getOuterExpDesc();
+        auto instExpDesc = expungedDescendents.find(inst.getModuleName());
+
+        if (inst.getModuleName() == expunging) {
+          // Handle the directly expunged module
+          // input maps also useful for directly expunged instance
+
+          auto singletonPath = pathFactory.create(instName);
+
+          auto designatedPrefix =
+              designatedPrefixes.find({processing.getName(), singletonPath});
+          std::string prefix = designatedPrefix != designatedPrefixes.end()
+                                   ? designatedPrefix->getSecond().str()
+                                   : (instName + "__").str();
+
+          // Port name collision is still possible, but current relying on MLIR
+          // to automatically rename input arguments.
+          // TODO: name collision detect
+
+          createPortsOn(
+              processing, prefix,
+              [&](circt::hw::ModulePort port) {
+                // Generate output for outer module, so input for us
+                return instIMap.at(port.name);
+              },
+              [&](circt::hw::ModulePort port, mlir::Value val) {
+                // Generated input for outer module, replace inst results
+                assert(instOMap.contains(port.name));
+                instOMap[port.name].replaceAllUsesWith(val);
+              });
+
+          outerExpDesc.emplace_back(singletonPath, prefix);
+
+          assert(instExpDesc == expungedDescendents.end() ||
+                 instExpDesc->getSecond().size() == 0);
+          inst.erase();
+        } else if (instExpDesc != expungedDescendents.end()) {
+          // Handle all transitive descendents
+          if (instExpDesc->second.size() == 0)
+            return;
+          llvm::DenseMap<llvm::StringRef, mlir::Value> newInputs;
+          for (const auto &exp : instExpDesc->second) {
+            auto newPath = pathFactory.add(instName, exp.first);
+            auto designatedPrefix =
+                designatedPrefixes.find({processing.getName(), newPath});
+            std::string prefix = designatedPrefix != designatedPrefixes.end()
+                                     ? designatedPrefix->getSecond().str()
+                                     : defaultPrefix(newPath);
+
+            // TODO: name collision detect
+
+            createPortsOn(
+                processing, prefix,
+                [&](circt::hw::ModulePort port) {
+                  // Generate output for outer module, directly forward from
+                  // inner inst
+                  return instOMap.at((exp.second + port.name.getValue()).str());
+                },
+                [&](circt::hw::ModulePort port, mlir::Value val) {
+                  // Generated input for outer module, replace inst results.
+                  // The operand in question has to be an backedge
+                  auto in =
+                      instIMap.at((exp.second + port.name.getValue()).str());
+                  auto inDef = in.getDefiningOp();
+                  assert(llvm::isa<mlir::UnrealizedConversionCastOp>(inDef));
+                  in.replaceAllUsesWith(val);
+                  inDef->erase();
+                });
+
+            outerExpDesc.emplace_back(newPath, prefix);
+          }
+        }
+      });
+    }
+  }
+}
+
+std::unique_ptr<mlir::Pass> circt::hw::createHWExpungeModulePass() {
+  return std::make_unique<HWExpungeModulePass>();
+}
diff --git a/lib/Dialect/HW/Transforms/HWTreeShake.cpp b/lib/Dialect/HW/Transforms/HWTreeShake.cpp
diff --git a/test/Dialect/HW/expunge-module.mlir b/test/Dialect/HW/expunge-module.mlir
