ExtractInstances.cpp

Go to the documentation of this file.

//===- ExtractInstances.cpp - Move instances up the hierarchy ---*- 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
//
//===----------------------------------------------------------------------===//
//
// Moves annotated instances upwards in the module hierarchy. Corresponds to the
// `ExtractBlackBoxes`, `ExtractClockGates`, and `ExtractSeqMems` passes in the
// Scala FIRRTL implementation.
//
//===----------------------------------------------------------------------===//
 
#include "circt/Dialect/Emit/EmitOps.h"
#include "circt/Dialect/FIRRTL/AnnotationDetails.h"
#include "circt/Dialect/FIRRTL/FIRRTLAnnotations.h"
#include "circt/Dialect/FIRRTL/FIRRTLInstanceGraph.h"
#include "circt/Dialect/FIRRTL/FIRRTLOps.h"
#include "circt/Dialect/FIRRTL/FIRRTLUtils.h"
#include "circt/Dialect/FIRRTL/NLATable.h"
#include "circt/Dialect/FIRRTL/Namespace.h"
#include "circt/Dialect/FIRRTL/Passes.h"
#include "circt/Dialect/HW/HWAttributes.h"
#include "circt/Dialect/HW/HWDialect.h"
#include "circt/Dialect/HW/InnerSymbolNamespace.h"
#include "circt/Dialect/SV/SVOps.h"
#include "circt/Support/Path.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/ImplicitLocOpBuilder.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/FileUtilities.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/Path.h"
 
#define DEBUG_TYPE "firrtl-extract-instances"
 
namespace circt {
namespace firrtl {
#define GEN_PASS_DEF_EXTRACTINSTANCES
#include "circt/Dialect/FIRRTL/Passes.h.inc"
} // namespace firrtl
} // namespace circt
 
using namespace circt;
using namespace firrtl;
using hw::InnerRefAttr;
 
//===----------------------------------------------------------------------===//
// Pass Implementation
//===----------------------------------------------------------------------===//
 
namespace {
/// All information necessary to move instances about.
struct ExtractionInfo {
  /// A filename into which the performed hierarchy modifications are emitted.
  StringRef traceFilename;
  /// A prefix to attach to the wiring generated by the extraction.
  StringRef prefix;
  /// Optional name of the wrapper module that will hold the moved instance.
  StringRef wrapperModule;
  /// Whether the extraction should stop at the root of the DUT instead of going
  /// past that and extracting into the test harness.
  bool stopAtDUT;
};
 
struct ExtractInstancesPass
    : public circt::firrtl::impl::ExtractInstancesBase<ExtractInstancesPass> {
  void runOnOperation() override;
  void collectAnnos();
  void collectAnno(InstanceOp inst, Annotation anno);
  void extractInstances();
  void groupInstances();
  void createTraceFiles();
 
  /// Get the cached namespace for a module.
  hw::InnerSymbolNamespace &getModuleNamespace(FModuleLike module) {
    return moduleNamespaces.try_emplace(module, module).first->second;
  }
 
  /// Obtain an inner reference to an operation, possibly adding an `inner_sym`
  /// to that operation.
  InnerRefAttr getInnerRefTo(Operation *op) {
    return ::getInnerRefTo(op,
                           [&](FModuleLike mod) -> hw::InnerSymbolNamespace & {
                             return getModuleNamespace(mod);
                           });
  }
 
  /// Create a clone of a `HierPathOp` with a new uniquified name.
  hw::HierPathOp cloneWithNewNameAndPath(hw::HierPathOp pathOp,
                                         ArrayRef<Attribute> newPath) {
    OpBuilder builder(pathOp);
    auto newPathOp = builder.cloneWithoutRegions(pathOp);
    newPathOp.setSymNameAttr(builder.getStringAttr(
        circuitNamespace.newName(newPathOp.getSymName())));
    newPathOp.setNamepathAttr(builder.getArrayAttr(newPath));
    return newPathOp;
  }
 
  /// Return a handle to the unique instance of file with a given name.
  emit::FileOp getOrCreateFile(StringRef fileName) {
    auto [it, inserted] = files.try_emplace(fileName, emit::FileOp{});
    if (inserted) {
      auto builder = ImplicitLocOpBuilder::atBlockEnd(
          UnknownLoc::get(&getContext()), getOperation().getBodyBlock());
      it->second = builder.create<emit::FileOp>(fileName);
    }
    return it->second;
  }
 
  bool anythingChanged;
  bool anyFailures;
 
  InstanceGraph *instanceGraph = nullptr;
 
  /// The modules in the design that are annotated with one or more annotations
  /// relevant for instance extraction.
  DenseMap<Operation *, SmallVector<Annotation, 1>> annotatedModules;
 
  /// All modules that are marked as DUT themselves. Realistically this is only
  /// ever one module in the design.
  DenseSet<Operation *> dutRootModules;
  /// All modules that are marked as DUT themselves, or that have a DUT parent
  /// module.
  DenseSet<Operation *> dutModules;
  /// All DUT module names.
  DenseSet<Attribute> dutModuleNames;
  /// The prefix of the DUT module.  This is used when creating new modules
  /// under the DUT.
  StringRef dutPrefix = "";
 
  /// A worklist of instances that need to be moved.
  SmallVector<std::pair<InstanceOp, ExtractionInfo>> extractionWorklist;
 
  /// A mapping from file names to file ops for de-duplication.
  DenseMap<StringRef, emit::FileOp> files;
 
  /// The path along which instances have been extracted. This essentially
  /// documents the original location of the instance in reverse. Every push
  /// upwards in the hierarchy adds another entry to this path documenting along
  /// which instantiation path each instance was extracted.
  DenseMap<Operation *, SmallVector<InnerRefAttr>> extractionPaths;
 
  /// A map of the original parent modules of instances before they were
  /// extracted. This is used in a corner case during trace file emission.
  DenseMap<Operation *, StringAttr> originalInstanceParents;
 
  /// All extracted instances in their position after moving upwards in the
  /// hierarchy, but before being grouped into an optional submodule.
  SmallVector<std::pair<InstanceOp, ExtractionInfo>> extractedInstances;
 
  // The uniquified wiring prefix for each instance.
  DenseMap<Operation *, SmallString<16>> instPrefices;
 
  /// The current circuit namespace valid within the call to `runOnOperation`.
  CircuitNamespace circuitNamespace;
  /// Cached module namespaces.
  DenseMap<Operation *, hw::InnerSymbolNamespace> moduleNamespaces;
};
} // end anonymous namespace
 
