[mlir][dxsa] Add dcl_tgsm_raw instruction

Example:
  dxsa.dcl_tgsm_raw <type = thread_group_shared_memory, components = 0, index = [0]>, 40

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

Author: tagolog

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

@@ -475,13 +475,7 @@ def DXSA_InlineOperandAttr : AttrDef<DXSADialect, "InlineOperand"> {
       "uint32_t":$components,
       OptionalParameter<"::mlir::dxsa::ComponentMaskAttr">:$mask,
       OptionalParameter<"::mlir::DenseI64ArrayAttr">:$index);
-  let assemblyFormat = [{
-    `<` `type` `=` $type
-        `,` `components` `=` $components
-        (`,` `mask` `=` $mask^)?
-        (`,` `index` `=` $index^)?
-    `>`
-  }];
+  let assemblyFormat = "`<` struct(params) `>`";
 }
 
 def DXSA_DclGlobalFlags : DXSA_Op<"dcl_global_flags"> {
@@ -719,4 +713,27 @@ def DXSA_DclOutput : DXSA_Op<"dcl_output"> {
   let assemblyFormat = "$operand attr-dict";
 }
 
+def DXSA_DclTgsmRaw : DXSA_Op<"dcl_tgsm_raw"> {
+  let summary = "declares a reference to a region of shared memory space";
+  let description = [{
+    The `dxsa.dcl_tgsm_raw` operation declares a reference to a region
+    of shared memory space available to the Compute Shader's thread group.
+
+    The `$operand` is the `g#` register being declared.
+    The `$byte_count$ is the size of the block of untyped shared memory;
+    must be a multiple of 4.
+
+    Example:
+
+    ```mlir
+    dxsa.dcl_tgsm_raw <type = thread_group_shared_memory, components = 0, index = [0]>, 1024
+    ```
+  }];
+  let arguments = (ins DXSA_InlineOperandAttr:$operand,
+                       ConfinedAttr<I32Attr,
+                           [IntPositive, IntMaxValue<32768>]>:$byte_count);
+  let assemblyFormat = "$operand `,` $byte_count attr-dict";
+  let hasVerifier = 1;
+}
+
 #endif // DXSA_OPS

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

@@ -41,6 +41,17 @@ void DXSADialect::initialize() {
 #define GET_OP_CLASSES
 #include "mlir/Dialect/DXSA/IR/DXSAOps.cpp.inc"
 
+//===----------------------------------------------------------------------===//
+// Op verifiers
+//===----------------------------------------------------------------------===//
+
+LogicalResult DclTgsmRaw::verify() {
+  auto byteCount = getByteCount();
+  if (byteCount % 4 != 0)
+    return emitOpError("byte count must be a multiple of 4, got ") << byteCount;
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd attribute method definitions
 //===----------------------------------------------------------------------===//

mlir/lib/Target/DXSA/BinaryParser.cpp

@@ -635,6 +635,12 @@ class DXBuilder {
     return dxsa::DclOutput::create(builder, loc, operand);
   }
 
+  Instruction buildDclTgsmRaw(dxsa::InlineOperandAttr operand,
+                              uint32_t byteCount, Location loc) {
+    return dxsa::DclTgsmRaw::create(builder, loc, operand,
+                                    builder.getI32IntegerAttr(byteCount));
+  }
+
 private:
   MLIRContext *context;
   ModuleOp module;
@@ -1174,6 +1180,14 @@ class Parser {
     return builder.buildDclOutput(*operand, loc);
   }
 
+  FailureOr<Instruction> parseDclTgsmRaw(Location loc) {
+    auto operand = parseInlineOperand();
+    FAILURE_IF_FAILED(operand);
+    auto byteCount = parseToken();
+    FAILURE_IF_FAILED(byteCount);
+    return builder.buildDclTgsmRaw(*operand, *byteCount, loc);
+  }
+
   OptionalParseResult parseDclInstruction(uint32_t opcodeToken, Location loc,
                                           Instruction &out) {
     FailureOr<Instruction> result;
@@ -1220,6 +1234,9 @@ class Parser {
     case D3D10_SB_OPCODE_DCL_OUTPUT:
       result = parseDclOutput(loc);
       break;
+    case D3D11_SB_OPCODE_DCL_THREAD_GROUP_SHARED_MEMORY_RAW:
+      result = parseDclTgsmRaw(loc);
+      break;
     default:
       return std::nullopt;
     }

mlir/test/Target/DXSA/dcl_tgsm_raw.mlir

@@ -0,0 +1,5 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_tgsm_raw.bin | FileCheck %s
+
+// CHECK:      module {
+// CHECK-NEXT:   dxsa.dcl_tgsm_raw <type = thread_group_shared_memory, components = 0, index = [0]>, 40
+// CHECK-NEXT: }

mlir/test/Target/DXSA/dcl_tgsm_raw_invalid.mlir

@@ -0,0 +1,14 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// expected-error@+1 {{'dxsa.dcl_tgsm_raw' op byte count must be a multiple of 4, got 42}}
+dxsa.dcl_tgsm_raw <type = thread_group_shared_memory, components = 0, index = [0]>, 42
+
+// -----
+
+// expected-error@+1 {{attribute 'byte_count' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 32768}}
+dxsa.dcl_tgsm_raw <type = thread_group_shared_memory, components = 0, index = [0]>, 0
+
+// -----
+
+// expected-error@+1 {{attribute 'byte_count' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 32768}}
+dxsa.dcl_tgsm_raw <type = thread_group_shared_memory, components = 0, index = [0]>, 32769