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.
//
// This pass will make no modifications if the circuit does not contain a
// design-under-test (DUT).  I.e., this pass does not use the "effecctive" DUT.
// If a DUT exists, then anything in the design is extracted.  Using the
// standard interpretation of passes like this, layers are not in the design.
// If a situation arise where a module is instantiated inside and outside the
// design that needs to be extracted, then it will be extracted in both up to
// the point where it no longer needs to be further extracted.  See the tests
// for examples of this.
//
//===----------------------------------------------------------------------===//
 
#include "circt/Analysis/FIRRTLInstanceInfo.h"
#include "circt/Dialect/Emit/EmitOps.h"
#include "circt/Dialect/FIRRTL/AnnotationDetails.h"
#include "circt/Dialect/FIRRTL/FIRRTLAnnotationHelper.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/InnerSymbolNamespace.h"
#include "circt/Dialect/SV/SVOps.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/ImplicitLocOpBuilder.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/FileUtilities.h"
#include "llvm/Support/Debug.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;
  /// Optional name of the instance of the wrapper module.
  StringRef wrapperInstanceName;
  /// 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(ClassOp &sifiveMetadata);
  void createSchema();
 
  /// 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.getName())));
    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 = emit::FileOp::create(builder, fileName);
    }
    return it->second;
  }
 
  bool anythingChanged;
  bool anyFailures;
 
  CircuitOp circuitOp;
  InstanceGraph *instanceGraph = nullptr;
  InstanceInfo *instanceInfo = nullptr;
  SymbolTable *symbolTable = nullptr;
 
  /// The modules in the design that are annotated with one or more annotations
  /// relevant for instance extraction.
  DenseMap<Operation *, SmallVector<Annotation, 1>> annotatedModules;
 
  /// 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 and original name for each instance.
  DenseMap<Operation *, std::pair<SmallString<16>, StringAttr>>
      instPrefixNamesPair;
 
  /// The current circuit namespace valid within the call to `runOnOperation`.
  CircuitNamespace circuitNamespace;
  /// Cached module namespaces.
  DenseMap<Operation *, hw::InnerSymbolNamespace> moduleNamespaces;
  /// The metadata class ops.
  ClassOp extractMetadataClass, schemaClass;
  const unsigned prefixNameFieldId = 0, pathFieldId = 2, fileNameFieldId = 4,
                 instNameFieldId = 6;
  /// Cache of the inner ref to the new instances created. Will be used to
  /// create a path to the instance
  DenseMap<InnerRefAttr, InstanceOp> innerRefToInstances;
  Type stringType, pathType;
 
  /// If set, this indicates that the `InjectDUTHierarchy` pass ran with the
  /// `moveDut` parameter enabled.  If true, then extraction continues outside
  /// the DUT even when extraction annotations have a wrapper module specified.
  bool moveDut;
};
} // end anonymous namespace
 
