[mlir][dxsa] Add dcl_indexable_temp instruction

Example:
  dxsa.dcl_indexable_temp x0[23], 2
  dxsa.dcl_indexable_temp x7[4096], 4

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

Author: tagolog

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

@@ -146,4 +146,14 @@ def DXSA_DclTemps : DXSA_Op<"dcl_temps"> {
   let assemblyFormat = [{ $count attr-dict }];
 }
 
+def DXSA_DclIndexableTemp : DXSA_Op<"dcl_indexable_temp"> {
+  let summary = "declare a temporary register";
+  let arguments = (ins I32Attr:$register_index,
+      ConfinedAttr<I32Attr, [IntPositive, IntMaxValue<4096>]>:$register_count,
+      ConfinedAttr<I32Attr, [IntMinValue<1>, IntMaxValue<4>]>:$num_components);
+  let assemblyFormat = [{
+    custom<IndexableTempRegister>($register_index) `[` $register_count `]` `,` $num_components attr-dict
+  }];
+}
+
 #endif // DXSA_OPS

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

@@ -34,6 +34,54 @@ void DXSADialect::initialize() {
       >();
 }
 
+//===----------------------------------------------------------------------===//
+// Custom assembly format helpers
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// Parse a register identifier of the form `<prefix><N>` (e.g. `x0`, `r3`).
+/// The whole identifier must be a single token (no whitespace between
+/// prefix and number).
+ParseResult parseRegisterName(OpAsmParser &parser, StringRef registerPrefix,
+                              IntegerAttr &registerIndex) {
+  StringRef name;
+  auto loc = parser.getCurrentLocation();
+  if (parser.parseKeyword(&name))
+    return failure();
+  if (!name.consume_front(registerPrefix))
+    return parser.emitError(loc)
+           << "expected register prefix '" << registerPrefix << "'";
+  unsigned value;
+  if (name.getAsInteger(10, value))
+    return parser.emitError(loc) << "expected integer after register prefix '"
+                                 << registerPrefix << "'";
+  registerIndex = parser.getBuilder().getI32IntegerAttr(value);
+  return success();
+}
+
+/// Print a register identifier of the form `<prefix><N>`.
+void printRegisterName(OpAsmPrinter &printer, StringRef registerPrefix,
+                       IntegerAttr registerIndex) {
+  printer << registerPrefix << registerIndex.getInt();
+}
+
+//===----------------------------------------------------------------------===//
+// Custom parser/printer for indexable temp registers (dcl_indexable_temp op)
+//===----------------------------------------------------------------------===//
+
+ParseResult parseIndexableTempRegister(OpAsmParser &parser,
+                                       IntegerAttr &registerIndex) {
+  return parseRegisterName(parser, "x", registerIndex);
+}
+
+void printIndexableTempRegister(OpAsmPrinter &printer, Operation *op,
+                                IntegerAttr registerIndex) {
+  printRegisterName(printer, "x", registerIndex);
+}
+
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//

mlir/lib/Target/DXSA/BinaryParser.cpp

@@ -515,6 +515,15 @@ class DXBuilder {
                                   builder.getI32IntegerAttr(count));
   }
 
+  Instruction buildDclIndexableTemp(uint32_t registerIndex,
+                                    uint32_t registerCount,
+                                    uint32_t numComponents, Location loc) {
+    return dxsa::DclIndexableTemp::create(
+        builder, loc, builder.getI32IntegerAttr(registerIndex),
+        builder.getI32IntegerAttr(registerCount),
+        builder.getI32IntegerAttr(numComponents));
+  }
+
 private:
   MLIRContext *context;
   ModuleOp module;
@@ -853,6 +862,40 @@ class Parser {
     return builder.buildDclTemps(count, loc);
   }
 
+  FailureOr<Instruction> parseDclIndexableTemp(Location loc) {
+    auto registerIndexToken = parseToken();
+    if (failed(registerIndexToken))
+      return failure();
+    auto registerCountToken = parseToken();
+    if (failed(registerCountToken))
+      return failure();
+    auto numComponentsToken = parseToken();
+    if (failed(numComponentsToken))
+      return failure();
+
+    auto registerCount = *registerCountToken;
+    if (registerCount == 0) {
+      emitError(getLocation(), "indexable temp array size cannot be zero");
+      return failure();
+    }
+    if (registerCount > 4096) {
+      emitError(getLocation(), "invalid indexable temp array size: ")
+          << registerCount << " (max 4096)";
+      return failure();
+    }
+
+    auto numComponents = *numComponentsToken;
+    if (numComponents < 1 || numComponents > 4) {
+      emitError(getLocation(), "indexable temp num components must be in 1..4, "
+                               "got ")
+          << numComponents;
+      return failure();
+    }
+
+    return builder.buildDclIndexableTemp(*registerIndexToken, registerCount,
+                                         numComponents, loc);
+  }
+
   OptionalParseResult parseDclInstruction(uint32_t opcodeToken, Location loc,
                                           Instruction &out) {
     FailureOr<Instruction> result;
@@ -863,6 +906,9 @@ class Parser {
     case D3D10_SB_OPCODE_DCL_TEMPS:
       result = parseDclTemps(loc);
       break;
+    case D3D10_SB_OPCODE_DCL_INDEXABLE_TEMP:
+      result = parseDclIndexableTemp(loc);
+      break;
     default:
       return std::nullopt;
     }

mlir/test/Target/DXSA/dcl_indexable_temp.mlir

@@ -0,0 +1,8 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_indexable_temp.bin | FileCheck %s
+
+// CHECK:      module {
+// CHECK-NEXT:   dxsa.dcl_indexable_temp x0[1], 1
+// CHECK-NEXT:   dxsa.dcl_indexable_temp x1[23], 2
+// CHECK-NEXT:   dxsa.dcl_indexable_temp x2[16], 4
+// CHECK-NEXT:   dxsa.dcl_indexable_temp x7[4096], 4
+// CHECK-NEXT: }

mlir/test/Target/DXSA/dcl_indexable_temp_invalid.mlir

@@ -0,0 +1,29 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// expected-error@+1 {{attribute 'register_count' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 4096}}
+dxsa.dcl_indexable_temp x0[0], 2
+
+// -----
+
+// expected-error@+1 {{attribute 'register_count' failed to satisfy constraint: 32-bit signless integer attribute whose value is positive whose maximum value is 4096}}
+dxsa.dcl_indexable_temp x0[4097], 2
+
+// -----
+
+// expected-error@+1 {{attribute 'num_components' failed to satisfy constraint: 32-bit signless integer attribute whose minimum value is 1 whose maximum value is 4}}
+dxsa.dcl_indexable_temp x0[16], 0
+
+// -----
+
+// expected-error@+1 {{attribute 'num_components' failed to satisfy constraint: 32-bit signless integer attribute whose minimum value is 1 whose maximum value is 4}}
+dxsa.dcl_indexable_temp x0[16], 5
+
+// -----
+
+// expected-error@+1 {{expected register prefix 'x'}}
+dxsa.dcl_indexable_temp r0[16], 4
+
+// -----
+
+// expected-error@+1 {{expected integer after register prefix 'x'}}
+dxsa.dcl_indexable_temp xfoo[16], 4