[FIRRTL] Prerequisites for Layer-Associated Probes

Miscellaneous fixes to get layer-associated probes working.

Signed-off-by: Schuyler Eldridge <schuyler.eldridge@sifive.com>

include/circt/Dialect/FIRRTL/FIRRTLDeclarations.td

@@ -572,7 +572,10 @@ def RegResetOp : HardwareDeclOp<"regreset", [Forceable]> {
   let hasVerifier = 1;
 }
 
-def WireOp : HardwareDeclOp<"wire", [Forceable]> {
+def WireOp : HardwareDeclOp<"wire", [
+    Forceable,
+    DeclareOpInterfaceMethods<SymbolUserOpInterface>
+]> {
   let summary = "Define a new wire";
   let description = [{
     Declare a new wire:

include/circt/Dialect/FIRRTL/FIRRTLExpressions.td

@@ -1205,7 +1205,8 @@ def DoubleConstantOp : FIRRTLOp<"double", [Pure, ConstantLike]> {
 def RefCastOp : FIRRTLOp<"ref.cast",
     [HasCustomSSAName,
      Pure,
-     CompatibleRefTypes<"result","input">]> {
+     CompatibleRefTypes<"result","input">,
+     DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
   let summary = "Cast between compatible reference types";
   let description = [{
     Losslessly cast between compatible reference types.
@@ -1220,6 +1221,7 @@ def RefCastOp : FIRRTLOp<"ref.cast",
   let results = (outs RefType:$result);
 
   let hasFolder = 1;
+  let hasVerifier = 1;
 
   let assemblyFormat =
      "$input attr-dict `:` functional-type($input, $result)";

lib/Dialect/FIRRTL/FIRRTLOps.cpp

@@ -1455,6 +1455,17 @@ static LogicalResult verifyPortSymbolUses(FModuleLike module,
   // verify types in ports.
   for (size_t i = 0, e = module.getNumPorts(); i < e; ++i) {
     auto type = module.getPortType(i);
+
+    if (auto refType = dyn_cast<RefType>(type)) {
+      auto layer = refType.getLayer();
+      if (layer && !dyn_cast_or_null<LayerOp>(
+                       symbolTable.lookupSymbolIn(circuitOp, layer)))
+        return emitError(module.getPortLocation(i))
+               << "probe port '" << module.getPortName(i)
+               << "' is associated with layer '" << layer
+               << "', but this layer was not defined";
+    }
+
     auto classType = dyn_cast<ClassType>(type);
     if (!classType)
       continue;
@@ -3045,6 +3056,20 @@ void WireOp::getAsmResultNames(OpAsmSetValueNameFn setNameFn) {
 
 std::optional<size_t> WireOp::getTargetResultIndex() { return 0; }
 
+LogicalResult WireOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  auto circuitOp = getOperation()->getParentOfType<CircuitOp>();
+  auto refType = dyn_cast<RefType>(getType(0));
+  if (!refType)
+    return success();
+
+  auto layer = refType.getLayer();
+  if (layer && !symbolTable.lookupSymbolIn(circuitOp, layer))
+    return emitOpError() << "has a probe type associated with layer '" << layer
+                         << "', but this layer is not defined";
+
+  return success();
+}
+
 void ObjectOp::build(OpBuilder &builder, OperationState &state, ClassLike klass,
                      StringRef name) {
   build(builder, state, klass.getInstanceType(),
@@ -5678,7 +5703,7 @@ FIRRTLType RefSubOp::inferReturnType(ValueRange operands,
       return RefType::get(
           vectorType.getElementType().getConstType(
               vectorType.isConst() || vectorType.getElementType().isConst()),
-          refType.getForceable());
+          refType.getForceable(), refType.getLayer());
     return emitInferRetTypeError(loc, "out of range index '", fieldIdx,
                                  "' in RefType of vector type ", refType);
   }
@@ -5691,13 +5716,64 @@ FIRRTLType RefSubOp::inferReturnType(ValueRange operands,
     return RefType::get(bundleType.getElement(fieldIdx).type.getConstType(
                             bundleType.isConst() ||
                             bundleType.getElement(fieldIdx).type.isConst()),
-                        refType.getForceable());
+                        refType.getForceable(), refType.getLayer());
   }
 
   return emitInferRetTypeError(
       loc, "ref.sub op requires a RefType of vector or bundle base type");
 }
 
+LogicalResult RefCastOp::verify() {
+
+  SymbolRefAttr layer;
+  auto layerBlockOp =
+      dyn_cast<LayerBlockOp>((*this)->getBlock()->getParentOp());
+  if (layerBlockOp)
+    layer = layerBlockOp.getLayerName();
+
+  // The dest layer must be the same as the source layer or a parent of it.
+  SymbolRefAttr layerPointer = layer;
+  auto destLayer = getType().getLayer();
+  for (;;) {
+    if (layerPointer == destLayer)
+      break;
+
+    if (!layerPointer) {
+      if (!layer)
+        return emitOpError()
+               << "cannot cast to a layer from outside a layerblock";
+      auto diag = emitOpError()
+                  << "casts to a probe associated with layer '" << destLayer
+                  << "' from a layerblock associated with layer '" << layer
+                  << "'. The cast op must be in a layerblock associated with "
+                     "or a child of layer '"
+                  << destLayer << "'.";
+      return diag.attachNote(layerBlockOp.getLoc())
+             << "the layerblock was declared here";
+    }
+
+    if (layerPointer.getNestedReferences().empty()) {
+      layerPointer = {};
+      continue;
+    }
+    layerPointer =
+        SymbolRefAttr::get(layerPointer.getRootReference(),
+                           layerPointer.getNestedReferences().drop_back());
+  }
+
+  return success();
+}
+
+LogicalResult RefCastOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  auto circuitOp = getOperation()->getParentOfType<CircuitOp>();
+  auto layer = getType().getLayer();
+  if (layer && !symbolTable.lookupSymbolIn(circuitOp, layer))
+    return emitOpError() << "casts to a probe type associated with layer '"
+                         << layer << "', but this layer is not defined";
+
+  return success();
+}
+
 LogicalResult RWProbeOp::verifyInnerRefs(hw::InnerRefNamespace &ns) {
   auto targetRef = getTarget();
   if (targetRef.getModule() !=
@@ -5793,16 +5869,10 @@ LogicalResult LayerBlockOp::verify() {
       // Any value captured from the current layer block is fine.
       if (operand.getParentBlock() == body)
         continue;
-      // Capture of a non-base type, e.g., reference is illegal.
+      // Capture of a non-base type, e.g., reference, is allowed.
       FIRRTLBaseType baseType = dyn_cast<FIRRTLBaseType>(operand.getType());
-      if (!baseType) {
-        auto diag = emitOpError()
-                    << "captures an operand which is not a FIRRTL base type";
-        diag.attachNote(operand.getLoc()) << "operand is defined here";
-        diag.attachNote(op->getLoc()) << "operand is used here";
-        failed = true;
+      if (!baseType)
         return WalkResult::advance();
-      }
       // Capturing a non-passive type is illegal.
       if (!baseType.isPassive()) {
         auto diag = emitOpError()

lib/Dialect/FIRRTL/Transforms/LowerTypes.cpp

@@ -1432,7 +1432,10 @@ bool TypeLoweringVisitor::visitExpr(RefCastOp op) {
   auto clone = [&](const FlatBundleFieldEntry &field,
                    ArrayAttr attrs) -> Value {
     auto input = getSubWhatever(op.getInput(), field.index);
-    return builder->create<RefCastOp>(RefType::get(field.type), input);
+    return builder->create<RefCastOp>(RefType::get(field.type,
+                                                   op.getType().getForceable(),
+                                                   op.getType().getLayer()),
+                                      input);
   };
   return lowerProducer(op, clone);
 }

test/Dialect/FIRRTL/errors.mlir

@@ -1869,22 +1869,6 @@ firrtl.circuit "WrongLayerBlockNesting" {
 
 // -----
 
-// A layer block captures a type which is not a FIRRTL base type.
-firrtl.circuit "NonBaseTypeCapture" {
-  firrtl.layer @A bind {}
-  firrtl.module @NonBaseTypeCapture(in %in: !firrtl.uint<1>) {
-    // expected-note @below {{operand is defined here}}
-     %ref = firrtl.ref.send %in : !firrtl.uint<1>
-    // expected-error @below {{'firrtl.layerblock' op captures an operand which is not a FIRRTL base type}}
-    firrtl.layerblock @A {
-      // expected-note @below {{operand is used here}}
-      %b = firrtl.ref.resolve %ref : !firrtl.probe<uint<1>>
-    }
-  }
-}
-
-// -----
-
 // A layer block captures a non-passive type.
 firrtl.circuit "NonPassiveCapture" {
   firrtl.layer @A bind {}
@@ -1920,6 +1904,53 @@ firrtl.circuit "LayerBlockDrivesSinksOutside" {
   }
 }
 
+// -----
+
+firrtl.circuit "IllegalRefCastDestModule" {
+  firrtl.module @IllegalRefCastDestModule() {
+    %a = firrtl.wire : !firrtl.uint<1>
+    %0 = firrtl.ref.send %a : !firrtl.uint<1>
+    // expected-error @below {{cannot cast to a layer from outside a layerblock}}
+    %1 = firrtl.ref.cast %0 : (!firrtl.probe<uint<1>>) -> !firrtl.probe<uint<1>, @A>
+  }
+}
+
+// -----
+
+firrtl.circuit "IllegalRefCastDestLayerBlock" {
+  firrtl.layer @A bind {
+  }
+  firrtl.layer @B bind {
+  }
+  firrtl.module @IllegalRefCastDestLayerBlock() {
+    // expected-note @below {{the layerblock was declared here}}
+    firrtl.layerblock @A {
+      %a = firrtl.wire : !firrtl.uint<1>
+      %0 = firrtl.ref.send %a : !firrtl.uint<1>
+      // expected-error @below {{'firrtl.ref.cast' op casts to a probe associated with layer '@B' from a layerblock associated with layer '@A'.}}
+      %1 = firrtl.ref.cast %0 : (!firrtl.probe<uint<1>>) -> !firrtl.probe<uint<1>, @B>
+    }
+  }
+}
+
+// -----
+
+firrtl.circuit "InvalidProbeAssociationPort" {
+  // expected-error @below {{probe port 'a' is associated with layer '@B', but this layer was not defined}}
+  firrtl.module @InvalidProbeAssociationPort(out %a: !firrtl.probe<uint<1>, @B>) {
+  }
+}
+
+// -----
+
+firrtl.circuit "InvalidProbeAssociationWire" {
+  firrtl.module @InvalidProbeAssociationWire() {
+    // expected-error @below {{has a probe type associated with layer '@B', but this layer is not defined}}
+    %a = firrtl.wire : !firrtl.probe<uint<1>, @B>
+  }
+}
+
+
 // -----
 
 firrtl.circuit "RWProbeRemote" {