[FIRRTL] Add Domain Info to WireOp

include/circt/Dialect/FIRRTL/FIRRTLDeclarations.td

@@ -632,7 +632,8 @@ def WireOp : HardwareDeclOp<"wire", [
   let arguments = (ins StrAttr:$name, NameKindAttr:$nameKind,
                        AnnotationArrayAttr:$annotations,
                        OptionalAttr<InnerSymAttr>:$inner_sym,
-                       UnitAttr:$forceable); // ReferenceKinds
+                       UnitAttr:$forceable,
+                       OptionalAttr<DomainInfoAttr>:$domainInfo); // ReferenceKinds
   let results = (outs AnyType:$result, Optional<RefType>:$ref);
 
   let hasCanonicalizer = true;
@@ -652,13 +653,17 @@ def WireOp : HardwareDeclOp<"wire", [
                    "::circt::firrtl::NameKindEnumAttr":$nameKind,
                    "::mlir::ArrayAttr":$annotations,
                    "::circt::hw::InnerSymAttr":$innerSym,
-                   "::mlir::UnitAttr":$forceable), [{
+                   "::mlir::UnitAttr":$forceable,
+                   "::mlir::ArrayAttr":$domainInfo), [{
       $_state.addAttribute(getNameAttrName($_state.name), name);
       $_state.addAttribute(getNameKindAttrName($_state.name), nameKind);
       $_state.addAttribute(getAnnotationsAttrName($_state.name), annotations);
       if (innerSym) {
         $_state.addAttribute(getInnerSymAttrName($_state.name), innerSym);
       }
+      if (domainInfo) {
+        $_state.addAttribute(getDomainInfoAttrName($_state.name), domainInfo);
+      }
       $_state.addTypes(elementType);
       if (forceable) {
         $_state.addAttribute(getForceableAttrName($_state.name), forceable);
@@ -671,33 +676,38 @@ def WireOp : HardwareDeclOp<"wire", [
                    "::circt::firrtl::NameKindEnum":$nameKind,
                    "::mlir::ArrayAttr":$annotations,
                    "::circt::hw::InnerSymAttr":$innerSym,
-                   "bool":$forceable), [{
+                   "bool":$forceable,
+                   "::mlir::ArrayAttr":$domainInfo), [{
       return build($_builder, $_state, elementType,
                    $_builder.getStringAttr(name),
                    ::circt::firrtl::NameKindEnumAttr::get(odsBuilder.getContext(), nameKind),
                    annotations, innerSym,
-                   forceable ? $_builder.getUnitAttr() : nullptr);
+                   forceable ? $_builder.getUnitAttr() : nullptr, domainInfo);
     }]>,
     OpBuilder<(ins "::mlir::Type":$elementType,
                       CArg<"StringRef", "{}">:$name,
                       CArg<"NameKindEnum", "NameKindEnum::DroppableName">:$nameKind,
                       CArg<"ArrayRef<Attribute>","{}">:$annotations,
                       CArg<"StringAttr", "StringAttr()">:$innerSym,
-                      CArg<"bool", "false">:$forceable), [{
+                      CArg<"bool", "false">:$forceable,
+                      CArg<"ArrayAttr", "ArrayAttr()">:$domainInfo), [{
       return build($_builder, $_state, elementType,
                    name, nameKind,
                    $_builder.getArrayAttr(annotations),
                    innerSym ? hw::InnerSymAttr::get(innerSym) : hw::InnerSymAttr(),
-                   forceable);
+                   forceable,
+                   domainInfo);
     }]>,
     OpBuilder<(ins "::mlir::Type":$elementType, "StringRef":$name,
                    "NameKindEnum":$nameKind, "::mlir::ArrayAttr":$annotations,
                    CArg<"StringAttr", "StringAttr()">:$innerSym,
-                   CArg<"bool", "false">:$forceable), [{
+                   CArg<"bool", "false">:$forceable,
+                   CArg<"ArrayAttr", "ArrayAttr()">:$domainInfo), [{
       return build($_builder, $_state, elementType,
                    name, nameKind, annotations,
                    innerSym ? hw::InnerSymAttr::get(innerSym) : hw::InnerSymAttr(),
-                   forceable);
+                   forceable,
+                   domainInfo);
     }]>
   ];
 

lib/Dialect/FIRRTL/FIRRTLOps.cpp

@@ -4290,6 +4290,18 @@ static FlatSymbolRefAttr getDomainTypeName(Value value) {
       return getDomainTypeNameOfResult(instance, result.getResultNumber());
     if (auto anonDomain = dyn_cast<DomainCreateAnonOp>(op))
       return anonDomain.getDomainAttr();
+    if (auto wireOp = dyn_cast<WireOp>(op)) {
+      std::optional<ArrayAttr> domainInfo = wireOp.getDomainInfo();
+
+      if (!domainInfo || domainInfo->empty()) {
+        return {};
+      }
+
+      if (domainInfo->size() != 1) {
+        return {};
+      }
+      return dyn_cast<FlatSymbolRefAttr>((*domainInfo)[0]);
+    }
     return {};
   }
 
@@ -4306,19 +4318,6 @@ LogicalResult DomainDefineOp::verify() {
   auto dst = getDest();
   auto src = getSrc();
 
-  // As wires cannot have domain information, don't do any checking when a wire
-  // is involved.  This weakens the verification.
-  //
-  // TOOD: Remove this by adding Domain Info to wires [1].
-  //
-  // [1] https://github.com/llvm/circt/issues/9398
-  if (auto *srcDefOp = src.getDefiningOp())
-    if (isa<WireOp>(srcDefOp))
-      return success();
-  if (auto *dstDefOp = dst.getDefiningOp())
-    if (isa<WireOp>(dstDefOp))
-      return success();
-
   auto dstDomain = getDomainTypeName(dst);
   if (!dstDomain)
     return emitError("could not determine domain-type of destination");

lib/Dialect/FIRRTL/Transforms/IMDeadCodeElim.cpp

@@ -623,29 +623,28 @@ void IMDeadCodeElimPass::rewriteModuleSignature(FModuleOp module) {
   LLVM_DEBUG(llvm::dbgs() << "Prune ports of module: " << module.getName()
                           << "\n");
 
-  auto replaceInstanceResultWithWire = [&](ImplicitLocOpBuilder &builder,
-                                           unsigned index,
-                                           InstanceOp instance) {
-    auto result = instance.getResult(index);
-    if (isAssumedDead(result)) {
-      // If the result is dead, replace the result with an unrealized conversion
-      // cast which works as a dummy placeholder.
-      auto wire =
-          mlir::UnrealizedConversionCastOp::create(
-              builder, ArrayRef<Type>{result.getType()}, ArrayRef<Value>{})
-              ->getResult(0);
-      result.replaceAllUsesWith(wire);
-      return;
-    }
+  auto replaceInstanceResultWithWire =
+      [&](ImplicitLocOpBuilder &builder, unsigned index, InstanceOp instance) {
+        auto result = instance.getResult(index);
+        if (isAssumedDead(result)) {
+          // If the result is dead, replace the result with an unrealized
+          // conversion cast which works as a dummy placeholder.
+          auto wire =
+              mlir::UnrealizedConversionCastOp::create(
+                  builder, ArrayRef<Type>{result.getType()}, ArrayRef<Value>{})
+                  ->getResult(0);
+          result.replaceAllUsesWith(wire);
+          return;
+        }
 
-    Value wire = WireOp::create(builder, result.getType()).getResult();
-    result.replaceAllUsesWith(wire);
-    // If a module port is dead but its instance result is alive, the port
-    // is used as a temporary wire so make sure that a replaced wire is
-    // putted into `liveSet`.
-    liveElements.erase(result);
-    liveElements.insert(wire);
-  };
+        Value wire = WireOp::create(builder, result.getType()).getResult();
+        result.replaceAllUsesWith(wire);
+        // If a module port is dead but its instance result is alive, the port
+        // is used as a temporary wire so make sure that a replaced wire is
+        // putted into `liveSet`.
+        liveElements.erase(result);
+        liveElements.insert(wire);
+      };
 
   // First, delete dead instances.
   for (auto *use : llvm::make_early_inc_range(instanceGraphNode->uses())) {

lib/Dialect/FIRRTL/Transforms/InferResets.cpp

@@ -1771,7 +1771,7 @@ LogicalResult InferResetsPass::implementFullReset(FModuleOp module,
         auto wireOp = WireOp::create(
             builder, nodeOp.getResult().getType(), nodeOp.getNameAttr(),
             nodeOp.getNameKindAttr(), nodeOp.getAnnotationsAttr(),
-            nodeOp.getInnerSymAttr(), nodeOp.getForceableAttr());
+            nodeOp.getInnerSymAttr(), nodeOp.getForceableAttr(), ArrayAttr());
         // Don't delete the node, since it might be in use in worklists.
         nodeOp->replaceAllUsesWith(wireOp);
         nodeOp->removeAttr(nodeOp.getInnerSymAttrName());

lib/Dialect/FIRRTL/Transforms/LowerOpenAggs.cpp

@@ -665,19 +665,27 @@ LogicalResult Visitor::visitDecl(WireOp op) {
     return mlir::emitError(op.getLoc())
            << "annotations on open aggregates not handled yet";
 
+  ArrayAttr domainAttr;
+  if (auto domainInfo = op.getDomainInfo())
+    domainAttr = *domainInfo;
+
   // Create the new HW wire.
   if (mappings.hwType)
     hwOnlyAggMap[op.getResult()] =
         WireOp::create(builder, mappings.hwType, op.getName(), op.getNameKind(),
-                       op.getAnnotations(), mappings.newSym, op.getForceable())
+                       op.getAnnotations(), mappings.newSym, op.getForceable(),
+                       domainAttr)
             .getResult();
 
   // Create the non-HW wires.  Non-HW wire names are always droppable.
   for (auto &[type, fieldID, _, suffix] : mappings.fields)
     nonHWValues[FieldRef(op.getResult(), fieldID)] =
         WireOp::create(builder, type,
                        builder.getStringAttr(Twine(op.getName()) + suffix),
-                       NameKindEnum::DroppableName)
+                       NameKindEnum::DroppableName,
+                       /*annotations=*/ArrayRef<Attribute>{},
+                       /*innerSym=*/StringAttr(),
+                       /*forceable=*/false, domainAttr)
             .getResult();
 
   for (auto fieldID : mappings.mapToNullInteriors)

lib/Dialect/FIRRTL/Transforms/LowerTypes.cpp

@@ -1299,11 +1299,16 @@ bool TypeLoweringVisitor::visitDecl(WireOp op) {
   if (op.isForceable())
     return false;
 
+  // Get the domain info from the original wire.
+  auto domain = op.getDomainInfoAttr();
+
   auto clone = [&](const FlatBundleFieldEntry &field,
                    ArrayAttr attrs) -> Value {
     return WireOp::create(*builder,
                           mapLoweredType(op.getDataRaw().getType(), field.type),
-                          "", NameKindEnum::DroppableName, attrs, StringAttr{})
+                          "", NameKindEnum::DroppableName, attrs,
+                          /*innerSym=*/StringAttr{}, /*forceable=*/false,
+                          /*domainInfo=*/domain)
         .getResult();
   };
   return lowerProducer(op, clone);

lib/Dialect/FIRRTL/Transforms/ModuleInliner.cpp

@@ -864,13 +864,23 @@ void Inliner::mapPortsToWires(StringRef prefix, InliningLevel &il,
       newAnnotations.push_back(anno.getAttr());
     }
 
+    ArrayAttr domainAttr;
+    if (auto domains = portInfo[i].domains) {
+      if (auto symRef = dyn_cast<FlatSymbolRefAttr>(domains))
+        // Domain-type ports: FlatSymbolRefAttr -> [@DomainName]
+        domainAttr = ArrayAttr::get(context, {symRef});
+      else if (auto array = dyn_cast<ArrayAttr>(domains))
+        // Non-domain ports with domain associations: [0] -> [[0]]
+        domainAttr = ArrayAttr::get(context, {array});
+    }
+
     Value wire =
         WireOp::create(
             il.mic.b, target.getLoc(), type,
             StringAttr::get(context, (prefix + portInfo[i].getName())),
             NameKindEnumAttr::get(context, NameKindEnum::DroppableName),
             ArrayAttr::get(context, newAnnotations), newSymAttr,
-            /*forceable=*/UnitAttr{})
+            /*forceable=*/UnitAttr{}, domainAttr)
             .getResult();
     il.wires.push_back(wire);
     mapper.map(arg, wire);

test/Dialect/FIRRTL/errors.mlir

@@ -3244,3 +3244,54 @@ firrtl.circuit "UndefinedDomainInAnonDomain" {
     %0 = firrtl.domain.anon : !firrtl.domain of @Foo
   }
 }
+
+// -----
+
+// Unable to determine domain type of wire without domain info.
+firrtl.circuit "WireWithoutDomainInfo" {
+  firrtl.domain @ClockDomain
+  firrtl.domain @PowerDomain
+  firrtl.module @WireWithoutDomainInfo(
+    in %in: !firrtl.domain of @ClockDomain,
+    out %out: !firrtl.domain of @PowerDomain
+  ) {
+    %w = firrtl.wire : !firrtl.domain
+    // expected-error @below {{could not determine domain-type of destination}}
+    firrtl.domain.define %w, %in
+    firrtl.domain.define %out, %w
+  }
+}
+
+// -----
+
+// Wire with domainInfo should fail when source has mismatched domain type.
+firrtl.circuit "WireWithMismatchedDomainInfo" {
+  firrtl.domain @ClockDomain
+  firrtl.domain @PowerDomain
+  firrtl.module @WireWithMismatchedDomainInfo(
+    in %in: !firrtl.domain of @ClockDomain,
+    out %out: !firrtl.domain of @PowerDomain
+  ) {
+    %w = firrtl.wire {domainInfo = [@PowerDomain]} : !firrtl.domain
+    // expected-error @below {{source domain type @ClockDomain does not match destination domain type @PowerDomain}}
+    firrtl.domain.define %w, %in
+    firrtl.domain.define %out, %w
+  }
+}
+
+// ----
+
+// Wire with domainInfo should fail when source has mismatched domain type.
+firrtl.circuit "WireWithMismatchedDomainInfo" {
+  firrtl.domain @ClockDomain
+  firrtl.domain @PowerDomain
+  firrtl.module @WireWithMismatchedDomainInfo(
+    in %in: !firrtl.domain of @PowerDomain,
+    out %out: !firrtl.domain of @PowerDomain
+  ) {
+    %w = firrtl.wire {domainInfo = [@ClockDomain]} : !firrtl.domain
+    // expected-error @below {{source domain type @PowerDomain does not match destination domain type @ClockDomain}}
+    firrtl.domain.define %w, %in
+    firrtl.domain.define %out, %w
+  }
+}

test/Dialect/FIRRTL/inliner.mlir

@@ -1440,3 +1440,46 @@ firrtl.circuit "RemoveNonLocalFromLocal" {
     firrtl.instance bar sym @sym {annotations = [{circt.nonlocal = @dutNLA, class = "circt.tracker", id = distinct[0]<>}]} @Bar()
   }
 }
+
+// -----
+
+// Test that domain info is preserved when inlining a module with domain-type ports.
+firrtl.circuit "InlineDomainTypePorts" {
+  firrtl.domain @ClockDomain
+  // CHECK-LABEL: firrtl.module @InlineDomainTypePorts
+  firrtl.module @InlineDomainTypePorts(
+    in %A: !firrtl.domain of @ClockDomain
+  ) {
+    // CHECK: %child_A = firrtl.wire {domainInfo = [@ClockDomain]} : !firrtl.domain
+    // CHECK: firrtl.domain.define %child_A, %A
+    %child_A = firrtl.instance child @Child(in A: !firrtl.domain of @ClockDomain)
+    firrtl.domain.define %child_A, %A
+  }
+  firrtl.module private @Child(
+    in %A: !firrtl.domain of @ClockDomain
+  ) attributes {annotations = [{class = "firrtl.passes.InlineAnnotation"}]} {
+  }
+}
+
+// -----
+
+// Test that domain associations are preserved when inlining a module with non-domain ports that have domain associations.
+firrtl.circuit "InlineDomainAssociations" {
+  firrtl.domain @ClockDomain
+  // CHECK-LABEL: firrtl.module @InlineDomainAssociations
+  firrtl.module @InlineDomainAssociations(
+    in %A: !firrtl.domain of @ClockDomain,
+    in %clock: !firrtl.clock domains [%A]
+  ) {
+    // CHECK: %child_A = firrtl.wire {domainInfo = [@ClockDomain]} : !firrtl.domain
+    // CHECK: %child_clock = firrtl.wire {domainInfo = {{\[}}[0 : ui32]]} : !firrtl.clock
+    %child_A, %child_clock = firrtl.instance child @Child(in A: !firrtl.domain of @ClockDomain, in clock: !firrtl.clock domains [A])
+    firrtl.domain.define %child_A, %A
+    firrtl.matchingconnect %child_clock, %clock : !firrtl.clock
+  }
+  firrtl.module private @Child(
+    in %A: !firrtl.domain of @ClockDomain,
+    in %clock: !firrtl.clock domains [%A]
+  ) attributes {annotations = [{class = "firrtl.passes.InlineAnnotation"}]} {
+  }
+}