/// Emit the annotated source code for black boxes in a circuit.
void ExtractInstancesPass::runOnOperation() {
  circuitOp = getOperation();
  anythingChanged = false;
  anyFailures = false;
  annotatedModules.clear();
  extractionWorklist.clear();
  files.clear();
  extractionPaths.clear();
  originalInstanceParents.clear();
  extractedInstances.clear();
  instPrefixNamesPair.clear();
  moduleNamespaces.clear();
  circuitNamespace.clear();
  circuitNamespace.add(circuitOp);
  innerRefToInstances.clear();
  extractMetadataClass = {};
  schemaClass = {};
  auto *context = circuitOp->getContext();
  stringType = StringType::get(context);
  pathType = PathType::get(context);
  moveDut = false;
 
  // Walk the IR and gather all the annotations relevant for extraction that
  // appear on instances and the instantiated modules.
  instanceGraph = &getAnalysis<InstanceGraph>();
  instanceInfo = &getAnalysis<InstanceInfo>();
  symbolTable = &getAnalysis<SymbolTable>();
  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();
 
  ClassOp sifiveMetadata =
      dyn_cast_or_null<ClassOp>(symbolTable->lookup("SiFive_Metadata"));
 
  // Generate the trace files that list where each instance was extracted from.
  createTraceFiles(sifiveMetadata);
  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();
 
  // Find an optional `InjectDUTHierarchyAnnotation`.  If it exists, inspect the
  // `moveDut` field.  If this is `true`, then we moved the DUT from the
  // original DUT to the wrapper.  If this occurred, then this affecst the
  // behavior of whether or not we stop at the DUT (now the wrapper) when we
  // extract instances.
  //
  // TODO: This is tech debt.  This was accepted on condition that work is done
  // to remove this pass.
  AnnotationSet::removeAnnotations(circuit, [&](Annotation anno) {
    if (!anno.isClass(injectDUTHierarchyAnnoClass))
      return false;
 
    if (auto moveDutAnnoAttr = anno.getMember<BoolAttr>("moveDut"))
      moveDut = moveDutAnnoAttr.getValue();
    return true;
  });
 
  // Grab the clock gate extraction annotation on the circuit.
  StringRef clkgateFileName;
  StringRef clkgateWrapperModule;
  StringRef clkgateWrapperInstanceName;
  AnnotationSet::removeAnnotations(circuit, [&](Annotation anno) {
    if (!anno.isClass(extractClockGatesFileAnnoClass))
      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");
    auto instanceNameAttr = anno.getMember<StringAttr>("instanceName");
    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();
    if (instanceNameAttr)
      clkgateWrapperInstanceName = instanceNameAttr.getValue();
    return true;
  });
 
  // Grab the memory extraction annotation on the circuit.
  StringRef memoryFileName;
  StringRef memoryWrapperModule;
  StringRef memoryWrapperInstanceName;
  AnnotationSet::removeAnnotations(circuit, [&](Annotation anno) {
    if (!anno.isClass(extractSeqMemsFileAnnoClass))
      return false;
    LLVM_DEBUG(llvm::dbgs()
               << "Memory extraction config: " << anno.getDict() << "\n");
    auto filenameAttr = anno.getMember<StringAttr>("filename");
    auto groupAttr = anno.getMember<StringAttr>("group");
    auto instanceNameAttr = anno.getMember<StringAttr>("instanceName");
    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();
    if (instanceNameAttr)
      memoryWrapperInstanceName = instanceNameAttr.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 (!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]);
  });
 
  // If clock gate extraction is requested, find instances of extmodules which
  // have a defname that ends with "EICG_wrapper".  This also allows this to
  // compose with Chisel-time module prefixing.
  //
  // TODO: This defname matching is a terrible hack and should be replaced with
  // something better.
  if (!clkgateFileName.empty()) {
    for (auto module : circuit.getOps<FExtModuleOp>()) {
      if (!module.getDefnameAttr().getValue().ends_with("EICG_wrapper"))
        continue;
      LLVM_DEBUG(llvm::dbgs()
                 << "Clock gate `" << module.getModuleName() << "`\n");
      if (!instanceInfo->anyInstanceInDesign(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.wrapperInstanceName = clkgateWrapperInstanceName;
      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});
        } else {
          instRecord->getInstance()->emitError()
              << "cannot extract clock gate instances through non-InstanceOp";
          anyFailures = true;
        }
      }
    }
  }
 
  // 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 (!instanceInfo->anyInstanceInDesign(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.wrapperInstanceName = memoryWrapperInstanceName;
      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});
        } else {
          instRecord->getInstance()->emitError()
              << "cannot extract memory instances through non-InstanceOp";
          anyFailures = true;
        }
      }
    }
  }
}
 
