From e76d2644a7782f87497668efbc57e76365b887c1 Mon Sep 17 00:00:00 2001
From: Andrew Young <youngar17@gmail.com>
Date: Mon, 7 Oct 2024 21:55:31 -0700
Subject: [PATCH] [FIRRTL] InferResets: properly lower FART'd registers

This is quick fix for https://github.com/llvm/circt/issues/7675.

The way this is supposed to work is, FRT is not supposed to add a reset
if a register already has one, but it isn't smart enough to understand
that an invalid value corresponds to a register without a reset. As a
result, it will not add a reset to this register. Later on, in
SFCCompat, when choosing how to lower the Invalid value, it should "see"
that the registers should all have a reset, by looking for the FRT
annotation, and choose to lower the reset value to 0 or remove the reset
accordingly. The "bug" is that the FRT annotation on a parent module
won't be discovered on the child module.

This attaches the annotation on to every module in a reset domain, so
that it can be found.
---
 lib/Dialect/FIRRTL/Transforms/InferResets.cpp |  8 ++++++
 test/Dialect/FIRRTL/infer-resets.mlir         | 27 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/lib/Dialect/FIRRTL/Transforms/InferResets.cpp b/lib/Dialect/FIRRTL/Transforms/InferResets.cpp
index ca52ee6c470a..1d6f8b943aa7 100644
--- a/lib/Dialect/FIRRTL/Transforms/InferResets.cpp
+++ b/lib/Dialect/FIRRTL/Transforms/InferResets.cpp
@@ -1685,6 +1685,14 @@ LogicalResult InferResetsPass::implementFullReset(FModuleOp module,
     return success();
   }
 
+  // Add an annotation indicating that this module belongs to a reset domain.
+  auto *context = module.getContext();
+  AnnotationSet annotations(module);
+  annotations.addAnnotations(DictionaryAttr::get(
+      context, NamedAttribute(StringAttr::get(context, "class"),
+                              StringAttr::get(context, fullResetAnnoClass))));
+  annotations.applyToOperation(module);
+
   // If needed, add a reset port to the module.
   Value actualReset = domain.existingValue;
   if (domain.newPortName) {
diff --git a/test/Dialect/FIRRTL/infer-resets.mlir b/test/Dialect/FIRRTL/infer-resets.mlir
index 1eb7d0e0a77d..9a64ea1265fc 100644
--- a/test/Dialect/FIRRTL/infer-resets.mlir
+++ b/test/Dialect/FIRRTL/infer-resets.mlir
@@ -845,6 +845,33 @@ firrtl.circuit "ZeroWidthRegister" {
 
 // -----
 
+// Every module which is contained inside a reset domain should be annotated as
+// such, so that internal registers can be lowered later correctly.
+// https://github.com/llvm/circt/issues/7675
+firrtl.circuit "top" {
+  // CHECK: firrtl.module private @test
+  // CHECK-SAME: annotations = [{class = "circt.FullResetAnnotation"}]
+  firrtl.module private @test(in %clock: !firrtl.clock, in %reset: !firrtl.asyncreset, in %in: !firrtl.uint<8>, out %out: !firrtl.uint<8>) {
+    %resetvalue = firrtl.wire : !firrtl.uint<8>
+    %invalid_ui8 = firrtl.invalidvalue : !firrtl.uint<8>
+    firrtl.matchingconnect %resetvalue, %invalid_ui8 : !firrtl.uint<8>
+    %reg1 = firrtl.regreset %clock, %reset, %resetvalue : !firrtl.clock, !firrtl.asyncreset, !firrtl.uint<8>, !firrtl.uint<8>
+    firrtl.matchingconnect %reg1, %in : !firrtl.uint<8>
+    firrtl.matchingconnect %out, %reg1 : !firrtl.uint<8>
+  }
+  // CHECK: firrtl.module @top
+  // CHECK-SAME: annotations = [{class = "circt.FullResetAnnotation"}]
+  firrtl.module @top(in %clock: !firrtl.clock, in %reset: !firrtl.asyncreset [{class = "circt.FullResetAnnotation", resetType = "async"}], in %in: !firrtl.uint<8>, out %out: !firrtl.uint<8>) attributes {convention = #firrtl<convention scalarized>} {
+    %child_clock, %child_reset, %child_in, %child_out = firrtl.instance child @test(in clock: !firrtl.clock, in reset: !firrtl.asyncreset, in in: !firrtl.uint<8>, out out: !firrtl.uint<8>)
+    firrtl.matchingconnect %child_clock, %clock : !firrtl.clock
+    firrtl.matchingconnect %child_reset, %reset : !firrtl.asyncreset
+    firrtl.matchingconnect %child_in, %in : !firrtl.uint<8>
+    firrtl.matchingconnect %out, %child_out : !firrtl.uint<8>
+  }
+}
+
+// -----
+
 // Check that unaffected fields ("data") are not being affected by width
 // inference. See https://github.com/llvm/circt/issues/2857.
 // CHECK-LABEL: firrtl.module @ZeroLengthVectorInBundle1