/// Emit the annotated source code for black boxes in a circuit.
void ExtractInstancesPass::runOnOperation() {
  CircuitOp circuitOp = getOperation();
  anythingChanged = false;
  anyFailures = false;
  annotatedModules.clear();
  dutRootModules.clear();
  dutModules.clear();
  dutModuleNames.clear();
  dutPrefix = "";
  extractionWorklist.clear();
  files.clear();
  extractionPaths.clear();
  originalInstanceParents.clear();
  extractedInstances.clear();
  instPrefices.clear();
  moduleNamespaces.clear();
  circuitNamespace.clear();
  circuitNamespace.add(circuitOp);
 
  // Walk the IR and gather all the annotations relevant for extraction that
  // appear on instances and the instantiated modules.
  instanceGraph = &getAnalysis<InstanceGraph>();
  collectAnnos();
  if (anyFailures)
    return signalPassFailure();
 
  // Actually move instances upwards.
  extractInstances();
  if (anyFailures)
    return signalPassFailure();
 
  // Group instances into submodules, if requested.
  groupInstances();
  if (anyFailures)
    return signalPassFailure();
 
  // Generate the trace files that list where each instance was extracted from.
  createTraceFiles();
  if (anyFailures)
    return signalPassFailure();
 
  // If nothing has changed we can preserve the analysis.
  LLVM_DEBUG(llvm::dbgs() << "\n");
  if (!anythingChanged)
    markAllAnalysesPreserved();
}
 
static bool isAnnoInteresting(Annotation anno) {
  return anno.isClass(extractBlackBoxAnnoClass);
}
 
