[FIRRTL][IMCP] Overdefine ports of modules with unknown symbol uses

include/circt/Support/InstanceGraph.h

@@ -200,11 +200,30 @@ class InstanceGraph {
   InstanceGraph(const InstanceGraph &) = delete;
   virtual ~InstanceGraph() = default;
 
-  /// Look up an InstanceGraphNode for a module.
-  InstanceGraphNode *lookup(ModuleOpInterface op);
+  /// Lookup an module by name. Returns null if no module with the given name
+  /// exists in the instance graph.
+  InstanceGraphNode *lookupOrNull(StringAttr name);
+
+  /// Look up an InstanceGraphNode for a module. Returns null if the module has
+  /// not been added to the instance graph.
+  InstanceGraphNode *lookupOrNull(ModuleOpInterface op) {
+    return lookup(op.getModuleNameAttr());
+  }
+
+  /// Look up an InstanceGraphNode for a module. Aborts if the module does not
+  /// exist.
+  InstanceGraphNode *lookup(ModuleOpInterface op) {
+    auto *node = lookupOrNull(op);
+    assert(node != nullptr && "Module not in InstanceGraph!");
+    return node;
+  }
 
-  /// Lookup an module by name.
-  InstanceGraphNode *lookup(StringAttr name);
+  /// Lookup an module by name. Aborts if the module does not exist.
+  InstanceGraphNode *lookup(StringAttr name) {
+    auto *node = lookupOrNull(name);
+    assert(node != nullptr && "Module not in InstanceGraph!");
+    return node;
+  }
 
   /// Lookup an InstanceGraphNode for a module.
   InstanceGraphNode *operator[](ModuleOpInterface op) { return lookup(op); }

lib/Dialect/FIRRTL/Transforms/IMConstProp.cpp

@@ -369,12 +369,44 @@ void IMConstPropPass::runOnOperation() {
 
   instanceGraph = &getAnalysis<InstanceGraph>();
 
-  // Mark the input ports of public modules as being overdefined.
-  for (auto module : circuit.getBodyBlock()->getOps<FModuleOp>()) {
-    if (module.isPublic()) {
-      markBlockExecutable(module.getBodyBlock());
-      for (auto port : module.getBodyBlock()->getArguments())
-        markOverdefined(port);
+  // Mark input ports as overdefined where appropriate.
+  for (auto &op : circuit.getOps()) {
+    // Inputs of public modules are overdefined.
+    if (auto module = dyn_cast<FModuleOp>(op)) {
+      if (module.isPublic()) {
+        markBlockExecutable(module.getBodyBlock());
+        for (auto port : module.getBodyBlock()->getArguments())
+          markOverdefined(port);
+      }
+      continue;
+    }
+
+    // Otherwise we check whether the top-level operation contains any
+    // references to modules. Symbol uses in NLAs are ignored.
+    if (isa<hw::HierPathOp>(op))
+      continue;
+
+    // Inputs of modules referenced by unknown operations are overdefined, since
+    // we don't know how those operations affect the input port values. This
+    // handles things like `firrtl.formal`, which may may assign symbolic values
+    // to input ports of a private module.
+    auto symbolUses = SymbolTable::getSymbolUses(&op);
+    if (!symbolUses)
+      continue;
+    for (const auto &use : *symbolUses) {
+      if (auto symRef = dyn_cast<FlatSymbolRefAttr>(use.getSymbolRef())) {
+        if (auto *igNode = instanceGraph->lookupOrNull(symRef.getAttr())) {
+          if (auto module = dyn_cast<FModuleOp>(*igNode->getModule())) {
+            LLVM_DEBUG(llvm::dbgs()
+                       << "Unknown use of " << module.getModuleNameAttr()
+                       << " in " << op.getName()
+                       << ", marking inputs as overdefined\n");
+            markBlockExecutable(module.getBodyBlock());
+            for (auto port : module.getBodyBlock()->getArguments())
+              markOverdefined(port);
+          }
+        }
+      }
     }
   }
 

lib/Support/InstanceGraph.cpp

@@ -106,16 +106,13 @@ void InstanceGraph::erase(InstanceGraphNode *node) {
   nodes.erase(node);
 }
 
-InstanceGraphNode *InstanceGraph::lookup(StringAttr name) {
+InstanceGraphNode *InstanceGraph::lookupOrNull(StringAttr name) {
   auto it = nodeMap.find(name);
-  assert(it != nodeMap.end() && "Module not in InstanceGraph!");
+  if (it == nodeMap.end())
+    return nullptr;
   return it->second;
 }
 
-InstanceGraphNode *InstanceGraph::lookup(ModuleOpInterface op) {
-  return lookup(cast<ModuleOpInterface>(op).getModuleNameAttr());
-}
-
 void InstanceGraph::replaceInstance(InstanceOpInterface inst,
                                     InstanceOpInterface newInst) {
   assert(inst.getReferencedModuleNamesAttr() ==

test/Dialect/FIRRTL/imconstprop.mlir

@@ -1,4 +1,4 @@
-// RUN: circt-opt -pass-pipeline='builtin.module(firrtl.circuit(firrtl-imconstprop))' --split-input-file  %s | FileCheck %s
+// RUN: circt-opt --firrtl-imconstprop --split-input-file --allow-unregistered-dialect %s | FileCheck %s
 
 firrtl.circuit "Test" {
 
@@ -896,3 +896,38 @@ firrtl.circuit "Layers" {
     }
   }
 }
+
+// -----
+
+// CHECK-LABEL: firrtl.circuit "PublicTop"
+firrtl.circuit "PublicTop" {
+  // CHECK-LABEL: firrtl.module private @Foo
+  firrtl.module private @Foo(in %a: !firrtl.uint<1>) {
+    // CHECK-NOT: firrtl.constant 0
+    // CHECK: firrtl.int.verif.assert %a
+    firrtl.int.verif.assert %a : !firrtl.uint<1>
+  }
+  // CHECK-LABEL: firrtl.module private @Bar
+  firrtl.module private @Bar(in %a: !firrtl.uint<1>) {
+    // CHECK-NOT: firrtl.constant 0
+    // CHECK: firrtl.int.verif.assert %a
+    firrtl.int.verif.assert %a : !firrtl.uint<1>
+  }
+  firrtl.module @PublicTop() {
+    %c0_ui1 = firrtl.constant 0 : !firrtl.uint<1>
+    %foo_a = firrtl.instance foo @Foo(in a: !firrtl.uint<1>)
+    %bar_a = firrtl.instance bar @Bar(in a: !firrtl.uint<1>)
+    firrtl.matchingconnect %foo_a, %c0_ui1 : !firrtl.uint<1>
+    firrtl.matchingconnect %bar_a, %c0_ui1 : !firrtl.uint<1>
+  }
+  firrtl.module private @PrivateTop1(in %a: !firrtl.uint<1>) {
+    %foo_a = firrtl.instance foo @Foo(in a: !firrtl.uint<1>)
+    firrtl.matchingconnect %foo_a, %a : !firrtl.uint<1>
+  }
+  firrtl.module private @PrivateTop2(in %a: !firrtl.uint<1>) {
+    %bar_a = firrtl.instance bar @Bar(in a: !firrtl.uint<1>)
+    firrtl.matchingconnect %bar_a, %a : !firrtl.uint<1>
+  }
+  firrtl.formal @Test, @PrivateTop1 {}
+  "some_unknown_dialect.op"() { magic = @PrivateTop2 } : () -> ()
+}