From 4fadf7bd0ac6f5af1e66ba3db4fbe980c6bfd6c0 Mon Sep 17 00:00:00 2001
From: Mike Urbach <mikeurbach@gmail.com>
Date: Mon, 12 Feb 2024 19:47:33 -0700
Subject: [PATCH] [LowerClasses] Lower classes that instantiate properties.
 (#6688)

The previous logic would create classes for any FIRRTL modules that
had property ports. However, if some intermediate modules had property
ports, but their parents did not, we would not convert their parents,
which would break pass invariants and lead to failed assertions.

This adds logic to also convert any FIRRTL modules that instantiate
FIRRTL modules that will be converted. This upholds the pass
invariants, and makes it robust in case property ports are used in
part of the hierarchy but not exported from the top-level module of
the circuit.
---
 lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp | 16 +++++++++++++++-
 test/Dialect/FIRRTL/lower-classes.mlir         | 11 +++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp b/lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp
index 25ca0304299..ad66f37aed4 100644
--- a/lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp
+++ b/lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp
@@ -441,9 +441,23 @@ bool LowerClassesPass::shouldCreateClass(FModuleLike moduleLike) {
   if (moduleLike.isPublic())
     return true;
 
-  return llvm::any_of(moduleLike.getPorts(), [](PortInfo port) {
+  // Create a class for modules with property ports.
+  bool hasClassPorts = llvm::any_of(moduleLike.getPorts(), [](PortInfo port) {
     return isa<PropertyType>(port.type);
   });
+
+  if (hasClassPorts)
+    return true;
+
+  // Create a class for modules that instantiate classes or modules with
+  // property ports.
+  for (auto op :
+       moduleLike.getOperation()->getRegion(0).getOps<FInstanceLike>())
+    for (auto result : op->getResults())
+      if (type_isa<PropertyType>(result.getType()))
+        return true;
+
+  return false;
 }
 
 // Create an OM Class op from a FIRRTL Class op or Module op with properties.
diff --git a/test/Dialect/FIRRTL/lower-classes.mlir b/test/Dialect/FIRRTL/lower-classes.mlir
index e30aff21a45..516106a3a22 100644
--- a/test/Dialect/FIRRTL/lower-classes.mlir
+++ b/test/Dialect/FIRRTL/lower-classes.mlir
@@ -327,3 +327,14 @@ firrtl.circuit "AnyCast" {
     firrtl.propassign %foo, %0 : !firrtl.anyref
   }
 }
+
+// CHECK-LABEL: firrtl.circuit "ModuleWithPropertySubmodule"
+firrtl.circuit "ModuleWithPropertySubmodule" {
+  firrtl.module private @ModuleWithPropertySubmodule() {
+    %c0 = firrtl.integer 0
+    %inst.prop = firrtl.instance inst @SubmoduleWithProperty(in prop: !firrtl.integer)
+    firrtl.propassign %inst.prop, %c0 : !firrtl.integer
+  }
+  firrtl.module private @SubmoduleWithProperty(in %prop: !firrtl.integer) {
+  }
+}