[FIRRTL/ExpandWhens] Support aggregate type registers (#2305)

This commit changes ExpandWhens to handle aggregate type registers. We need to expand the connection into individual ground type elements.
llvm · Dec 10, 2021 · 7b649c5 · 7b649c5
1 parent 6a6992a
commit 7b649c5
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 28 deletions.
diff --git a/lib/Dialect/FIRRTL/Transforms/ExpandWhens.cpp b/lib/Dialect/FIRRTL/Transforms/ExpandWhens.cpp
@@ -204,28 +204,48 @@ class LastConnectResolver : public FIRRTLVisitor<ConcreteT> {
 
   void visitDecl(WireOp op) { declareSinks(op.result(), Flow::Duplex); }
 
+  /// Take an aggregate value and construct ground subelements recursively.
+  /// And then apply function `fn`.
+  void foreachSubelement(OpBuilder &builder, Value value,
+                         llvm::function_ref<void(Value)> fn) {
+    TypeSwitch<Type>(value.getType())
+        .template Case<BundleType>([&](BundleType bundle) {
+          for (auto i : llvm::seq(0u, (unsigned)bundle.getNumElements())) {
+            auto subfield =
+                builder.create<SubfieldOp>(value.getLoc(), value, i);
+            foreachSubelement(builder, subfield, fn);
+          }
+        })
+        .template Case<FVectorType>([&](FVectorType vector) {
+          for (auto i : llvm::seq(0u, vector.getNumElements())) {
+            auto subindex =
+                builder.create<SubindexOp>(value.getLoc(), value, i);
+            foreachSubelement(builder, subindex, fn);
+          }
+        })
+        .Default([&](auto) { fn(value); });
+  }
+
   void visitDecl(RegOp op) {
-    // Registers are initialized to themselves.
-    // TODO: register of aggregate types are not supported.
-    if (!op.getType().cast<FIRRTLType>().isGround()) {
-      op.emitError() << "aggegate type register is not supported";
-      return;
-    }
-    auto connect = OpBuilder(op->getBlock(), ++Block::iterator(op))
-                       .create<ConnectOp>(op.getLoc(), op, op);
-    driverMap[getFieldRefFromValue(op.result())] = connect;
+    // Registers are initialized to themselves. If the register has an
+    // aggergate type, connect each ground type element.
+    auto builder = OpBuilder(op->getBlock(), ++Block::iterator(op));
+    auto fn = [&](Value value) {
+      auto connect = builder.create<ConnectOp>(value.getLoc(), value, value);
+      driverMap[getFieldRefFromValue(value)] = connect;
+    };
+    foreachSubelement(builder, op.result(), fn);
   }
 
   void visitDecl(RegResetOp op) {
-    // Registers are initialized to themselves.
-    // TODO: register of aggregate types are not supported.
-    if (!op.getType().cast<FIRRTLType>().isGround()) {
-      op.emitError() << "aggegate type register is not supported";
-      return;
-    }
-    auto connect = OpBuilder(op->getBlock(), ++Block::iterator(op))
-                       .create<ConnectOp>(op.getLoc(), op, op);
-    driverMap[getFieldRefFromValue(op.result())] = connect;
+    // Registers are initialized to themselves. If the register has an
+    // aggergate type, connect each ground type element.
+    auto builder = OpBuilder(op->getBlock(), ++Block::iterator(op));
+    auto fn = [&](Value value) {
+      auto connect = builder.create<ConnectOp>(value.getLoc(), value, value);
+      driverMap[getFieldRefFromValue(value)] = connect;
+    };
+    foreachSubelement(builder, op.result(), fn);
   }
 
   void visitDecl(InstanceOp op) {

diff --git a/test/Dialect/FIRRTL/expand-whens-errors.mlir b/test/Dialect/FIRRTL/expand-whens-errors.mlir
@@ -128,13 +128,4 @@ firrtl.circuit "CheckInitialization" {
 firrtl.module @CheckInitialization(in %p : !firrtl.uint<1>, out %out: !firrtl.vector<bundle<a:uint<1>, b:uint<1>>, 1>) {
   // expected-error @-1 {{sink "out[0].a" not fully initialized}}
 }
-}
-
-// -----
-
-firrtl.circuit "CheckInitialization" {
-firrtl.module @CheckInitialization(in %clock: !firrtl.clock) {
-  // expected-error @+1 {{aggegate type register is not supported}}
-  %reg0 = firrtl.reg %clock : !firrtl.vector<uint<1>, 1>
-}
-}
+}
diff --git a/test/Dialect/FIRRTL/expand-whens.mlir b/test/Dialect/FIRRTL/expand-whens.mlir
@@ -512,4 +512,22 @@ firrtl.module @vector_of_bundle(in %p : !firrtl.uint<1>, out %ret: !firrtl.vecto
   // CHECK-NOT: firrtl.connect %1, %c0_ui1 : !firrtl.uint<1>, !firrtl.uint<1>
   // CHECK:     firrtl.connect %1, %c1_ui1 : !firrtl.uint<1>, !firrtl.uint<1>
 }
+
+// CHECK-LABEL: @aggregate_register
+firrtl.module @aggregate_register(in %clock: !firrtl.clock) {
+  %0 = firrtl.reg %clock : !firrtl.bundle<a : uint<1>, b : uint<1>>
+  // CHECK:      %1 = firrtl.subfield %0(0)
+  // CHECK-NEXT: firrtl.connect %1, %1
+  // CHECK-NEXT: %2 = firrtl.subfield %0(1)
+  // CHECK-NEXT: firrtl.connect %2, %2
+}
+
+// CHECK-LABEL: @aggregate_regreset
+firrtl.module @aggregate_regreset(in %clock: !firrtl.clock, in %reset: !firrtl.uint<1>, in %resetval: !firrtl.vector<uint<1>, 2>) {
+  %0 = firrtl.regreset %clock, %reset, %resetval : !firrtl.uint<1>, !firrtl.vector<uint<1>, 2>, !firrtl.vector<uint<1>, 2>
+  // CHECK:      %1 = firrtl.subindex %0[0]
+  // CHECK-NEXT: firrtl.connect %1, %1
+  // CHECK-NEXT: %2 = firrtl.subindex %0[1]
+  // CHECK-NEXT: firrtl.connect %2, %2
+}
 }