[mlir][dxsa] Add dcl_immediate_constant_buffer instruction (variant of customdata)

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

@@ -1156,6 +1156,29 @@ def DXSA_DclTgsmStructured : DXSA_Op<"dcl_tgsm_structured"> {
   let hasVerifier = 1;
 }
 
+def DXSA_DclImmediateConstantBuffer : DXSA_Op<"dcl_immediate_constant_buffer"> {
+  let summary = "declares an immediate constant buffer";
+  let description = [{
+    The `dxsa.dcl_immediate_constant_buffer` operation declares
+    an immediate constant buffer holding an array of 32-bit floating-point
+    values grouped into 4-component vectors.
+
+    Example:
+
+    ```mlir
+    dxsa.dcl_immediate_constant_buffer
+                   [
+                     1.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00,
+                     0.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00,
+                     0.000000e+00, 0.000000e+00, 1.000000e+00, 0.000000e+00,
+                     0.000000e+00, 0.000000e+00, 0.000000e+00, 1.000000e+00
+                   ]
+    ```
+  }];
+  let arguments = (ins DenseF32ArrayAttr:$values);
+  let assemblyFormat = "$values attr-dict";
+}
+
 def DXSA_DclConstantBuffer : DXSA_Op<"dcl_constant_buffer"> {
   let summary = "declares a constant buffer";
   let description = [{

mlir/lib/Target/DXSA/BinaryParser.cpp

@@ -892,6 +892,11 @@ class DXBuilder {
         optionalToAttr(space), accessPattern);
   }
 
+  Instruction buildDclImmediateConstantBuffer(ArrayRef<float> values,
+                                              Location loc) {
+    return dxsa::DclImmediateConstantBuffer::create(builder, loc, values);
+  }
+
   Instruction buildDclSampler(uint32_t id, std::optional<uint32_t> lbound,
                               std::optional<uint32_t> ubound,
                               std::optional<uint32_t> space,
@@ -1830,6 +1835,21 @@ class Parser {
                                           *structCount, loc);
   }
 
+  FailureOr<Instruction> parseDclImmediateConstantBuffer(uint32_t numTokens,
+                                                         Location loc) {
+    uint32_t numDataTokens = numTokens >= 2 ? numTokens - 2 : 0;
+
+    auto dataTokens = parseTokens(numDataTokens);
+    FAILURE_IF_FAILED(dataTokens);
+
+    SmallVector<float, 16> values;
+    values.reserve(numDataTokens);
+    for (uint32_t token : *dataTokens)
+      values.push_back(llvm::bit_cast<float>(token));
+
+    return builder.buildDclImmediateConstantBuffer(values, loc);
+  }
+
   FailureOr<Instruction> parseDclConstantBuffer(uint32_t opcodeToken,
                                                 Location loc) {
     auto rawAccessPattern =
@@ -2296,12 +2316,29 @@ class Parser {
     uint32_t opcode = DECODE_D3D10_SB_OPCODE_TYPE(*opcodeToken0);
 
     // CUSTOMDATA carries its total token count (>= 2) in token1.
-    // Just set the instruction length for the unknown fallback.
     if (opcode == D3D10_SB_OPCODE_CUSTOMDATA) {
       auto numTokensToken = parseToken();
       FAILURE_IF_FAILED(numTokensToken);
       instructionLengthInTokens = std::max(*numTokensToken, 2u);
-      return emitError(getLocation(), "customdata is not supported yet");
+
+      if (failed(verifyInstructionLengthFitsBufferSize(
+              beginOffset, instructionLengthInTokens)))
+        return failure();
+
+      auto customDataClass = DECODE_D3D10_SB_CUSTOMDATA_CLASS(*opcodeToken0);
+      switch (customDataClass) {
+      case D3D10_SB_CUSTOMDATA_DCL_IMMEDIATE_CONSTANT_BUFFER: {
+        auto result =
+            parseDclImmediateConstantBuffer(*numTokensToken, getLocation());
+        FAILURE_IF_FAILED(result);
+        if (failed(verifyInstructionLength(beginOffset,
+                                           instructionLengthInTokens)))
+          return failure();
+        return *result;
+      }
+      default:
+        return emitError(getLocation(), "customdata is not supported yet");
+      }
     }
 
     instructionLengthInTokens = std::max(
@@ -2329,8 +2366,7 @@ class Parser {
     auto parseResult =
         parseDclInstruction(*opcodeToken0, getLocation(), dclInstruction);
     if (parseResult.has_value()) {
-      if (failed(*parseResult))
-        return failure();
+      FAILURE_IF_FAILED(*parseResult);
       if (failed(
               verifyInstructionLength(beginOffset, instructionLengthInTokens)))
         return failure();

mlir/test/Target/DXSA/dcl_immediate_constant_buffer.test

@@ -0,0 +1,21 @@
+// RUN: mlir-translate --import-dxsa-hex %s | FileCheck %s
+// RUN: mlir-translate --import-dxsa-hex %s | mlir-opt --verify-roundtrip
+
+// Check customdata dcl_immediate_constant_buffer parsing.
+
+// customdata may be confused with program header, so begin our module with program header.
+// CHECK:      dxsa.module pixel_shader 5 0 {
+0x00000050, 0x00000014,
+// CHECK-NEXT:   dxsa.dcl_immediate_constant_buffer
+// CHECK-SAME:     [
+// CHECK-SAME:       1.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00,
+// CHECK-SAME:       0.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00,
+// CHECK-SAME:       0.000000e+00, 0.000000e+00, 1.000000e+00, 0.000000e+00,
+// CHECK-SAME:       0.000000e+00, 0.000000e+00, 0.000000e+00, 1.000000e+00
+// CHECK-SAME:     ]
+0x00001835, 0x00000012,
+0x3F800000, 0x00000000, 0x00000000, 0x00000000,
+0x00000000, 0x3F800000, 0x00000000, 0x00000000,
+0x00000000, 0x00000000, 0x3F800000, 0x00000000,
+0x00000000, 0x00000000, 0x00000000, 0x3F800000
+// CHECK-NEXT: }