[mlir][dxsa] Add dcl_thread_group instruction

Example:
  dxsa.dcl_thread_group 8, 8, 1

Signed-off-by: Vladimir Shiryaev <tagolog@users.noreply.github.com>

Author: tagolog

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

@@ -158,4 +158,14 @@ def DXSA_DclTemps : DXSA_Op<"dcl_temps"> {
   let assemblyFormat = [{ $count attr-dict }];
 }
 
+def DXSA_DclThreadGroup : DXSA_Op<"dcl_thread_group"> {
+  let summary = "declare thread group size";
+  let arguments = (ins
+      ConfinedAttr<I32Attr, [IntPositive, IntMaxValue<1024>]>:$x,
+      ConfinedAttr<I32Attr, [IntPositive, IntMaxValue<1024>]>:$y,
+      ConfinedAttr<I32Attr, [IntPositive, IntMaxValue<64>]>:$z);
+  let assemblyFormat = [{ $x `,` $y `,` $z attr-dict }];
+  let hasVerifier = 1;
+}
+
 #endif // DXSA_OPS

mlir/lib/Dialect/DXSA/IR/DXSA.cpp

@@ -44,6 +44,18 @@ LogicalResult DclGlobalFlags::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// DclThreadGroup
+//===----------------------------------------------------------------------===//
+
+LogicalResult DclThreadGroup::verify() {
+  auto total = uint64_t(getX()) * uint64_t(getY()) * uint64_t(getZ());
+  if (total > 1024)
+    return emitOpError("thread group size x*y*z must be <= 1024, got ")
+           << total;
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//

mlir/lib/Target/DXSA/BinaryParser.cpp

@@ -516,6 +516,13 @@ class DXBuilder {
                                   builder.getI32IntegerAttr(count));
   }
 
+  Instruction buildDclThreadGroup(uint32_t x, uint32_t y, uint32_t z,
+                                  Location loc) {
+    return dxsa::DclThreadGroup::create(
+        builder, loc, builder.getI32IntegerAttr(x),
+        builder.getI32IntegerAttr(y), builder.getI32IntegerAttr(z));
+  }
+
 private:
   MLIRContext *context;
   ModuleOp module;
@@ -831,6 +838,41 @@ class Parser {
     return builder.buildDclTemps(count, loc);
   }
 
+  FailureOr<Instruction> parseDclThreadGroup(Location loc) {
+    auto parseDimension = [&](StringRef dimensionName,
+                              uint32_t maxValue) -> FailureOr<uint32_t> {
+      auto token = parseToken();
+      if (failed(token))
+        return failure();
+      auto value = *token;
+      if (value == 0 || value > maxValue) {
+        emitError(getLocation(), "thread group ")
+            << dimensionName << " dimension must be in [1, " << maxValue
+            << "], got " << value;
+        return failure();
+      }
+      return value;
+    };
+
+    auto x = parseDimension("x", 1024);
+    if (failed(x))
+      return failure();
+    auto y = parseDimension("y", 1024);
+    if (failed(y))
+      return failure();
+    auto z = parseDimension("z", 64);
+    if (failed(z))
+      return failure();
+
+    auto total = uint64_t(*x) * uint64_t(*y) * uint64_t(*z);
+    if (total > 1024) {
+      emitError(getLocation(), "thread group size x*y*z must be <= 1024, got ")
+          << total;
+      return failure();
+    }
+    return builder.buildDclThreadGroup(*x, *y, *z, loc);
+  }
+
   OptionalParseResult parseDclInstruction(uint32_t opcodeToken, Location loc,
                                           Instruction &out) {
     FailureOr<Instruction> result;
@@ -841,6 +883,9 @@ class Parser {
     case D3D10_SB_OPCODE_DCL_TEMPS:
       result = parseDclTemps(loc);
       break;
+    case D3D11_SB_OPCODE_DCL_THREAD_GROUP:
+      result = parseDclThreadGroup(loc);
+      break;
     default:
       return std::nullopt;
     }

mlir/test/Target/DXSA/dcl_thread_group.mlir

@@ -0,0 +1,10 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_thread_group.bin | FileCheck %s
+
+// CHECK:      module {
+// CHECK-NEXT:   dxsa.dcl_thread_group 1, 1, 1
+// CHECK-NEXT:   dxsa.dcl_thread_group 8, 8, 1
+// CHECK-NEXT:   dxsa.dcl_thread_group 32, 32, 1
+// CHECK-NEXT:   dxsa.dcl_thread_group 16, 1, 64
+// CHECK-NEXT:   dxsa.dcl_thread_group 1024, 1, 1
+// CHECK-NEXT:   dxsa.dcl_thread_group 1, 1024, 1
+// CHECK-NEXT: }

mlir/test/Target/DXSA/dcl_thread_group_invalid.mlir

@@ -0,0 +1,35 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// expected-error@+1 {{attribute 'x' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 1024}}
+dxsa.dcl_thread_group 0, 1, 1
+
+// -----
+
+// expected-error@+1 {{attribute 'x' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 1024}}
+dxsa.dcl_thread_group 1025, 1, 1
+
+// -----
+
+// expected-error@+1 {{attribute 'y' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 1024}}
+dxsa.dcl_thread_group 1, 0, 1
+
+// -----
+
+// expected-error@+1 {{attribute 'y' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 1024}}
+dxsa.dcl_thread_group 1, 1025, 1
+
+// -----
+
+// expected-error@+1 {{attribute 'z' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 64}}
+dxsa.dcl_thread_group 1, 1, 0
+
+// -----
+
+// expected-error@+1 {{attribute 'z' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 64}}
+dxsa.dcl_thread_group 1, 1, 65
+
+// -----
+
+// 64 * 8 * 4 == 2048
+// expected-error@+1 {{'dxsa.dcl_thread_group' op thread group size x*y*z must be <= 1024, got 2048}}
+dxsa.dcl_thread_group 64, 8, 4