/// Gather the modules and instances annotated to be moved by this pass. This
/// populates the corresponding lists and maps of the pass.
void ExtractInstancesPass::collectAnnos() {
  CircuitOp circuit = getOperation();
 
  // Grab the clock gate extraction annotation on the circuit.
  StringRef clkgateFileName;
  StringRef clkgateWrapperModule;
  AnnotationSet::removeAnnotations(circuit, [&](Annotation anno) {
    if (!anno.isClass(extractClockGatesAnnoClass))
      return false;
    LLVM_DEBUG(llvm::dbgs()
               << "Clock gate extraction config: " << anno.getDict() << "\n");
    auto filenameAttr = anno.getMember<StringAttr>("filename");
    auto groupAttr = anno.getMember<StringAttr>("group");
    if (!filenameAttr) {
      circuit.emitError("missing `filename` attribute in `")
          << anno.getClass() << "` annotation";
      anyFailures = true;
      return true;
    }
 
    if (!clkgateFileName.empty()) {
      circuit.emitError("multiple `")
          << anno.getClass() << "` annotations on circuit";
      anyFailures = true;
      return true;
    }
 
    clkgateFileName = filenameAttr.getValue();
    if (groupAttr)
      clkgateWrapperModule = groupAttr.getValue();
    return true;
  });
 
  // Grab the memory extraction annotation on the circuit.
  StringRef memoryFileName;
  StringRef memoryWrapperModule;
  AnnotationSet::removeAnnotations(circuit, [&](Annotation anno) {
    if (!anno.isClass(extractSeqMemsAnnoClass))
      return false;
    LLVM_DEBUG(llvm::dbgs()
               << "Memory extraction config: " << anno.getDict() << "\n");
    auto filenameAttr = anno.getMember<StringAttr>("filename");
    auto groupAttr = anno.getMember<StringAttr>("group");
    if (!filenameAttr) {
      circuit.emitError("missing `filename` attribute in `")
          << anno.getClass() << "` annotation";
      anyFailures = true;
      return true;
    }
 
    if (!memoryFileName.empty()) {
      circuit.emitError("multiple `")
          << anno.getClass() << "` annotations on circuit";
      anyFailures = true;
      return true;
    }
 
    memoryFileName = filenameAttr.getValue();
    if (groupAttr)
      memoryWrapperModule = groupAttr.getValue();
    return true;
  });
 
  // Gather the annotations on modules. These complement the later per-instance
  // annotations.
  for (auto module : circuit.getOps<FModuleLike>()) {
    AnnotationSet::removeAnnotations(module, [&](Annotation anno) {
      if (anno.isClass(dutAnnoClass)) {
        LLVM_DEBUG(llvm::dbgs()
                   << "Marking DUT `" << module.getModuleName() << "`\n");
        dutRootModules.insert(module);
        dutModules.insert(module);
        if (auto prefix = anno.getMember<StringAttr>("prefix"))
          dutPrefix = prefix;
        return false; // other passes may rely on this anno; keep it
      }
      if (!isAnnoInteresting(anno))
        return false;
      LLVM_DEBUG(llvm::dbgs() << "Annotated module `" << module.getModuleName()
                              << "`:\n  " << anno.getDict() << "\n");
      annotatedModules[module].push_back(anno);
      return true;
    });
  }
 
  // Gather the annotations on instances to be extracted.
  circuit.walk([&](InstanceOp inst) {
    SmallVector<Annotation, 1> instAnnos;
    Operation *module = inst.getReferencedModule(*instanceGraph);
 
    // Module-level annotations.
    auto it = annotatedModules.find(module);
    if (it != annotatedModules.end())
      instAnnos.append(it->second);
 
    // Instance-level annotations.
    AnnotationSet::removeAnnotations(inst, [&](Annotation anno) {
      if (!isAnnoInteresting(anno))
        return false;
      LLVM_DEBUG(llvm::dbgs() << "Annotated instance `" << inst.getName()
                              << "`:\n  " << anno.getDict() << "\n");
      instAnnos.push_back(anno);
      return true;
    });
 
    // No need to do anything about unannotated instances.
    if (instAnnos.empty())
      return;
 
    // Ensure there are no conflicting annotations.
    if (instAnnos.size() > 1) {
      auto d = inst.emitError("multiple extraction annotations on instance `")
               << inst.getName() << "`";
      d.attachNote(inst.getLoc()) << "instance has the following annotations, "
                                     "but at most one is allowed:";
      for (auto anno : instAnnos)
        d.attachNote(inst.getLoc()) << anno.getDict();
      anyFailures = true;
      return;
    }
 
    // Process the annotation.
    collectAnno(inst, instAnnos[0]);
  });
 
  // Propagate the DUT marking to all arbitrarily nested submodules of the DUT.
  LLVM_DEBUG(llvm::dbgs() << "Marking DUT hierarchy\n");
  SmallVector<InstanceGraphNode *> worklist;
  for (Operation *op : dutModules)
    worklist.push_back(
        instanceGraph->lookup(cast<igraph::ModuleOpInterface>(op)));
  while (!worklist.empty()) {
    auto *module = worklist.pop_back_val();
    dutModuleNames.insert(module->getModule().getModuleNameAttr());
    LLVM_DEBUG(llvm::dbgs()
               << "- " << module->getModule().getModuleName() << "\n");
    for (auto *instRecord : *module) {
      auto *target = instRecord->getTarget();
      if (dutModules.insert(target->getModule()).second)
        worklist.push_back(target);
    }
  }
 
  // If clock gate extraction is requested, find instances of extmodules with
  // the corresponding `defname` and mark them as to be extracted.
  // TODO: This defname really shouldn't be hardcoded here. Make this at least
  // somewhat configurable.
  if (!clkgateFileName.empty()) {
    auto clkgateDefNameAttr = StringAttr::get(&getContext(), "EICG_wrapper");
    for (auto module : circuit.getOps<FExtModuleOp>()) {
      if (module.getDefnameAttr() != clkgateDefNameAttr)
        continue;
      LLVM_DEBUG(llvm::dbgs()
                 << "Clock gate `" << module.getModuleName() << "`\n");
      if (!dutModules.contains(module)) {
        LLVM_DEBUG(llvm::dbgs() << "- Ignored (outside DUT)\n");
        continue;
      }
 
      ExtractionInfo info;
      info.traceFilename = clkgateFileName;
      info.prefix = "clock_gate"; // TODO: Don't hardcode this
      info.wrapperModule = clkgateWrapperModule;
      info.stopAtDUT = !info.wrapperModule.empty();
      for (auto *instRecord : instanceGraph->lookup(module)->uses()) {
        if (auto inst = dyn_cast<InstanceOp>(*instRecord->getInstance())) {
          LLVM_DEBUG(llvm::dbgs()
                     << "- Marking `"
                     << inst->getParentOfType<FModuleLike>().getModuleName()
                     << "." << inst.getName() << "`\n");
          extractionWorklist.push_back({inst, info});
        }
      }
    }
  }
 
  // If memory extraction is requested, find instances of `FMemModuleOp` and
  // mark them as to be extracted.
  // somewhat configurable.
  if (!memoryFileName.empty()) {
    // Create a potentially empty file if a name is specified. This is done to
    // align with the SFC implementation of this pass where the file is always
    // created.  This does introduce an additional leading newline in the file.
    getOrCreateFile(memoryFileName);
 
    for (auto module : circuit.getOps<FMemModuleOp>()) {
      LLVM_DEBUG(llvm::dbgs() << "Memory `" << module.getModuleName() << "`\n");
      if (!dutModules.contains(module)) {
        LLVM_DEBUG(llvm::dbgs() << "- Ignored (outside DUT)\n");
        continue;
      }
 
      ExtractionInfo info;
      info.traceFilename = memoryFileName;
      info.prefix = "mem_wiring"; // TODO: Don't hardcode this
      info.wrapperModule = memoryWrapperModule;
      info.stopAtDUT = !info.wrapperModule.empty();
      for (auto *instRecord : instanceGraph->lookup(module)->uses()) {
        if (auto inst = dyn_cast<InstanceOp>(*instRecord->getInstance())) {
          LLVM_DEBUG(llvm::dbgs()
                     << "- Marking `"
                     << inst->getParentOfType<FModuleLike>().getModuleName()
                     << "." << inst.getName() << "`\n");
          extractionWorklist.push_back({inst, info});
        }
      }
    }
  }
}
 
/// Process an extraction annotation on an instance into a corresponding
/// `ExtractionInfo` and a spot on the worklist for later moving things around.
void ExtractInstancesPass::collectAnno(InstanceOp inst, Annotation anno) {
  LLVM_DEBUG(llvm::dbgs() << "Processing instance `" << inst.getName() << "` "
                          << anno.getDict() << "\n");
 
  auto getStringOrError = [&](StringRef member) {
    auto attr = anno.getMember<StringAttr>(member);
    if (!attr) {
      inst.emitError("missing `")
          << member << "` attribute in `" << anno.getClass() << "` annotation";
      anyFailures = true;
    }
    return attr;
  };
 
  if (anno.isClass(extractBlackBoxAnnoClass)) {
    auto filename = getStringOrError("filename");
    auto prefix = getStringOrError("prefix");
    auto dest = anno.getMember<StringAttr>("dest"); // optional
    if (anyFailures)
      return;
 
    ExtractionInfo info;
    info.traceFilename = filename;
    info.prefix = prefix;
    info.wrapperModule = (dest ? dest.getValue() : "");
 
    // CAVEAT: If the instance has a wrapper module configured then extraction
    // should stop at the DUT module instead of extracting past the DUT into the
    // surrounding test harness. This is all very ugly and hacky.
    info.stopAtDUT = !info.wrapperModule.empty();
 
    extractionWorklist.push_back({inst, info});
    return;
  }
}
 
