[InferWidths] Add support for inferring through ref.sub. (#5000)

llvm · Apr 12, 2023 · a72ed75 · a72ed75
1 parent 282d804
commit a72ed75
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 2 deletions.
diff --git a/lib/Dialect/FIRRTL/Transforms/InferWidths.cpp b/lib/Dialect/FIRRTL/Transforms/InferWidths.cpp
@@ -1353,6 +1353,17 @@ LogicalResult InferenceMapping::mapOperation(Operation *op) {
                    op.getType());
       })
 
+      .Case<RefSubOp>([&](RefSubOp op) {
+        uint64_t fieldID = TypeSwitch<FIRRTLBaseType, uint64_t>(
+                               op.getInput().getType().getType())
+                               .Case<FVectorType>([](auto _) { return 1; })
+                               .Case<BundleType>([&](auto type) {
+                                 return type.getFieldID(op.getIndex());
+                               });
+        unifyTypes(FieldRef(op.getResult(), 0),
+                   FieldRef(op.getInput(), fieldID), op.getType());
+      })
+
       // Arithmetic and Logical Binary Primitives
       .Case<AddPrimOp, SubPrimOp>([&](auto op) {
         auto lhs = getExpr(op.getLhs());

diff --git a/test/Dialect/FIRRTL/infer-widths.mlir b/test/Dialect/FIRRTL/infer-widths.mlir
@@ -819,27 +819,55 @@ firrtl.circuit "Foo" {
   // CHECK-LABEL: @SubRef
   // CHECK-SAME: out %x: !firrtl.probe<uint<2>>
   // CHECK-SAME: out %y: !firrtl.rwprobe<uint<2>>
-  firrtl.module private @SubRef(out %x: !firrtl.probe<uint>, out %y : !firrtl.rwprobe<uint>) {
+  // CHECK-SAME: out %bov_ref: !firrtl.rwprobe<bundle<a: vector<uint<2>, 2>, b: uint<2>>>
+  firrtl.module private @SubRef(out %x: !firrtl.probe<uint>, out %y : !firrtl.rwprobe<uint>, out %bov_ref : !firrtl.rwprobe<bundle<a: vector<uint, 2>, b : uint>>) {
     // CHECK: firrtl.wire forceable : !firrtl.uint<2>, !firrtl.rwprobe<uint<2>>
     %w, %w_rw = firrtl.wire forceable : !firrtl.uint, !firrtl.rwprobe<uint>
+    %bov, %bov_rw = firrtl.wire forceable : !firrtl.bundle<a: vector<uint, 2>, b flip: uint>, !firrtl.rwprobe<bundle<a: vector<uint, 2>, b : uint>>
+    firrtl.ref.define %bov_ref, %bov_rw : !firrtl.rwprobe<bundle<a: vector<uint, 2>, b : uint>>
+
     %ref_w = firrtl.ref.send %w : !firrtl.uint
     firrtl.ref.define %x, %ref_w : !firrtl.probe<uint>
     firrtl.ref.define %y, %w_rw : !firrtl.rwprobe<uint>
 
     %c0_ui2 = firrtl.constant 0 : !firrtl.uint<2>
     firrtl.connect %w, %c0_ui2 : !firrtl.uint, !firrtl.uint<2>
+
+    %bov_a = firrtl.subfield %bov[a] : !firrtl.bundle<a: vector<uint, 2>, b flip: uint>
+    %bov_a_1 = firrtl.subindex %bov_a[1] : !firrtl.vector<uint, 2>
+    %bov_b = firrtl.subfield %bov[b] : !firrtl.bundle<a: vector<uint, 2>, b flip: uint>
+
+    firrtl.connect %w, %c0_ui2 : !firrtl.uint, !firrtl.uint<2>
+    firrtl.connect %bov_a_1, %c0_ui2 : !firrtl.uint, !firrtl.uint<2>
+    firrtl.connect %bov_b, %c0_ui2 : !firrtl.uint, !firrtl.uint<2>
   }
   // CHECK-LABEL: @Ref
   // CHECK: out x: !firrtl.probe<uint<2>>
   // CHECK-SAME: out y: !firrtl.rwprobe<uint<2>>
   // CHECK: firrtl.ref.resolve %sub_x : !firrtl.probe<uint<2>>
   // CHECK: firrtl.ref.resolve %sub_y : !firrtl.rwprobe<uint<2>>
   firrtl.module @Ref(out %r : !firrtl.uint, out %s : !firrtl.uint) {
-    %sub_x, %sub_y = firrtl.instance sub @SubRef(out x: !firrtl.probe<uint>, out y: !firrtl.rwprobe<uint>)
+    %sub_x, %sub_y, %sub_bov_ref = firrtl.instance sub @SubRef(out x: !firrtl.probe<uint>, out y: !firrtl.rwprobe<uint>, out bov_ref : !firrtl.rwprobe<bundle<a: vector<uint, 2>, b : uint>>)
     %res_x = firrtl.ref.resolve %sub_x : !firrtl.probe<uint>
     %res_y = firrtl.ref.resolve %sub_y : !firrtl.rwprobe<uint>
     firrtl.connect %r, %res_x : !firrtl.uint, !firrtl.uint
     firrtl.connect %s, %res_y : !firrtl.uint, !firrtl.uint
+
+    // CHECK: !firrtl.rwprobe<bundle<a: vector<uint<2>, 2>, b: uint<2>>>
+    %read_bov = firrtl.ref.resolve %sub_bov_ref : !firrtl.rwprobe<bundle<a: vector<uint, 2>, b : uint>>
+    // CHECK: !firrtl.rwprobe<bundle<a: vector<uint<2>, 2>, b: uint<2>>>
+    %bov_ref_a = firrtl.ref.sub %sub_bov_ref[0] : !firrtl.rwprobe<bundle<a: vector<uint, 2>, b : uint>>
+    // CHECK: !firrtl.rwprobe<vector<uint<2>, 2>>
+    %bov_ref_a_1 = firrtl.ref.sub %bov_ref_a[1] : !firrtl.rwprobe<vector<uint, 2>>
+    // CHECK: !firrtl.rwprobe<bundle<a: vector<uint<2>, 2>, b: uint<2>>>
+    %bov_ref_b  = firrtl.ref.sub %sub_bov_ref[1] : !firrtl.rwprobe<bundle<a: vector<uint, 2>, b : uint>>
+
+    // CHECK: !firrtl.rwprobe<vector<uint<2>, 2>>
+    %bov_a = firrtl.ref.resolve %bov_ref_a : !firrtl.rwprobe<vector<uint,2>>
+    // CHECK: !firrtl.rwprobe<uint<2>>
+    %bov_a_1 = firrtl.ref.resolve %bov_ref_a_1 : !firrtl.rwprobe<uint>
+    // CHECK: !firrtl.rwprobe<uint<2>>
+    %bov_b = firrtl.ref.resolve %bov_ref_b : !firrtl.rwprobe<uint>
   }
 
   // CHECK-LABEL: @ForeignTypes