[ImportVerilog] Capture analysis should skip reprocessing identical module instances (#10338)

Currently capture analysis unconditionally descends into all module instance
bodies. This forces slang to instantiate complete ASTs for all module
instances, transitively. For large designs this does not scale. For example
with one 7M LOC Verilog example, ImportVerilog consumes more than 200GB of
memory before crashing with OOM.

Slang lets clients avoid this problem via "canonical instance bodies"
that express deduplication. Setting `VisitCanonical` to true in the
`ASTVisitor` template causes the `InstanceBodySymbol` visit to use
the deduplicated canonical instance body (when available). For this
analysis we don't need to re-analyze duplicate instances, so we can
just keep track of which bodies we've visited and avoid descending
into already-visited bodies.

Co-authored-by: Jacques Pienaar <jacques+gh@japienaar.info>

Author: rocallahan

lib/Conversion/ImportVerilog/CaptureAnalysis.cpp

@@ -80,7 +80,9 @@ namespace {
 /// visiting enabled so that we recurse into all function bodies.
 struct CaptureWalker
     : public ASTVisitor<CaptureWalker, /*VisitStatements=*/true,
-                        /*VisitExpressions=*/true> {
+                        /*VisitExpressions=*/true,
+                        /*VisitBad=*/false,
+                        /*VisitCanonical=*/true> {
 
   /// The function whose body we are currently inside, or nullptr if we are at
   /// a scope outside any function.
@@ -89,6 +91,8 @@ struct CaptureWalker
   /// Captured variables per function.
   CaptureMap capturedVars;
 
+  mlir::DenseSet<const InstanceBodySymbol *> visitedInstanceBodies;
+
   /// Inverse call graph: maps each callee to the set of callers that call it.
   /// Used to propagate captures from callees to their callers. Uses MapVector
   /// for deterministic iteration order during propagation.
@@ -150,6 +154,15 @@ struct CaptureWalker
     visitDefault(expr);
   }
 
+  /// `VisitCanonical` is set above, so this visits the canonical instance body
+  /// if there is one. If there is and we've already visited it via some
+  /// other instance, don't process it again.
+  void handle(const InstanceBodySymbol &instance) {
+    if (visitedInstanceBodies.insert(&instance).second) {
+      visitDefault(instance);
+    }
+  }
+
   /// Propagate captures transitively through the call graph. For each callee
   /// that has captures, push each captured variable upward through all
   /// transitive callers using a worklist. A captured variable is only

lib/Conversion/ImportVerilog/ImportVerilogInternals.h

@@ -33,6 +33,14 @@ namespace ImportVerilog {
 
 using moore::Domain;
 
+/// Get the slang canonical body for the given `instance`, if there is one.
+/// If there isn't one, fall back to the normal body.
+inline const slang::ast::InstanceBodySymbol *
+getCanonicalBody(const slang::ast::InstanceSymbol &inst) {
+  const slang::ast::InstanceBodySymbol *body = inst.getCanonicalBody();
+  return body == nullptr ? &inst.body : body;
+}
+
 /// Port lowering information.
 struct PortLowering {
   const slang::ast::PortSymbol *
@@ -216,8 +224,16 @@ struct Context {
 
   /// Convert hierarchy and structure AST nodes to MLIR ops.
   LogicalResult convertCompilation();
+  /// Convert a module and its ports to an empty module op in the IR. Also adds
+  /// the op to the worklist of module bodies to be lowered. This acts like a
+  /// module "declaration", allowing instances to already refer to a module even
+  /// before its body has been lowered.
+  /// `module` must be the canonical module body if there is one.
   ModuleLowering *
   convertModuleHeader(const slang::ast::InstanceBodySymbol *module);
+  /// Convert a module's body to the corresponding IR ops. The module op must
+  /// have already been created earlier through a `convertModuleHeader` call.
+  /// `module` must be the canonical module body if there is one.
   LogicalResult convertModuleBody(const slang::ast::InstanceBodySymbol *module);
   LogicalResult convertPackage(const slang::ast::PackageSymbol &package);
   FunctionLowering *
@@ -419,6 +435,7 @@ struct Context {
   std::map<LocationKey, Operation *> orderedRootOps;
 
   /// How we have lowered modules to MLIR.
+  /// The keys must be the slang canonical module bodies where they exist.
   DenseMap<const slang::ast::InstanceBodySymbol *,
            std::unique_ptr<ModuleLowering>>
       modules;

lib/Conversion/ImportVerilog/Structure.cpp

@@ -407,13 +407,19 @@ struct ModuleVisitor : public BaseVisitor {
     using slang::ast::MultiPortSymbol;
     using slang::ast::PortSymbol;
 
+    // Always operate on the canonical instance body if there is one.
+    // This means any symbols we record will be the symbols from the
+    // canonical body, which will match up with the symbols encountered
+    // by analyses which visit the canonical bodies.
+    const slang::ast::InstanceBodySymbol *body = getCanonicalBody(instNode);
+
     // Interface instances are expanded inline into individual variable/net ops
     // rather than creating a moore.instance op.
-    auto defKind = instNode.body.getDefinition().definitionKind;
+    auto defKind = body->getDefinition().definitionKind;
     if (defKind == slang::ast::DefinitionKind::Interface)
       return expandInterfaceInstance(instNode);
 
-    auto *moduleLowering = context.convertModuleHeader(&instNode.body);
+    auto *moduleLowering = context.convertModuleHeader(body);
     if (!moduleLowering)
       return failure();
     auto module = moduleLowering->op;
@@ -952,11 +958,13 @@ LogicalResult Context::convertCompilation() {
   // Interfaces are not lowered as modules; they are expanded inline at each
   // use site, so skip them here.
   SmallVector<const slang::ast::InstanceSymbol *> topInstances;
-  for (auto *inst : root.topInstances)
-    if (inst->body.getDefinition().definitionKind !=
+  for (auto *inst : root.topInstances) {
+    const slang::ast::InstanceBodySymbol *body = getCanonicalBody(*inst);
+    if (body->getDefinition().definitionKind !=
         slang::ast::DefinitionKind::Interface)
-      if (!convertModuleHeader(&inst->body))
+      if (!convertModuleHeader(body))
         return failure();
+  }
 
   // Convert all the root module definitions.
   while (!moduleWorklist.empty()) {
@@ -1006,10 +1014,6 @@ LogicalResult Context::convertCompilation() {
   return success();
 }
 
-/// Convert a module and its ports to an empty module op in the IR. Also adds
-/// the op to the worklist of module bodies to be lowered. This acts like a
-/// module "declaration", allowing instances to already refer to a module even
-/// before its body has been lowered.
 ModuleLowering *
 Context::convertModuleHeader(const slang::ast::InstanceBodySymbol *module) {
   using slang::ast::ArgumentDirection;
@@ -1024,53 +1028,8 @@ Context::convertModuleHeader(const slang::ast::InstanceBodySymbol *module) {
   timeScale = module->getTimeScale().value_or(slang::TimeScale());
   llvm::scope_exit timeScaleGuard([&] { timeScale = prevTimeScale; });
 
-  auto parameters = module->getParameters();
-  bool hasModuleSame = false;
-  // If there is already exist a module that has the same name with this
-  // module ,has the same parent scope and has the same parameters we can
-  // define this module is a duplicate module
-  for (auto const &existingModule : modules) {
-    if (module->getDeclaringDefinition() ==
-        existingModule.getFirst()->getDeclaringDefinition()) {
-      auto moduleParameters = existingModule.getFirst()->getParameters();
-      hasModuleSame = true;
-      for (auto it1 = parameters.begin(), it2 = moduleParameters.begin();
-           it1 != parameters.end() && it2 != moduleParameters.end();
-           it1++, it2++) {
-        // Parameters size different
-        if (it1 == parameters.end() || it2 == moduleParameters.end()) {
-          hasModuleSame = false;
-          break;
-        }
-        const auto *para1 = (*it1)->symbol.as_if<ParameterSymbol>();
-        const auto *para2 = (*it2)->symbol.as_if<ParameterSymbol>();
-        // Parameters kind different
-        if ((para1 == nullptr) ^ (para2 == nullptr)) {
-          hasModuleSame = false;
-          break;
-        }
-        // Compare ParameterSymbol
-        if (para1 != nullptr) {
-          hasModuleSame = para1->getValue() == para2->getValue();
-        }
-        // Compare TypeParameterSymbol
-        if (para1 == nullptr) {
-          auto para1Type = convertType(
-              (*it1)->symbol.as<TypeParameterSymbol>().getTypeAlias());
-          auto para2Type = convertType(
-              (*it2)->symbol.as<TypeParameterSymbol>().getTypeAlias());
-          hasModuleSame = para1Type == para2Type;
-        }
-        if (!hasModuleSame)
-          break;
-      }
-      if (hasModuleSame) {
-        module = existingModule.first;
-        break;
-      }
-    }
-  }
-
+  // `module` is the canonical module body if it exists (i.e. deduplicated by
+  // slang).
   auto &slot = modules[module];
   if (slot)
     return slot.get();
@@ -1273,8 +1232,6 @@ Context::convertModuleHeader(const slang::ast::InstanceBodySymbol *module) {
   return &lowering;
 }
 
-/// Convert a module's body to the corresponding IR ops. The module op must have
-/// already been created earlier through a `convertModuleHeader` call.
 LogicalResult
 Context::convertModuleBody(const slang::ast::InstanceBodySymbol *module) {
   auto &lowering = *modules[module];