/// Find the location in an NLA that corresponds to a given instance (either by
/// mentioning exactly the instance, or the instance's parent module). Returns a
/// position within the NLA's path, or the length of the path if the instances
/// was not found.
static unsigned findInstanceInNLA(InstanceOp inst, hw::HierPathOp nla) {
  unsigned nlaLen = nla.getNamepath().size();
  auto instName = getInnerSymName(inst);
  auto parentName = cast<FModuleOp>(inst->getParentOp()).getModuleNameAttr();
  for (unsigned nlaIdx = 0; nlaIdx < nlaLen; ++nlaIdx) {
    auto refPart = nla.refPart(nlaIdx);
    if (nla.modPart(nlaIdx) == parentName && (!refPart || refPart == instName))
      return nlaIdx;
  }
  return nlaLen;
}
 
/// Move instances in the extraction worklist upwards in the hierarchy. This
/// iteratively pushes instances up one level of hierarchy until they have
/// arrived in the desired container module.
void ExtractInstancesPass::extractInstances() {
  // The list of ports to be added to an instance's parent module. Cleared and
  // reused across instances.
  SmallVector<std::pair<unsigned, PortInfo>> newPorts;
  // The number of instances with the same prefix. Used to uniquify prefices.
  DenseMap<StringRef, unsigned> prefixUniqueIDs;
 
  SmallPtrSet<Operation *, 4> nlasToRemove;
 
  auto &nlaTable = getAnalysis<NLATable>();
 
  // Keep track of where the instance was originally.
  for (auto &[inst, info] : extractionWorklist)
    originalInstanceParents[inst] =
        inst->getParentOfType<FModuleLike>().getModuleNameAttr();
 
  while (!extractionWorklist.empty()) {
    InstanceOp inst;
    ExtractionInfo info;
    std::tie(inst, info) = extractionWorklist.pop_back_val();
    auto parent = inst->getParentOfType<FModuleOp>();
 
    // Figure out the wiring prefix to use for this instance. If we are supposed
    // to use a wiring prefix (`info.prefix` is non-empty), we assemble a
    // `<prefix>_<N>` string, where `N` is an unsigned integer used to uniquifiy
    // the prefix. This is very close to what the original Scala implementation
    // of the pass does, which would group instances to be extracted by prefix
    // and then iterate over them with the index in the group being used as `N`.
    StringRef prefix;
    if (!info.prefix.empty()) {
      auto &prefixSlot = instPrefices[inst];
      if (prefixSlot.empty()) {
        auto idx = prefixUniqueIDs[info.prefix]++;
        (Twine(info.prefix) + "_" + Twine(idx)).toVector(prefixSlot);
      }
      prefix = prefixSlot;
    }
 
    // If the instance is already in the right place (outside the DUT or already
    // in the root module), there's nothing left for us to do. Otherwise we
    // proceed to bubble it up one level in the hierarchy and add the resulting
    // instances back to the worklist.
    if (!dutModules.contains(parent) ||
        instanceGraph->lookup(parent)->noUses() ||
        (info.stopAtDUT && dutRootModules.contains(parent))) {
      LLVM_DEBUG(llvm::dbgs() << "\nNo need to further move " << inst << "\n");
      extractedInstances.push_back({inst, info});
      continue;
    }
    LLVM_DEBUG({
      llvm::dbgs() << "\nMoving ";
      if (!prefix.empty())
        llvm::dbgs() << "`" << prefix << "` ";
      llvm::dbgs() << inst << "\n";
    });
 
    // Add additional ports to the parent module as a replacement for the
    // instance port signals once the instance is extracted.
    unsigned numParentPorts = parent.getNumPorts();
    unsigned numInstPorts = inst.getNumResults();
 
    for (unsigned portIdx = 0; portIdx < numInstPorts; ++portIdx) {
      // Assemble the new port name as "<prefix>_<name>", where the prefix is
      // provided by the extraction annotation.
      auto name = inst.getPortNameStr(portIdx);
      auto nameAttr = StringAttr::get(
          &getContext(),
          prefix.empty() ? Twine(name) : Twine(prefix) + "_" + name);
 
      PortInfo newPort{nameAttr,
                       type_cast<FIRRTLType>(inst.getResult(portIdx).getType()),
                       direction::flip(inst.getPortDirection(portIdx))};
      newPort.loc = inst.getResult(portIdx).getLoc();
      newPorts.push_back({numParentPorts, newPort});
      LLVM_DEBUG(llvm::dbgs()
                 << "- Adding port " << newPort.direction << " "
                 << newPort.name.getValue() << ": " << newPort.type << "\n");
    }
    parent.insertPorts(newPorts);
    anythingChanged = true;
 
    // Replace all uses of the existing instance ports with the newly-created
    // module ports.
    for (unsigned portIdx = 0; portIdx < numInstPorts; ++portIdx) {
      inst.getResult(portIdx).replaceAllUsesWith(
          parent.getArgument(numParentPorts + portIdx));
    }
    assert(inst.use_empty() && "instance ports should have been detached");
    DenseSet<hw::HierPathOp> instanceNLAs;
    // Get the NLAs that pass through the InstanceOp `inst`.
    // This does not returns NLAs that have the `inst` as the leaf.
    nlaTable.getInstanceNLAs(inst, instanceNLAs);
    // Map of the NLAs, that are applied to the InstanceOp. That is the NLA
    // terminates on the InstanceOp.
    DenseMap<hw::HierPathOp, SmallVector<Annotation>> instNonlocalAnnos;
    AnnotationSet::removeAnnotations(inst, [&](Annotation anno) {
      // Only consider annotations with a `circt.nonlocal` field.
      auto nlaName = anno.getMember<FlatSymbolRefAttr>("circt.nonlocal");
      if (!nlaName)
        return false;
      // Track the NLA.
      if (hw::HierPathOp nla = nlaTable.getNLA(nlaName.getAttr())) {
        instNonlocalAnnos[nla].push_back(anno);
        instanceNLAs.insert(nla);
      }
      return true;
    });
 
    // Sort the instance NLAs we've collected by the NLA name to have a
    // deterministic output.
    SmallVector<hw::HierPathOp> sortedInstanceNLAs(instanceNLAs.begin(),
                                                   instanceNLAs.end());
    llvm::sort(sortedInstanceNLAs,
               [](auto a, auto b) { return a.getSymName() < b.getSymName(); });
 
    // Move the original instance one level up such that it is right next to
    // the instances of the parent module, and wire the instance ports up to
    // the newly added parent module ports.
    auto *instParentNode =
        instanceGraph->lookup(cast<igraph::ModuleOpInterface>(*parent));
    for (auto *instRecord : instParentNode->uses()) {
      auto oldParentInst = cast<InstanceOp>(*instRecord->getInstance());
      auto newParent = oldParentInst->getParentOfType<FModuleLike>();
      LLVM_DEBUG(llvm::dbgs() << "- Updating " << oldParentInst << "\n");
      auto newParentInst = oldParentInst.cloneAndInsertPorts(newPorts);
 
      // Migrate connections to existing ports.
      for (unsigned portIdx = 0; portIdx < numParentPorts; ++portIdx)
        oldParentInst.getResult(portIdx).replaceAllUsesWith(
            newParentInst.getResult(portIdx));
 
      // Clone the existing instance and remove it from its current parent, such
      // that we can insert it at its extracted location.
      auto newInst = inst.cloneAndInsertPorts({});
      newInst->remove();
 
      // Ensure that the `inner_sym` of the instance is unique within the parent
      // module we're extracting it to.
      if (auto instSym = getInnerSymName(inst)) {
        auto newName =
            getModuleNamespace(newParent).newName(instSym.getValue());
        if (newName != instSym.getValue())
          newInst.setInnerSymAttr(
              hw::InnerSymAttr::get(StringAttr::get(&getContext(), newName)));
      }
 
      // Add the moved instance and hook it up to the added ports.
      ImplicitLocOpBuilder builder(inst.getLoc(), newParentInst);
      builder.setInsertionPointAfter(newParentInst);
      builder.insert(newInst);
      for (unsigned portIdx = 0; portIdx < numInstPorts; ++portIdx) {
        auto dst = newInst.getResult(portIdx);
        auto src = newParentInst.getResult(numParentPorts + portIdx);
        if (newPorts[portIdx].second.direction == Direction::In)
          std::swap(src, dst);
        builder.create<MatchingConnectOp>(dst, src);
      }
 
      // Move the wiring prefix from the old to the new instance. We just look
      // up the prefix for the old instance and if it exists, we remove it and
      // assign it to the new instance. This has the effect of making the first
      // new instance we create inherit the wiring prefix, and all additional
      // new instances (e.g. through multiple instantiation of the parent) will
      // pick a new prefix.
      auto oldPrefix = instPrefices.find(inst);
      if (oldPrefix != instPrefices.end()) {
        LLVM_DEBUG(llvm::dbgs()
                   << "  - Reusing prefix `" << oldPrefix->second << "`\n");
        auto newPrefix = std::move(oldPrefix->second);
        instPrefices.erase(oldPrefix);
        instPrefices.insert({newInst, newPrefix});
      }
 
      // Inherit the old instance's extraction path.
      extractionPaths.try_emplace(newInst); // (create entry first)
      auto &extractionPath = (extractionPaths[newInst] = extractionPaths[inst]);
      extractionPath.push_back(getInnerRefTo(newParentInst));
      originalInstanceParents.try_emplace(newInst); // (create entry first)
      originalInstanceParents[newInst] = originalInstanceParents[inst];
      // Record the Nonlocal annotations that need to be applied to the new
      // Inst.
      SmallVector<Annotation> newInstNonlocalAnnos;
 
      // Update all NLAs that touch the moved instance.
      for (auto nla : sortedInstanceNLAs) {
        LLVM_DEBUG(llvm::dbgs() << "  - Updating " << nla << "\n");
 
        // Find the position of the instance in the NLA path. This is going to
        // be the position at which we have to modify the NLA.
        SmallVector<Attribute> nlaPath(nla.getNamepath().begin(),
                                       nla.getNamepath().end());
        unsigned nlaIdx = findInstanceInNLA(inst, nla);
 
        // Handle the case where the instance no longer shows up in the NLA's
        // path. This usually happens if the instance is extracted into multiple
        // parents (because the current parent module is multiply instantiated).
        // In that case NLAs that were specific to one instance may have been
        // moved when we arrive at the second instance, and the NLA is already
        // updated.
        if (nlaIdx >= nlaPath.size()) {
          LLVM_DEBUG(llvm::dbgs() << "    - Instance no longer in path\n");
          continue;
        }
        LLVM_DEBUG(llvm::dbgs() << "    - Position " << nlaIdx << "\n");
 
        // Handle the case where the NLA's path doesn't go through the
        // instance's new parent module, which happens if the current parent
        // module is multiply instantiated. In that case, we only move over NLAs
        // that actually affect the instance through the new parent module.
        if (nlaIdx > 0) {
          auto innerRef = dyn_cast<InnerRefAttr>(nlaPath[nlaIdx - 1]);
          if (innerRef &&
              !(innerRef.getModule() == newParent.getModuleNameAttr() &&
                innerRef.getName() == getInnerSymName(newParentInst))) {
            LLVM_DEBUG(llvm::dbgs()
                       << "    - Ignored since NLA parent " << innerRef
                       << " does not pass through extraction parent\n");
            continue;
          }
        }
 
        // There are two interesting cases now:
        // - If `nlaIdx == 0`, the NLA is rooted at the module the instance was
        //   located in prior to extraction. This indicates that the NLA applies
        //   to all instances of that parent module. Since we are extracting
        //   *out* of that module, we have to create a new NLA rooted at the new
        //   parent module after extraction.
        // - If `nlaIdx > 0`, the NLA is rooted further up in the hierarchy and
        //   we can simply remove the old parent module from the path.
 
        // Handle the case where we need to come up with a new NLA for this
        // instance since we've moved it past the module at which the old NLA
        // was rooted at.
        if (nlaIdx == 0) {
          LLVM_DEBUG(llvm::dbgs() << "    - Re-rooting " << nlaPath[0] << "\n");
          assert(isa<InnerRefAttr>(nlaPath[0]) &&
                 "head of hierpath must be an InnerRefAttr");
          nlaPath[0] = InnerRefAttr::get(newParent.getModuleNameAttr(),
                                         getInnerSymName(newInst));
 
          if (instParentNode->hasOneUse()) {
            // Simply update the existing NLA since our parent is only
            // instantiated once, and we therefore are not creating multiple
            // instances through the extraction.
            nlaTable.erase(nla);
            nla.setNamepathAttr(builder.getArrayAttr(nlaPath));
            for (auto anno : instNonlocalAnnos.lookup(nla))
              newInstNonlocalAnnos.push_back(anno);
            nlaTable.addNLA(nla);
            LLVM_DEBUG(llvm::dbgs() << "    - Modified to " << nla << "\n");
          } else {
            // Since we are extracting to multiple parent locations, create a
            // new NLA for each instantiation site.
            auto newNla = cloneWithNewNameAndPath(nla, nlaPath);
            for (auto anno : instNonlocalAnnos.lookup(nla)) {
              anno.setMember("circt.nonlocal",
                             FlatSymbolRefAttr::get(newNla.getSymNameAttr()));
              newInstNonlocalAnnos.push_back(anno);
            }
 
            nlaTable.addNLA(newNla);
            LLVM_DEBUG(llvm::dbgs() << "    - Created " << newNla << "\n");
            // CAVEAT(fschuiki): This results in annotations in the subhierarchy
            // below `inst` with the old NLA symbol name, instead of those
            // annotations duplicated for each of the newly-created NLAs. This
            // shouldn't come up in our current use cases, but is a weakness of
            // the current implementation. Instead, we should keep an NLA
            // replication table that we fill with mappings from old NLA names
            // to lists of new NLA names. A post-pass would then traverse the
            // entire subhierarchy and go replicate all annotations with the old
            // names.
            inst.emitWarning("extraction of instance `")
                << inst.getInstanceName()
                << "` could break non-local annotations rooted at `"
                << parent.getModuleName() << "`";
          }
          continue;
        }
 
        // In the subequent code block we are going to remove one element from
        // the NLA path, corresponding to the fact that the extracted instance
        // has moved up in the hierarchy by one level. Removing that element may
        // leave the NLA in a degenerate state, with only a single element in
        // its path. If that is the case we have to convert the NLA into a
        // regular local annotation.
        if (nlaPath.size() == 2) {
          for (auto anno : instNonlocalAnnos.lookup(nla)) {
            anno.removeMember("circt.nonlocal");
            newInstNonlocalAnnos.push_back(anno);
            LLVM_DEBUG(llvm::dbgs() << "    - Converted to local "
                                    << anno.getDict() << "\n");
          }
          nlaTable.erase(nla);
          nlasToRemove.insert(nla);
          continue;
        }
 
        // At this point the NLA looks like `NewParent::X, OldParent::BB`, and
        // the `nlaIdx` points at `OldParent::BB`. To make our lives easier,
        // since we know that `nlaIdx` is a `InnerRefAttr`, we'll modify
        // `OldParent::BB` to be `NewParent::BB` and delete `NewParent::X`.
        StringAttr parentName =
            cast<InnerRefAttr>(nlaPath[nlaIdx - 1]).getModule();
        Attribute newRef;
        if (isa<InnerRefAttr>(nlaPath[nlaIdx]))
          newRef = InnerRefAttr::get(parentName, getInnerSymName(newInst));
        else
          newRef = FlatSymbolRefAttr::get(parentName);
        LLVM_DEBUG(llvm::dbgs()
                   << "    - Replacing " << nlaPath[nlaIdx - 1] << " and "
                   << nlaPath[nlaIdx] << " with " << newRef << "\n");
        nlaPath[nlaIdx] = newRef;
        nlaPath.erase(nlaPath.begin() + nlaIdx - 1);
 
        if (isa<FlatSymbolRefAttr>(newRef)) {
          // Since the original NLA ended at the instance's parent module, there
          // is no guarantee that the instance is the sole user of the NLA (as
          // opposed to the original NLA explicitly naming the instance). Create
          // a new NLA.
          auto newNla = cloneWithNewNameAndPath(nla, nlaPath);
          nlaTable.addNLA(newNla);
          LLVM_DEBUG(llvm::dbgs() << "    - Created " << newNla << "\n");
          for (auto anno : instNonlocalAnnos.lookup(nla)) {
            anno.setMember("circt.nonlocal",
                           FlatSymbolRefAttr::get(newNla.getSymNameAttr()));
            newInstNonlocalAnnos.push_back(anno);
          }
        } else {
          nla.setNamepathAttr(builder.getArrayAttr(nlaPath));
          LLVM_DEBUG(llvm::dbgs() << "    - Modified to " << nla << "\n");
          for (auto anno : instNonlocalAnnos.lookup(nla))
            newInstNonlocalAnnos.push_back(anno);
        }
 
        // No update to NLATable required, since it will be deleted from the
        // parent, and it should already exist in the new parent module.
        continue;
      }
      AnnotationSet newInstAnnos(newInst);
      newInstAnnos.addAnnotations(newInstNonlocalAnnos);
      newInstAnnos.applyToOperation(newInst);
 
      // Add the moved instance to the extraction worklist such that it gets
      // bubbled up further if needed.
      extractionWorklist.push_back({newInst, info});
      LLVM_DEBUG(llvm::dbgs() << "  - Updated to " << newInst << "\n");
 
      // Keep instance graph up-to-date.
      instanceGraph->replaceInstance(oldParentInst, newParentInst);
      oldParentInst.erase();
    }
    // Remove the obsolete NLAs from the instance of the parent module, since
    // the extracted instance no longer resides in that module and any NLAs to
    // it no longer go through the parent module.
    nlaTable.removeNLAsfromModule(instanceNLAs, parent.getNameAttr());
 
    // Clean up the original instance.
    inst.erase();
    newPorts.clear();
  }
 
  // Remove unused NLAs.
  for (Operation *op : nlasToRemove) {
    LLVM_DEBUG(llvm::dbgs() << "Removing obsolete " << *op << "\n");
    op->erase();
  }
}
 