/// 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.
 
    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;
  // Track the index used to unique a given wiring prefix (e.g., "mem_wiring")
  // in a given module.  This needs to be isolated per-module to avoid causing
  // pollution of the namespace across the circuit.
  DenseMap<std::pair<Operation *, 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. The prefix is recomputed at every level of the hierarchy the
    // instance is bubbled up through, with `N` counting from zero within the
    // module the instance is currently in. This keeps the names of the ports
    // added to a module independent of the surrounding circuit.
    StringRef prefix;
    auto &instPrefixEntry = instPrefixNamesPair[inst];
    instPrefixEntry.second = inst.getInstanceNameAttr();
    if (!info.prefix.empty()) {
      auto &prefixSlot = instPrefixEntry.first;
      prefixSlot.clear();
      auto idx = prefixUniqueIDs[{parent, info.prefix}]++;
      (Twine(info.prefix) + "_" + Twine(idx)).toVector(prefixSlot);
      prefix = prefixSlot;
    }
 
    /// Return true if this extraction should stop at the DUT or if it should
    /// continue beyond it.
    bool stopAtDUT = !moveDut && !info.wrapperModule.empty();
 
    // If the instance is already in the right place (outside the DUT, already
    // in the root module, or has hit a layer), 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 (inst->getParentOfType<LayerBlockOp>() ||
        !instanceInfo->anyInstanceInDesign(parent) ||
        instanceGraph->lookup(parent)->noUses() ||
        (stopAtDUT && instanceInfo->isDut(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.getPortName(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.getName() < b.getName(); });
 
    // 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 = dyn_cast<InstanceOp>(*instRecord->getInstance());
      if (!oldParentInst) {
        inst.emitError("cannot extract instance `")
            << inst.getName() << "` through a non-InstanceOp parent";
        anyFailures = true;
        continue;
      }
      auto newParent = oldParentInst->getParentOfType<FModuleLike>();
      LLVM_DEBUG(llvm::dbgs() << "- Updating " << oldParentInst << "\n");
      auto newParentInst = cast<InstanceOp>(
          oldParentInst.cloneWithInsertedPortsAndReplaceUses(newPorts));
      if (newParentInst.getInnerSymAttr())
        innerRefToInstances[getInnerRefTo(newParentInst)] = newParentInst;
 
      auto newInst = cast<InstanceOp>(inst->clone());
 
      // 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);
      if (newParentInst.getInnerSymAttr())
        innerRefToInstances[getInnerRefTo(newInst)] = 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);
        MatchingConnectOp::create(builder, dst, src);
      }
 
      // Move the wiring prefix bookkeeping from the old to the new instance.
      // The prefix itself is recomputed once the new instance is popped off the
      // worklist, since the numbering is local to the module the instance is
      // being extracted out of. This mainly serves to keep the map from growing
      // stale entries for instances that are about to be erased.
      auto oldPrefix = instPrefixNamesPair.find(inst);
      if (oldPrefix != instPrefixNamesPair.end()) {
        LLVM_DEBUG(llvm::dbgs() << "  - Moving prefix `"
                                << oldPrefix->second.first << "`\n");
        auto newPrefix = std::move(oldPrefix->second);
        instPrefixNamesPair.erase(oldPrefix);
        instPrefixNamesPair.insert({newInst, newPrefix});
      }
 
      // Inherit the old instance's extraction path.
      extractionPaths.try_emplace(newInst); // (create entry first)
      auto &extractionPath = (extractionPaths[newInst] = extractionPaths[inst]);
      auto instInnerRef = getInnerRefTo(newParentInst);
      innerRefToInstances[instInnerRef] = newParentInst;
      extractionPath.push_back(instInnerRef);
      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.getNameAttr()));
              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.getNameAttr()));
            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 extracted instances by their wrapper module, requested instance name,
  // and parent module. Instances in different parent modules cannot share a
  // wrapper submodule.
  using WrapperKey = std::tuple<Operation *, StringRef, StringRef>;
  llvm::MapVector<WrapperKey, SmallVector<InstanceOp>> instsByWrapper;
  for (auto &[inst, info] : extractedInstances) {
    if (info.wrapperModule.empty())
      continue;
    auto instanceName = info.wrapperInstanceName.empty()
                            ? info.wrapperModule
                            : info.wrapperInstanceName;
    instsByWrapper[{inst->getParentOfType<FModuleOp>(), info.wrapperModule,
                    instanceName}]
        .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 &[wrapperKey, insts] : instsByWrapper) {
    auto [parentOp, wrapperName, requestedInstanceName] = wrapperKey;
    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(wrapperName));
    auto wrapperInstName = builder.getStringAttr(
        getModuleNamespace(parent).newName(requestedInstanceName));
 
    // 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(instPrefixNamesPair[inst].first);
      unsigned portNum = inst.getNumResults();
      for (unsigned portIdx = 0; portIdx < portNum; ++portIdx) {
        auto name = inst.getPortName(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 = FModuleOp::create(
        builder, 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 = InstanceOp::create(
        builder, wrapper.getLoc(), wrapper, requestedInstanceName,
        NameKindEnum::DroppableName, ArrayRef<Attribute>{},
        /*portAnnotations=*/ArrayRef<Attribute>{}, /*lowerToBind=*/false,
        /*doNotPrint=*/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);
        MatchingConnectOp::create(builder, 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(ClassOp &sifiveMetadataClass) {
  LLVM_DEBUG(llvm::dbgs() << "\nGenerating trace files\n");
 
  // Group the extracted instances by their trace file name.
  llvm::MapVector<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;
  if (sifiveMetadataClass && !extractMetadataClass)
    createSchema();
 
  auto addPortsToClass = [&](ArrayRef<std::pair<Value, Twine>> objFields,
                             ClassOp classOp) {
    auto builderOM = mlir::ImplicitLocOpBuilder::atBlockEnd(
        classOp.getLoc(), classOp.getBodyBlock());
    auto portIndex = classOp.getNumPorts();
    SmallVector<std::pair<unsigned, PortInfo>> newPorts;
    for (auto [index, port] : enumerate(objFields)) {
      portIndex += index;
      auto obj = port.first;
      newPorts.emplace_back(
          portIndex,
          PortInfo(builderOM.getStringAttr(port.second + Twine(portIndex)),
                   obj.getType(), Direction::Out));
      auto blockarg =
          classOp.getBodyBlock()->addArgument(obj.getType(), obj.getLoc());
      PropAssignOp::create(builderOM, blockarg, obj);
    }
    classOp.insertPorts(newPorts);
  };
 
  HierPathCache pathCache(circuitOp, *symbolTable);
  SmallVector<std::pair<Value, Twine>> classFields;
  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(instPrefixNamesPair[inst].first);
      StringAttr origInstName(instPrefixNamesPair[inst].second);
      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 << " -> ";
 
      if (sifiveMetadataClass) {
        // Create the entry for this extracted instance in the metadata class.
        auto builderOM = mlir::ImplicitLocOpBuilder::atBlockEnd(
            inst.getLoc(), extractMetadataClass.getBodyBlock());
        auto prefixName = StringConstantOp::create(builderOM, prefix);
        auto object = ObjectOp::create(builderOM, schemaClass, prefix);
        auto fPrefix =
            ObjectSubfieldOp::create(builderOM, object, prefixNameFieldId);
        PropAssignOp::create(builderOM, fPrefix, prefixName);
 
        auto targetInstance = innerRefToInstances[path.front()];
        SmallVector<Attribute> pathOpAttr(llvm::reverse(path));
        auto nla = pathCache.getOpFor(
            ArrayAttr::get(circuitOp->getContext(), pathOpAttr));
 
        auto pathOp = createPathRef(targetInstance, nla, builderOM);
        auto fPath = ObjectSubfieldOp::create(builderOM, object, pathFieldId);
        PropAssignOp::create(builderOM, fPath, pathOp);
        auto fFile =
            ObjectSubfieldOp::create(builderOM, object, fileNameFieldId);
        auto fileNameOp = StringConstantOp::create(
            builderOM, builder.getStringAttr(fileName));
        PropAssignOp::create(builderOM, fFile, fileNameOp);
 
        auto finstName =
            ObjectSubfieldOp::create(builderOM, object, instNameFieldId);
        auto instNameOp = StringConstantOp::create(builderOM, origInstName);
        PropAssignOp::create(builderOM, finstName, instNameOp);
 
        // Now add this to the output field of the class.
        classFields.emplace_back(object, prefix + "_field");
      }
      // 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() &&
             !instanceInfo->anyInstanceInDesign(cast<igraph::ModuleOpInterface>(
                 symbolTable->lookup(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);
      }
      os << "." << origInstName.getValue();
      // 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.
    sv::VerbatimOp::create(builder, builder.getUnknownLoc(), buffer,
                           ValueRange{}, builder.getArrayAttr(symbols));
  }
  if (!classFields.empty()) {
    addPortsToClass(classFields, extractMetadataClass);
    // This extract instances metadata class, now needs to be instantiated
    // inside the SifiveMetadata class. This also updates its signature, so keep
    // the object of the SifiveMetadata class updated.
    auto builderOM = mlir::ImplicitLocOpBuilder::atBlockEnd(
        sifiveMetadataClass->getLoc(), sifiveMetadataClass.getBodyBlock());
    SmallVector<std::pair<Value, Twine>> classFields = {
        {ObjectOp::create(
             builderOM, extractMetadataClass,
             builderOM.getStringAttr("extract_instances_metadata")),
         "extractedInstances_field"}};
 
    addPortsToClass(classFields, sifiveMetadataClass);
    auto *node = instanceGraph->lookup(sifiveMetadataClass);
    assert(node && node->hasOneUse());
    ObjectOp metadataObj = (*node->usesBegin())->getInstance<ObjectOp>();
    assert(metadataObj &&
           "expected the class to be instantiated by an object op");
    builderOM.setInsertionPoint(metadataObj);
    auto newObj =
        ObjectOp::create(builderOM, sifiveMetadataClass, metadataObj.getName());
    metadataObj->replaceAllUsesWith(newObj);
    metadataObj->remove();
  }
}
 
void ExtractInstancesPass::createSchema() {
 
  auto *context = circuitOp->getContext();
  auto unknownLoc = mlir::UnknownLoc::get(context);
  auto builderOM = mlir::ImplicitLocOpBuilder::atBlockEnd(
      unknownLoc, circuitOp.getBodyBlock());
  mlir::Type portsType[] = {
      stringType, // name
      pathType,   // extracted instance path
      stringType, // filename
      stringType  // instance name
  };
  StringRef portFields[] = {"name", "path", "filename", "inst_name"};
 
  schemaClass = ClassOp::create(builderOM, "ExtractInstancesSchema", portFields,
                                portsType);
 
  // Now create the class that will instantiate the schema objects.
  SmallVector<PortInfo> mports;
  extractMetadataClass = ClassOp::create(
      builderOM, builderOM.getStringAttr("ExtractInstancesMetadata"), mports);
}