[LowerTypes] Track public modules and lower ports properly (#4655)

This changes LowerTypes to look up module visibilities. `module.isPublic()` cannot be called from a different thread because of a race condition so this PR first creates a set of public modules and checks the visibility by querying to the set.
llvm · Feb 14, 2023 · b16d03f · b16d03f
1 parent cc1a5ee
commit b16d03f
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 9 deletions.
diff --git a/lib/Dialect/FIRRTL/Transforms/LowerTypes.cpp b/lib/Dialect/FIRRTL/Transforms/LowerTypes.cpp
@@ -38,6 +38,7 @@
 #include "circt/Dialect/FIRRTL/Namespace.h"
 #include "circt/Dialect/FIRRTL/Passes.h"
 #include "circt/Dialect/HW/HWAttributes.h"
+#include "circt/Dialect/HW/HWOpInterfaces.h"
 #include "circt/Dialect/SV/SVOps.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/Threading.h"
@@ -354,10 +355,11 @@ struct TypeLoweringVisitor : public FIRRTLVisitor<TypeLoweringVisitor, bool> {
   TypeLoweringVisitor(MLIRContext *context,
                       PreserveAggregate::PreserveMode preserveAggregate,
                       bool preservePublicTypes, SymbolTable &symTbl,
-                      const AttrCache &cache)
+                      const AttrCache &cache,
+                      const llvm::DenseSet<FModuleLike> &publicModuleSet)
       : context(context), aggregatePreservationMode(preserveAggregate),
-        preservePublicTypes(preservePublicTypes), symTbl(symTbl), cache(cache) {
-  }
+        preservePublicTypes(preservePublicTypes), symTbl(symTbl), cache(cache),
+        publicModuleSet(publicModuleSet) {}
   using FIRRTLVisitor<TypeLoweringVisitor, bool>::visitDecl;
   using FIRRTLVisitor<TypeLoweringVisitor, bool>::visitExpr;
   using FIRRTLVisitor<TypeLoweringVisitor, bool>::visitStmt;
@@ -440,6 +442,9 @@ struct TypeLoweringVisitor : public FIRRTLVisitor<TypeLoweringVisitor, bool> {
   // Cache some attributes
   const AttrCache &cache;
 
+  // Keep track of public modules.
+  const llvm::DenseSet<FModuleLike> &publicModuleSet;
+
   // Set true if the lowering failed.
   bool encounteredError = false;
 };
@@ -454,11 +459,12 @@ TypeLoweringVisitor::getPreservatinoModeForModule(FModuleLike module) {
   if (!isa<FModuleOp>(module))
     return PreserveAggregate::None;
 
-  // If `module` is a top-module, we have to lower ports. Don't read attributes
-  // of `module` since the attributes could be mutated in a different thread.
+  // If preservePublicTypes is true, we have to lower ports of public modules.
+  // Query the module visibility to `publicModuleSet`. Don't call
+  // `module.isPublic` since the attributes could be mutated in a different
+  // thread.
   if (aggregatePreservationMode != PreserveAggregate::None &&
-      preservePublicTypes &&
-      module->getParentOfType<CircuitOp>().getMainModule(&symTbl) == module)
+      preservePublicTypes && publicModuleSet.count(module))
     return PreserveAggregate::None;
   return aggregatePreservationMode;
 }
@@ -1395,6 +1401,7 @@ void LowerTypesPass::runOnOperation() {
       llvm::dbgs() << "===- Running LowerTypes Pass "
                       "------------------------------------------------===\n");
   std::vector<FModuleLike> ops;
+  llvm::DenseSet<FModuleLike> publicModuleSet;
   // Symbol Table
   SymbolTable symTbl(getOperation());
   // Cached attr
@@ -1405,16 +1412,20 @@ void LowerTypesPass::runOnOperation() {
                  [&](Operation &op) {
                    // Creating a map of all ops in the circt, but only modules
                    // are relevant.
-                   if (auto module = dyn_cast<FModuleLike>(op))
+                   if (auto module = dyn_cast<FModuleLike>(op)) {
                      ops.push_back(module);
+                     if (cast<hw::HWModuleLike>(op).isPublic())
+                       publicModuleSet.insert(module);
+                   }
                  });
 
   LLVM_DEBUG(llvm::dbgs() << "Recording Inner Symbol Renames:\n");
 
   // This lambda, executes in parallel for each Op within the circt.
   auto lowerModules = [&](FModuleLike op) -> LogicalResult {
     auto tl = TypeLoweringVisitor(&getContext(), preserveAggregate,
-                                  preservePublicTypes, symTbl, cache);
+                                  preservePublicTypes, symTbl, cache,
+                                  publicModuleSet);
     tl.lowerModule(op);
 
     return LogicalResult::failure(tl.isFailed());

diff --git a/test/Dialect/FIRRTL/lower-types-aggregate.mlir b/test/Dialect/FIRRTL/lower-types-aggregate.mlir
@@ -20,4 +20,10 @@ firrtl.circuit "TopLevel" {
   // 1D_VEC: %a_0: !firrtl.uint<1>
   firrtl.module private @Bar(in %a: !firrtl.vector<uint<1>, 1>) {
   }
+  // CHECK-LABEL: PublicModule
+  // CHECK-NOT: firrtl.bundle
+  // NOT_PRESERVE_PUBLIC_TYPES-LABEL: PublicModule
+  // NOT_PRESERVE_PUBLIC_TYPES: firrtl.bundle
+  firrtl.module public @PublicModule(in %source: !firrtl.bundle<valid: uint<1>>) {
+  }
 }