/// Group instances into submodules after they have been moved upwards. This
/// only occurs for instances that had the corresponding `dest` field of the
/// annotation set.
void ExtractInstancesPass::groupInstances() {
  // Group the extracted instances by their wrapper module name and their parent
  // module. Note that we cannot group instances that landed in different parent
  // modules into the same submodule, so we use that parent module as a grouping
  // key.
  llvm::MapVector<std::pair<Operation *, StringRef>, SmallVector<InstanceOp>>
      instsByWrapper;
  for (auto &[inst, info] : extractedInstances) {
    if (!info.wrapperModule.empty())
      instsByWrapper[{inst->getParentOfType<FModuleOp>(), info.wrapperModule}]
          .push_back(inst);
  }
  if (instsByWrapper.empty())
    return;
  LLVM_DEBUG(llvm::dbgs() << "\nGrouping instances into wrappers\n");
 
  // Generate the wrappers.
  SmallVector<PortInfo> ports;
  auto &nlaTable = getAnalysis<NLATable>();
 
  for (auto &[parentAndWrapperName, insts] : instsByWrapper) {
    auto [parentOp, wrapperName] = parentAndWrapperName;
    auto parent = cast<FModuleOp>(parentOp);
    LLVM_DEBUG(llvm::dbgs() << "- Wrapper `" << wrapperName << "` in `"
                            << parent.getModuleName() << "` with "
                            << insts.size() << " instances\n");
    OpBuilder builder(parentOp);
 
    // Uniquify the wrapper name.
    auto wrapperModuleName = builder.getStringAttr(
        circuitNamespace.newName(dutPrefix + wrapperName));
    auto wrapperInstName =
        builder.getStringAttr(getModuleNamespace(parent).newName(wrapperName));
 
    // Assemble a list of ports for the wrapper module, which is basically just
    // a concatenation of the wrapped instance ports. Also keep track of the
    // NLAs that target the grouped instances since these will have to pass
    // through the wrapper module.
    ports.clear();
    for (auto inst : insts) {
      // Determine the ports for the wrapper.
      StringRef prefix(instPrefices[inst]);
      unsigned portNum = inst.getNumResults();
      for (unsigned portIdx = 0; portIdx < portNum; ++portIdx) {
        auto name = inst.getPortNameStr(portIdx);
        auto nameAttr = builder.getStringAttr(
            prefix.empty() ? Twine(name) : Twine(prefix) + "_" + name);
        PortInfo port{nameAttr,
                      type_cast<FIRRTLType>(inst.getResult(portIdx).getType()),
                      inst.getPortDirection(portIdx)};
        port.loc = inst.getResult(portIdx).getLoc();
        ports.push_back(port);
      }
 
      // Set of NLAs that have a reference to this InstanceOp `inst`.
      DenseSet<hw::HierPathOp> instNlas;
      // Get the NLAs that pass through the `inst`, and not end at it.
      nlaTable.getInstanceNLAs(inst, instNlas);
      AnnotationSet instAnnos(inst);
      // Get the NLAs that end at the InstanceOp, that is the Nonlocal
      // annotations that apply to the InstanceOp.
      for (auto anno : instAnnos) {
        auto nlaName = anno.getMember<FlatSymbolRefAttr>("circt.nonlocal");
        if (!nlaName)
          continue;
        hw::HierPathOp nla = nlaTable.getNLA(nlaName.getAttr());
        if (nla)
          instNlas.insert(nla);
      }
      for (auto nla : instNlas) {
        LLVM_DEBUG(llvm::dbgs() << "  - Updating " << nla << "\n");
 
        // Find the position of the instance in the NLA path. This is going to
        // be the position at which we have to modify the NLA.
        SmallVector<Attribute> nlaPath(nla.getNamepath().begin(),
                                       nla.getNamepath().end());
        unsigned nlaIdx = findInstanceInNLA(inst, nla);
        assert(nlaIdx < nlaPath.size() && "instance not found in its own NLA");
        LLVM_DEBUG(llvm::dbgs() << "    - Position " << nlaIdx << "\n");
 
        // The relevant part of the NLA is of the form `Top::bb`, which we want
        // to expand to `Top::wrapperInst` and `Wrapper::bb`.
        auto ref1 =
            InnerRefAttr::get(parent.getModuleNameAttr(), wrapperInstName);
        Attribute ref2;
        if (auto innerRef = dyn_cast<InnerRefAttr>(nlaPath[nlaIdx]))
          ref2 = InnerRefAttr::get(wrapperModuleName, innerRef.getName());
        else
          ref2 = FlatSymbolRefAttr::get(wrapperModuleName);
        LLVM_DEBUG(llvm::dbgs() << "    - Expanding " << nlaPath[nlaIdx]
                                << " to (" << ref1 << ", " << ref2 << ")\n");
        nlaPath[nlaIdx] = ref1;
        nlaPath.insert(nlaPath.begin() + nlaIdx + 1, ref2);
        // CAVEAT: This is likely to conflict with additional users of `nla`
        // that have nothing to do with this instance. Might need some NLATable
        // machinery at some point to allow for these things to be updated.
        nla.setNamepathAttr(builder.getArrayAttr(nlaPath));
        LLVM_DEBUG(llvm::dbgs() << "    - Modified to " << nla << "\n");
        // Add the NLA to the wrapper module.
        nlaTable.addNLAtoModule(nla, wrapperModuleName);
      }
    }
 
    // Create the wrapper module.
    auto wrapper = builder.create<FModuleOp>(
        builder.getUnknownLoc(), wrapperModuleName,
        ConventionAttr::get(builder.getContext(), Convention::Internal), ports);
    SymbolTable::setSymbolVisibility(wrapper, SymbolTable::Visibility::Private);
 
    // Instantiate the wrapper module in the parent and replace uses of the
    // extracted instances' ports with the corresponding wrapper module ports.
    // This will essentially disconnect the extracted instances.
    builder.setInsertionPointToStart(parent.getBodyBlock());
    auto wrapperInst = builder.create<InstanceOp>(
        wrapper.getLoc(), wrapper, wrapperName, NameKindEnum::DroppableName,
        ArrayRef<Attribute>{},
        /*portAnnotations=*/ArrayRef<Attribute>{}, /*lowerToBind=*/false,
        hw::InnerSymAttr::get(wrapperInstName));
    unsigned portIdx = 0;
    for (auto inst : insts)
      for (auto result : inst.getResults())
        result.replaceAllUsesWith(wrapperInst.getResult(portIdx++));
 
    // Move all instances into the wrapper module and wire them up to the
    // wrapper ports.
    portIdx = 0;
    builder.setInsertionPointToStart(wrapper.getBodyBlock());
    for (auto inst : insts) {
      inst->remove();
      builder.insert(inst);
      for (auto result : inst.getResults()) {
        Value dst = result;
        Value src = wrapper.getArgument(portIdx);
        if (ports[portIdx].direction == Direction::Out)
          std::swap(dst, src);
        builder.create<MatchingConnectOp>(result.getLoc(), dst, src);
        ++portIdx;
      }
    }
  }
}
 
