From 3f0dcf4405fbcc3e95405c2429c8285641144e20 Mon Sep 17 00:00:00 2001
From: Will Dietz <will.dietz@sifive.com>
Date: Wed, 26 Jun 2024 08:35:10 -0500
Subject: [PATCH] [FIRRTL] Allow layers under when and match. (#7234)

Behavior is same as-if the contained operations were
not under a layer.
---
 .../circt/Dialect/FIRRTL/FIRRTLStatements.td  |  4 ++-
 lib/Dialect/FIRRTL/FIRRTLOps.cpp              |  4 +++
 lib/Dialect/FIRRTL/Transforms/ExpandWhens.cpp |  6 ++++
 test/Dialect/FIRRTL/expand-whens.mlir         | 14 +++++++++
 test/firtool/layers.fir                       | 29 +++++++++++++++++--
 5 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/include/circt/Dialect/FIRRTL/FIRRTLStatements.td b/include/circt/Dialect/FIRRTL/FIRRTLStatements.td
index 649801ad03..6f0674d3be 100644
--- a/include/circt/Dialect/FIRRTL/FIRRTLStatements.td
+++ b/include/circt/Dialect/FIRRTL/FIRRTLStatements.td
@@ -374,7 +374,9 @@ def VerifCoverIntrinsicOp : VerifIntrinsicOp<"cover">;
 def LayerBlockOp : FIRRTLOp<
   "layerblock",
   [SingleBlock, NoTerminator, NoRegionArguments,
-   ParentOneOf<["firrtl::FModuleOp", "firrtl::LayerBlockOp"]>,
+   ParentOneOf<[
+     "firrtl::FModuleOp", "firrtl::LayerBlockOp",
+     "firrtl::WhenOp", "firrtl::MatchOp"]>,
    DeclareOpInterfaceMethods<SymbolUserOpInterface>]
 > {
   let summary = "A definition of a layer block";
diff --git a/lib/Dialect/FIRRTL/FIRRTLOps.cpp b/lib/Dialect/FIRRTL/FIRRTLOps.cpp
index aa5adadc38..ff9f114959 100644
--- a/lib/Dialect/FIRRTL/FIRRTLOps.cpp
+++ b/lib/Dialect/FIRRTL/FIRRTLOps.cpp
@@ -6171,6 +6171,10 @@ LogicalResult LayerBlockOp::verify() {
   auto layerName = getLayerName();
   auto *parentOp = (*this)->getParentOp();
 
+  // Get parent operation that isn't a when or match.
+  while (isa<WhenOp, MatchOp>(parentOp))
+    parentOp = parentOp->getParentOp();
+
   // Verify the correctness of the symbol reference.  Only verify that this
   // layer block makes sense in its parent module or layer block.
   auto nestedReferences = layerName.getNestedReferences();
diff --git a/lib/Dialect/FIRRTL/Transforms/ExpandWhens.cpp b/lib/Dialect/FIRRTL/Transforms/ExpandWhens.cpp
index f785f92c8a..6009ac9efd 100644
--- a/lib/Dialect/FIRRTL/Transforms/ExpandWhens.cpp
+++ b/lib/Dialect/FIRRTL/Transforms/ExpandWhens.cpp
@@ -533,6 +533,7 @@ public:
   void visitStmt(PrintFOp op);
   void visitStmt(StopOp op);
   void visitStmt(WhenOp op);
+  void visitStmt(LayerBlockOp op);
   void visitStmt(RefForceOp op);
   void visitStmt(RefForceInitialOp op);
   void visitStmt(RefReleaseOp op);
@@ -619,6 +620,11 @@ void WhenOpVisitor::visitStmt(WhenOp whenOp) {
   processWhenOp(whenOp, condition);
 }
 
+// NOLINTNEXTLINE(misc-no-recursion)
+void WhenOpVisitor::visitStmt(LayerBlockOp layerBlockOp) {
+  process(*layerBlockOp.getBody());
+}
+
 void WhenOpVisitor::visitStmt(RefForceOp op) {
   op.getPredicateMutable().assign(andWithCondition(op, op.getPredicate()));
 }
diff --git a/test/Dialect/FIRRTL/expand-whens.mlir b/test/Dialect/FIRRTL/expand-whens.mlir
index 77a750119b..c07606f2c0 100644
--- a/test/Dialect/FIRRTL/expand-whens.mlir
+++ b/test/Dialect/FIRRTL/expand-whens.mlir
@@ -589,6 +589,20 @@ firrtl.module @WhenInGroup(in %cond : !firrtl.uint<1>) {
   }
 }
 
+// Check that expand whens works for layers under when's.
+firrtl.layer @Layer bind {}
+// CHECK-LABEL: firrtl.module @LayerUnderWhen(
+// CHECK-NEXT:  firrtl.layerblock @Layer
+// CHECK:       firrtl.printf %clock, %cond
+firrtl.module @LayerUnderWhen(in %cond : !firrtl.uint<1>, in %clock : !firrtl.clock) {
+  firrtl.when %cond : !firrtl.uint<1> {
+    firrtl.layerblock @Layer {
+      %c1_ui1 = firrtl.constant 1 : !firrtl.uint<1>
+      firrtl.printf %clock, %c1_ui1, "Condition is true" : !firrtl.clock, !firrtl.uint<1>
+    }
+  }
+}
+
 // CHECK: firrtl.class @ClassWithInput(in %in: !firrtl.string)
 firrtl.class @ClassWithInput(in %in: !firrtl.string) {}
 
diff --git a/test/firtool/layers.fir b/test/firtool/layers.fir
index 30eaffc2f4..2f8cdb306d 100644
--- a/test/firtool/layers.fir
+++ b/test/firtool/layers.fir
@@ -9,6 +9,10 @@ circuit Foo: %[[
   {
     "class": "firrtl.transforms.DontTouchAnnotation",
     "target": "~Foo|Foo>y"
+  },
+  {
+    "class": "firrtl.transforms.DontTouchAnnotation",
+    "target": "~Foo|Foo>z"
   }
 ]]
   layer A, bind:
@@ -17,16 +21,34 @@ circuit Foo: %[[
   public module Foo:
     input in: UInt<1>
 
+    input clock: Clock
+    input cond: UInt<1>
+    input enable: UInt<1>
+
     layerblock A:
       node x = in
 
       layerblock B:
         node y = x
 
+      when cond:
+        layerblock B:
+          when x:
+            node z = x
+            assert(clock, cond, enable, "Test")
+
 ; CHECK-LABEL: module Foo_A_B(
-; CHECK-NEXT:    input x
+; CHECK-NEXT:    input x,
+; CHECK-NEXT:          cond,
+; CHECK-NEXT:          enable,
+; CHECK-NEXT:          clock
 ; CHECK-NEXT:  );
 ; CHECK:         wire y = x;
+; CHECK:         wire z = x;
+; CHECK:         always @(posedge clock) begin
+; CHECK-NEXT:      if (cond & x & enable)
+; CHECK-NEXT:        assert(cond) else $error("Test");
+; CHECK-NEXT:    end // always @(posedge)
 ; CHECK-NEXT:  endmodule
 
 ; CHECK-LABEL: module Foo_A(
@@ -40,7 +62,10 @@ circuit Foo: %[[
 ; CHECK-NEXT:  `ifndef layers_Foo_A_B
 ; CHECK-NEXT:  `define layers_Foo_A_B
 ; CHECK-NEXT:  bind Foo Foo_A_B a_b (
-; CHECK-NEXT:    x (Foo.a.x_probe)
+; CHECK-NEXT:    .x (Foo.a.x_probe),
+; CHECK-NEXT:    .cond (cond),
+; CHECK-NEXT:    .enable (enable),
+; CHECK-NEXT:    .clock (clock)
 ; CHECK-NEXT:  );
 ; CHECK-NEXT:  `endif // layers_Foo_A_B