[FIRRTL] Add domain support to instance choice

Author: rwy7

lib/Dialect/FIRRTL/FIRRTLOps.cpp

@@ -2924,6 +2924,13 @@ void InstanceChoiceOp::build(
     portAnnotationsAttr = builder.getArrayAttr(portAnnotations);
   }
 
+  // Create the domain info attribute.
+  ArrayAttr domainInfoAttr = defaultModule.getDomainInfoAttr();
+  if (domainInfoAttr.empty()) {
+    domainInfoAttr = builder.getArrayAttr(SmallVector<Attribute, 16>(
+        resultTypes.size(), builder.getArrayAttr({})));
+  }
+
   // Gather the module & case names.
   SmallVector<Attribute> moduleNames, caseNames;
   moduleNames.push_back(SymbolRefAttr::get(defaultModule.getModuleNameAttr()));
@@ -2934,14 +2941,14 @@ void InstanceChoiceOp::build(
     moduleNames.push_back(SymbolRefAttr::get(caseModule.getModuleNameAttr()));
   }
 
-  return build(
-      builder, result, resultTypes, builder.getArrayAttr(moduleNames),
-      builder.getArrayAttr(caseNames), builder.getStringAttr(name),
-      NameKindEnumAttr::get(builder.getContext(), nameKind),
-      defaultModule.getPortDirectionsAttr(), defaultModule.getPortNamesAttr(),
-      defaultModule.getDomainInfoAttr(), builder.getArrayAttr(annotations),
-      portAnnotationsAttr, defaultModule.getLayersAttr(),
-      innerSym ? hw::InnerSymAttr::get(innerSym) : hw::InnerSymAttr());
+  return build(builder, result, resultTypes, builder.getArrayAttr(moduleNames),
+               builder.getArrayAttr(caseNames), builder.getStringAttr(name),
+               NameKindEnumAttr::get(builder.getContext(), nameKind),
+               defaultModule.getPortDirectionsAttr(),
+               defaultModule.getPortNamesAttr(), domainInfoAttr,
+               builder.getArrayAttr(annotations), portAnnotationsAttr,
+               defaultModule.getLayersAttr(),
+               innerSym ? hw::InnerSymAttr::get(innerSym) : hw::InnerSymAttr());
 }
 
 std::optional<size_t> InstanceChoiceOp::getTargetResultIndex() {
@@ -3067,7 +3074,7 @@ ParseResult InstanceChoiceOp::parse(OpAsmParser &parser,
   }
 
   if (parseModulePorts(parser, /*hasSSAIdentifiers=*/false,
-                       /*supportsSymbols=*/false, /*supportsDomains=*/false,
+                       /*supportsSymbols=*/false, /*supportsDomains=*/true,
                        entryArgs, portDirections, portNames, portTypes,
                        portAnnotations, portSyms, portLocs, domains))
     return failure();

test/Dialect/FIRRTL/round-trip.mlir

@@ -262,4 +262,18 @@ firrtl.module @DomainDefine(
   firrtl.domain.define %y, %x
 }
 
+firrtl.module private @DefaultTargetWithDomain(in %A: !firrtl.domain of @ClockDomain, in %clock: !firrtl.clock domains [%A]) {}
+firrtl.module private @FPGATargetWithDomain(in %A: !firrtl.domain of @ClockDomain, in %clock: !firrtl.clock domains [%A]) {}
+firrtl.module private @ASICTargetWithDomain(in %A: !firrtl.domain of @ClockDomain, in %clock: !firrtl.clock domains [%A]) {}
+
+// CHECK-LABEL: firrtl.module @InstanceChoiceWithDomain
+firrtl.module @InstanceChoiceWithDomain(in %A: !firrtl.domain of @ClockDomain, in %clock: !firrtl.clock domains [%A]) {
+  // CHECK:      %inst_A, %inst_clock = firrtl.instance_choice inst interesting_name @DefaultTargetWithDomain alternatives @Platform
+  // CHECK-SAME:   { @FPGA -> @FPGATargetWithDomain, @ASIC -> @ASICTargetWithDomain } (in A: !firrtl.domain of @ClockDomain, in clock: !firrtl.clock domains [A])
+  %inst_A, %inst_clock = firrtl.instance_choice inst interesting_name @DefaultTargetWithDomain alternatives @Platform
+    { @FPGA -> @FPGATargetWithDomain, @ASIC -> @ASICTargetWithDomain } (in A: !firrtl.domain of @ClockDomain, in clock: !firrtl.clock domains [A])
+  firrtl.matchingconnect %inst_clock, %clock : !firrtl.clock
+  firrtl.domain.define %inst_A, %A
+}
+
 }