/// Generate trace files, which are plain text metadata files that list the
/// hierarchical path where each instance was extracted from. The file lists one
/// instance per line in the form `<prefix> -> <original-path>`.
void ExtractInstancesPass::createTraceFiles() {
  LLVM_DEBUG(llvm::dbgs() << "\nGenerating trace files\n");
 
  // Group the extracted instances by their trace file name.
  SmallDenseMap<StringRef, SmallVector<InstanceOp>> instsByTraceFile;
  for (auto &[inst, info] : extractedInstances)
    if (!info.traceFilename.empty())
      instsByTraceFile[info.traceFilename].push_back(inst);
 
  // Generate the trace files.
  SmallVector<Attribute> symbols;
  SmallDenseMap<Attribute, unsigned> symbolIndices;
 
  for (auto &[fileName, insts] : instsByTraceFile) {
    LLVM_DEBUG(llvm::dbgs() << "- " << fileName << "\n");
    std::string buffer;
    llvm::raw_string_ostream os(buffer);
    symbols.clear();
    symbolIndices.clear();
 
    auto addSymbol = [&](Attribute symbol) {
      unsigned id;
      auto it = symbolIndices.find(symbol);
      if (it != symbolIndices.end()) {
        id = it->second;
      } else {
        id = symbols.size();
        symbols.push_back(symbol);
        symbolIndices.insert({symbol, id});
      }
      os << "{{" << id << "}}";
    };
 
    auto file = getOrCreateFile(fileName);
    auto builder = OpBuilder::atBlockEnd(file.getBody());
    for (auto inst : insts) {
      StringRef prefix(instPrefices[inst]);
      if (prefix.empty()) {
        LLVM_DEBUG(llvm::dbgs() << "  - Skipping `" << inst.getName()
                                << "` since it has no extraction prefix\n");
        continue;
      }
      ArrayRef<InnerRefAttr> path(extractionPaths[inst]);
      if (path.empty()) {
        LLVM_DEBUG(llvm::dbgs() << "  - Skipping `" << inst.getName()
                                << "` since it has not been moved\n");
        continue;
      }
      LLVM_DEBUG(llvm::dbgs()
                 << "  - " << prefix << ": " << inst.getName() << "\n");
      os << prefix << " -> ";
 
      // HACK: To match the Scala implementation, we strip all non-DUT modules
      // from the path and make the path look like it's rooted at the first DUT
      // module (so `TestHarness.dut.foo.bar` becomes `DUTModule.foo.bar`).
      while (!path.empty() &&
             !dutModuleNames.contains(path.back().getModule())) {
        LLVM_DEBUG(llvm::dbgs()
                   << "    - Dropping non-DUT segment " << path.back() << "\n");
        path = path.drop_back();
      }
      // HACK: This is extremely ugly. In case the instance was just moved by a
      // single level, the path may become empty. In that case we simply use the
      // instance's original parent before it was moved.
      addSymbol(FlatSymbolRefAttr::get(path.empty()
                                           ? originalInstanceParents[inst]
                                           : path.back().getModule()));
      for (auto sym : llvm::reverse(path)) {
        os << ".";
        addSymbol(sym);
      }
      // The final instance name is excluded as this does not provide useful
      // additional information and could conflict with a name inside the final
      // module.
      os << "\n";
    }
 
    // Put the information in a verbatim operation.
    builder.create<sv::VerbatimOp>(builder.getUnknownLoc(), buffer,
                                   ValueRange{}, builder.getArrayAttr(symbols));
  }
}
 
//===----------------------------------------------------------------------===//
// Pass Creation
//===----------------------------------------------------------------------===//
 
std::unique_ptr<mlir::Pass> circt::firrtl::createExtractInstancesPass() {
  return std::make_unique<ExtractInstancesPass>();
}