[FIRRTL] Add Layer Association to Probes

include/circt/Dialect/FIRRTL/FIRRTLTypesImpl.td

@@ -391,7 +391,8 @@ def RefImpl : FIRRTLImplType<"Ref",
   }];
   let parameters = (ins TypeParameter<"::circt::firrtl::FIRRTLBaseType",
                                       "Type of reference target">:$type,
-                        "bool":$forceable);
+                        "bool":$forceable,
+                        OptionalParameter<"::mlir::SymbolRefAttr">:$layer);
   let genAccessors = true;
   let genStorageClass = true;
   let genVerifyDecl = true;
@@ -400,7 +401,8 @@ def RefImpl : FIRRTLImplType<"Ref",
   let skipDefaultBuilders = true;
   let builders = [
     TypeBuilderWithInferredContext<(ins "::circt::firrtl::FIRRTLBaseType":$type,
-                                        CArg<"bool", "false">:$forceable)>
+                                        CArg<"bool", "false">:$forceable,
+                                        CArg<"::mlir::SymbolRefAttr", "{}">:$layer)>
   ];
 
   let extraClassDeclaration = [{

include/circt/Dialect/FIRRTL/FIRRTLUtils.h

@@ -234,7 +234,8 @@ inline FIRRTLType mapBaseType(FIRRTLType type,
   return TypeSwitch<FIRRTLType, FIRRTLType>(type)
       .Case<FIRRTLBaseType>([&](auto base) { return fn(base); })
       .Case<RefType>([&](auto ref) {
-        return RefType::get(fn(ref.getType()), ref.getForceable());
+        return RefType::get(fn(ref.getType()), ref.getForceable(),
+                            ref.getLayer());
       });
 }
 
@@ -250,7 +251,7 @@ mapBaseTypeNullable(FIRRTLType type,
         auto result = fn(ref.getType());
         if (!result)
           return {};
-        return RefType::get(result, ref.getForceable());
+        return RefType::get(result, ref.getForceable(), ref.getLayer());
       });
 }
 

lib/Dialect/FIRRTL/FIRRTLTypes.cpp

@@ -105,11 +105,13 @@ static LogicalResult customTypePrinter(Type type, AsmPrinter &os) {
         printNestedType(vectorType.getElementType(), os);
         os << ", " << vectorType.getNumElements() << '>';
       })
-      .Case<RefType>([&](auto refType) {
+      .Case<RefType>([&](RefType refType) {
         if (refType.getForceable())
           os << "rw";
         os << "probe<";
         printNestedType(refType.getType(), os);
+        if (auto layer = refType.getLayer())
+          os << ", " << layer;
         os << '>';
       })
       .Case<StringType>([&](auto stringType) { os << "string"; })
@@ -329,31 +331,43 @@ static OptionalParseResult customTypeParser(AsmParser &parser, StringRef name,
   // For now, support both firrtl.ref and firrtl.probe.
   if (name.equals("ref") || name.equals("probe")) {
     FIRRTLBaseType type;
+    SymbolRefAttr layer;
     // Don't pass `isConst` to `parseNestedBaseType since `ref` can point to
     // either `const` or non-`const` types
-    if (parser.parseLess() || parseNestedBaseType(type, parser) ||
-        parser.parseGreater())
+    if (parser.parseLess() || parseNestedBaseType(type, parser))
+      return failure();
+    if (parser.parseOptionalComma().succeeded())
+      if (parser.parseOptionalAttribute(layer).value())
+        return parser.emitError(parser.getNameLoc(),
+                                "expected symbol reference");
+    if (parser.parseGreater())
       return failure();
 
     if (failed(RefType::verify(
             [&]() { return parser.emitError(parser.getNameLoc()); }, type,
-            false)))
+            false, layer)))
       return failure();
 
-    return result = RefType::get(type, false), success();
+    return result = RefType::get(type, false, layer), success();
   }
   if (name.equals("rwprobe")) {
     FIRRTLBaseType type;
-    if (parser.parseLess() || parseNestedBaseType(type, parser) ||
-        parser.parseGreater())
+    SymbolRefAttr layer;
+    if (parser.parseLess() || parseNestedBaseType(type, parser))
+      return failure();
+    if (parser.parseOptionalComma().succeeded())
+      if (parser.parseOptionalAttribute(layer).value())
+        return parser.emitError(parser.getNameLoc(),
+                                "expected symbol reference");
+    if (parser.parseGreater())
       return failure();
 
     if (failed(RefType::verify(
-            [&]() { return parser.emitError(parser.getNameLoc()); }, type,
-            true)))
+            [&]() { return parser.emitError(parser.getNameLoc()); }, type, true,
+            layer)))
       return failure();
 
-    return result = RefType::get(type, true), success();
+    return result = RefType::get(type, true, layer), success();
   }
   if (name.equals("class")) {
     if (isConst)
@@ -2400,12 +2414,14 @@ BaseTypeAliasType::getIndexAndSubfieldID(uint64_t fieldID) const {
 // RefType
 //===----------------------------------------------------------------------===//
 
-auto RefType::get(FIRRTLBaseType type, bool forceable) -> RefType {
-  return Base::get(type.getContext(), type, forceable);
+auto RefType::get(FIRRTLBaseType type, bool forceable, SymbolRefAttr layer)
+    -> RefType {
+  return Base::get(type.getContext(), type, forceable, layer);
 }
 
 auto RefType::verify(function_ref<InFlightDiagnostic()> emitErrorFn,
-                     FIRRTLBaseType base, bool forceable) -> LogicalResult {
+                     FIRRTLBaseType base, bool forceable, SymbolRefAttr layer)
+    -> LogicalResult {
   if (!base.isPassive())
     return emitErrorFn() << "reference base type must be passive";
   if (forceable && base.containsConst())

lib/Dialect/FIRRTL/Transforms/DropConst.cpp

@@ -35,7 +35,8 @@ static Type convertType(Type type) {
 
   if (auto refType = type_dyn_cast<RefType>(type)) {
     if (auto converted = convertType(refType.getType()))
-      return RefType::get(converted, refType.getForceable());
+      return RefType::get(converted, refType.getForceable(),
+                          refType.getLayer());
   }
 
   return {};

lib/Dialect/FIRRTL/Transforms/VBToBV.cpp

@@ -174,7 +174,8 @@ RefType Visitor::convertType(RefType type) {
   auto cached = typeMap.lookup(type);
   if (cached)
     return type_cast<RefType>(cached);
-  auto converted = RefType::get(convertType(type.getType()));
+  auto converted = RefType::get(convertType(type.getType()),
+                                type.getForceable(), type.getLayer());
   typeMap.insert({type, converted});
   return converted;
 }

test/Dialect/FIRRTL/round-trip.mlir

@@ -44,4 +44,16 @@ firrtl.module @Foo(in %clock: !firrtl.clock) {
   firrtl.strictconnect %inst_clock, %clock : !firrtl.clock
 }
 
+firrtl.layer @LayerA bind {
+  firrtl.layer @LayerB bind {}
+}
+
+// CHECK-LABEL: firrtl.module @Layers
+// CHECK-SAME:    out %a: !firrtl.probe<uint<1>, @LayerA>
+// CHECK-SAME:    out %b: !firrtl.rwprobe<uint<1>, @LayerA::@LayerB>
+firrtl.module @Layers(
+  out %a: !firrtl.probe<uint<1>, @LayerA>,
+  out %b: !firrtl.rwprobe<uint<1>, @LayerA::@LayerB>
+) {}
+
 }