[mlir][dxsa] Add dcl_gs_instance_count instruction

mlir/include/mlir/Dialect/DXSA/IR/DXSAOps.td

@@ -838,4 +838,23 @@ def DXSA_DclHsForkPhaseInstanceCount :
   let assemblyFormat = "$count attr-dict";
 }
 
+def DXSA_DclGsInstanceCount : DXSA_Op<"dcl_gs_instance_count"> {
+  let summary = "declares instance count for the geometry shader";
+  let description = [{
+    The `dxsa.dcl_gs_instance_count` operation declares declares instance count
+    for the geometry shader.
+
+    The `$count` must be in [1, 32].
+
+    Example:
+
+    ```mlir
+    dxsa.dcl_gs_instance_count 4
+    ```
+  }];
+  let arguments = (ins ConfinedAttr<I32Attr,
+      [IntPositive, IntMaxValue<32>]>:$count);
+  let assemblyFormat = [{ $count attr-dict }];
+}
+
 #endif // DXSA_OPS

mlir/lib/Target/DXSA/BinaryParser.cpp

@@ -586,6 +586,11 @@ class DXBuilder {
     return dxsa::DclInputPrimitive::create(builder, loc, inputPrimitiveAttr);
   }
 
+  Instruction buildDclGsInstanceCount(uint32_t count, Location loc) {
+    return dxsa::DclGsInstanceCount::create(builder, loc,
+                                            builder.getI32IntegerAttr(count));
+  }
+
   Instruction buildDclInputPs(dxsa::InterpolationMode interpolationMode,
                               Operand operand, Location loc) {
     auto interpolationModeAttr = dxsa::InterpolationModeAttr::get(
@@ -1085,6 +1090,17 @@ class Parser {
     return builder.buildDclInputPrimitive(*inputPrimitive, loc);
   }
 
+  FailureOr<Instruction> parseDclGsInstanceCount(Location loc) {
+    auto countToken = parseToken();
+    FAILURE_IF_FAILED(countToken);
+    auto count = *countToken;
+    if (count == 0)
+      return emitError(loc, "instance count cannot be zero");
+    if (count > 32)
+      return emitError(loc, "instance count must be <= 32, got ") << count;
+    return builder.buildDclGsInstanceCount(count, loc);
+  }
+
   FailureOr<dxsa::InterpolationMode>
   parseInterpolationMode(uint32_t opcodeToken, Location loc) {
     auto rawInterpolationMode =
@@ -1282,6 +1298,9 @@ class Parser {
     case D3D10_SB_OPCODE_DCL_GS_INPUT_PRIMITIVE:
       result = parseDclInputPrimitive(opcodeToken, loc);
       break;
+    case D3D11_SB_OPCODE_DCL_GS_INSTANCE_COUNT:
+      result = parseDclGsInstanceCount(loc);
+      break;
     case D3D10_SB_OPCODE_DCL_INPUT_PS:
       result = parseDclInputPs(opcodeToken, loc);
       break;

mlir/test/Target/DXSA/dcl_gs_instance_count.mlir

@@ -0,0 +1,6 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_gs_instance_count.bin | FileCheck %s
+
+// CHECK:      module {
+// CHECK-NEXT:   dxsa.dcl_gs_instance_count 1
+// CHECK-NEXT:   dxsa.dcl_gs_instance_count 32
+// CHECK-NEXT: }

mlir/test/Target/DXSA/dcl_gs_instance_count_invalid.mlir

@@ -0,0 +1,14 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// expected-error@+1 {{attribute 'count' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 32}}
+dxsa.dcl_gs_instance_count -1
+
+// -----
+
+// expected-error@+1 {{attribute 'count' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 32}}
+dxsa.dcl_gs_instance_count 0
+
+// -----
+
+// expected-error@+1 {{attribute 'count' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 32}}
+dxsa.dcl_gs_instance_count 33