[mlir][dxsa] Add dcl_temps instruction

Example:
  dxsa.dcl_temps {count = 1 : i32}

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

Author: tagolog

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

@@ -151,4 +151,11 @@ def DXSA_DclGlobalFlags : DXSA_Op<"dcl_global_flags"> {
   let hasVerifier = 1;
 }
 
+def DXSA_DclTemps : DXSA_Op<"dcl_temps"> {
+  let summary = "declares the number of temporary registers";
+  let arguments = (ins ConfinedAttr<I32Attr,
+      [IntPositive, IntMaxValue<4096>]>:$count);
+  let assemblyFormat = [{ $count attr-dict }];
+}
+
 #endif // DXSA_OPS

mlir/lib/Target/DXSA/BinaryParser.cpp

@@ -507,11 +507,15 @@ class DXBuilder {
 
   Instruction buildDclGlobalFlags(uint32_t flags, FileLineColLoc loc) {
     auto flagsAttr = dxsa::GlobalFlagsAttr::get(
-        builder.getContext(),
-        static_cast<dxsa::GlobalFlags>(flags));
+        builder.getContext(), static_cast<dxsa::GlobalFlags>(flags));
     return dxsa::DclGlobalFlags::create(builder, loc, flagsAttr);
   }
 
+  Instruction buildDclTemps(uint32_t count, FileLineColLoc loc) {
+    return dxsa::DclTemps::create(builder, loc,
+                                  builder.getI32IntegerAttr(count));
+  }
+
 private:
   MLIRContext *context;
   ModuleOp module;
@@ -804,6 +808,42 @@ class Parser {
                                 getLocation());
   }
 
+  FailureOr<Instruction> parseDclGlobalFlags(uint32_t opcodeToken,
+                                             FileLineColLoc loc) {
+    return builder.buildDclGlobalFlags(
+        DECODE_D3D10_SB_GLOBAL_FLAGS(opcodeToken), loc);
+  }
+
+  FailureOr<Instruction> parseDclTemps(FileLineColLoc loc) {
+    auto countToken = parseToken();
+    if (failed(countToken))
+      return failure();
+    auto count = *countToken;
+    if (count == 0) {
+      emitError(getLocation(), "temp register count cannot be zero");
+      return failure();
+    }
+    if (count > 4096) {
+      emitError(getLocation(), "invalid temp register count: ")
+          << count << " (max 4096)";
+      return failure();
+    }
+    return builder.buildDclTemps(count, loc);
+  }
+
+  std::optional<FailureOr<Instruction>>
+  tryParseAsDclInstruction(uint32_t opcode, uint32_t opcodeToken,
+                           FileLineColLoc loc) {
+    switch (opcode) {
+    case D3D10_SB_OPCODE_DCL_GLOBAL_FLAGS:
+      return parseDclGlobalFlags(opcodeToken, loc);
+    case D3D10_SB_OPCODE_DCL_TEMPS:
+      return parseDclTemps(loc);
+    default:
+      return std::nullopt;
+    }
+  }
+
   FailureOr<Instruction> parseInstruction() {
     size_t beginOffset = currentTokenOffset;
     Token token = parseToken();
@@ -829,12 +869,15 @@ class Parser {
       return failure();
     }
 
-    uint32_t opcodeToken = *token;
+    auto opcodeToken = *token;
 
-    if (opcode == D3D10_SB_OPCODE_DCL_GLOBAL_FLAGS) {
+    auto dclResult = tryParseAsDclInstruction(opcode, opcodeToken, loc);
+    if (dclResult) {
+      if (failed(*dclResult))
+        return failure();
       if (failed(verifyInstructionLength(beginOffset, length)))
         return failure();
-      return builder.buildDclGlobalFlags(DECODE_D3D10_SB_GLOBAL_FLAGS(opcodeToken), loc);
+      return *dclResult;
     }
 
     unsigned numOperands = instrInfo[opcode].numOperands;

mlir/test/Target/DXSA/dcl_temps.mlir

@@ -0,0 +1,7 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_temps.bin | FileCheck %s
+
+// CHECK:      module {
+// CHECK-NEXT:   dxsa.dcl_temps 1
+// CHECK-NEXT:   dxsa.dcl_temps 16
+// CHECK-NEXT:   dxsa.dcl_temps 4096
+// CHECK-NEXT: }

mlir/test/Target/DXSA/dcl_temps_invalid.mlir

@@ -0,0 +1,9 @@
+// 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 4096}}
+dxsa.dcl_temps 0
+
+// -----
+
+// expected-error@+1 {{attribute 'count' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 4096}}
+dxsa.dcl_temps 4097