[mlir][dxsa] Add dcl_function_body and dcl_function_table

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

@@ -719,4 +719,41 @@ def DXSA_DclOutput : DXSA_Op<"dcl_output"> {
   let assemblyFormat = "$operand attr-dict";
 }
 
+def DXSA_DclFunctionBody : DXSA_Op<"dcl_function_body"> {
+  let summary = "declares a function body number";
+  let description = [{
+    The `dxsa.dcl_function_body` declares a unique function body whose index
+    will appear later in the program at `label #` or `dcl_function_table`.
+
+    Example:
+    ```mlir
+    dxsa.dcl_function_body 0
+    ```
+  }];
+  let arguments = (ins ConfinedAttr<I32Attr, [IntNonNegative]>:$index);
+  let assemblyFormat = "$index attr-dict";
+}
+
+def DXSA_DclFunctionTable : DXSA_Op<"dcl_function_table"> {
+  let summary = "declares a function table";
+  let description = [{
+    Declare a function table as a set of function bodies that have
+    been declared earlier.
+
+    This is like a C++ vtable except there is an entry per call site
+    for an interface instead of per method.
+
+    Example:
+    ```mlir
+    dxsa.dcl_function_body 0
+    dxsa.dcl_function_body 1
+    dxsa.dcl_function_table 0, <functions = [0, 1]>
+    ```
+  }];
+  let arguments = (ins ConfinedAttr<I32Attr, [IntNonNegative]>:$index,
+                       DenseI32ArrayAttr:$functions);
+  let assemblyFormat = "$index `,` `<` `functions` `=` $functions `>` attr-dict";
+  let hasVerifier = 1;
+}
+
 #endif // DXSA_OPS

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

@@ -41,6 +41,16 @@ void DXSADialect::initialize() {
 #define GET_OP_CLASSES
 #include "mlir/Dialect/DXSA/IR/DXSAOps.cpp.inc"
 
+LogicalResult DclFunctionTable::verify() {
+  for (int32_t functionIndex : getFunctions()) {
+    if (functionIndex < 0) {
+      return emitOpError("function body index must not be negative, got ")
+             << functionIndex;
+    }
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd attribute method definitions
 //===----------------------------------------------------------------------===//

mlir/lib/Target/DXSA/BinaryParser.cpp

@@ -635,6 +635,17 @@ class DXBuilder {
     return dxsa::DclOutput::create(builder, loc, operand);
   }
 
+  Instruction buildDclFunctionBody(uint32_t index, Location loc) {
+    return dxsa::DclFunctionBody::create(builder, loc, index);
+  }
+
+  Instruction buildDclFunctionTable(uint32_t index, ArrayRef<int32_t> functions,
+                                    Location loc) {
+    auto *ctx = builder.getContext();
+    return dxsa::DclFunctionTable::create(
+        builder, loc, index, DenseI32ArrayAttr::get(ctx, functions));
+  }
+
 private:
   MLIRContext *context;
   ModuleOp module;
@@ -1174,6 +1185,30 @@ class Parser {
     return builder.buildDclOutput(*operand, loc);
   }
 
+  FailureOr<Instruction> parseDclFunctionBody(Location loc) {
+    auto index = parseToken();
+    FAILURE_IF_FAILED(index);
+    return builder.buildDclFunctionBody(*index, loc);
+  }
+
+  FailureOr<Instruction> parseDclFunctionTable(Location loc) {
+    auto index = parseToken();
+    FAILURE_IF_FAILED(index);
+
+    auto numFunctions = parseToken();
+    FAILURE_IF_FAILED(numFunctions);
+
+    SmallVector<int32_t, 16> functions;
+    functions.resize(*numFunctions);
+    for (uint32_t i = 0; i < *numFunctions; ++i) {
+      auto functionIndex = parseToken();
+      FAILURE_IF_FAILED(functionIndex);
+      functions[i] = *functionIndex;
+    }
+
+    return builder.buildDclFunctionTable(*index, functions, loc);
+  }
+
   OptionalParseResult parseDclInstruction(uint32_t opcodeToken, Location loc,
                                           Instruction &out) {
     FailureOr<Instruction> result;
@@ -1220,6 +1255,12 @@ class Parser {
     case D3D10_SB_OPCODE_DCL_OUTPUT:
       result = parseDclOutput(loc);
       break;
+    case D3D11_SB_OPCODE_DCL_FUNCTION_BODY:
+      result = parseDclFunctionBody(loc);
+      break;
+    case D3D11_SB_OPCODE_DCL_FUNCTION_TABLE:
+      result = parseDclFunctionTable(loc);
+      break;
     default:
       return std::nullopt;
     }
@@ -1229,6 +1270,18 @@ class Parser {
     return success();
   }
 
+  FailureOr<std::optional<uint32_t>> parseExtendedLength(uint32_t opcode) {
+    std::optional<uint32_t> result;
+    if (opcode != D3D11_SB_OPCODE_DCL_FUNCTION_TABLE) {
+      return success(result);
+    }
+
+    Token lengthToken = parseToken();
+    FAILURE_IF_FAILED(lengthToken);
+    result = *lengthToken;
+    return success(result);
+  }
+
   FailureOr<Instruction> parseInstruction() {
     size_t beginOffset = currentTokenOffset;
     Token token = parseToken();
@@ -1243,11 +1296,16 @@ class Parser {
     modifier.saturate = DECODE_IS_D3D10_SB_INSTRUCTION_SATURATE_ENABLED(*token);
 
     uint32_t length = DECODE_D3D10_SB_TOKENIZED_INSTRUCTION_LENGTH(*token);
+    if (DECODE_IS_D3D10_SB_OPCODE_EXTENDED(*token)) {
+      auto extLength = parseExtendedLength(opcode);
+      FAILURE_IF_FAILED(extLength);
+      if (*extLength) {
+        length = extLength->value();
+      }
+    }
 
     // TODO: extended instructions:
     // BOOL b51PlusShader =
-    // BOOL bExtended = DECODE_IS_D3D10_SB_OPCODE_EXTENDED(Token)
-    // ...
 
     if (opcode >= D3D10_SB_NUM_OPCODES) {
       emitError(getLocation(), "unknown opcode");

mlir/test/Target/DXSA/dcl_function_body.mlir

@@ -0,0 +1,4 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_function_body.bin | FileCheck %s
+//
+// CHECK-LABEL: module
+// CHECK-NEXT: dxsa.dcl_function_body 1

mlir/test/Target/DXSA/dcl_function_body_invalid.mlir

@@ -0,0 +1,4 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// expected-error@+1 {{'dxsa.dcl_function_body' op attribute 'index' failed to satisfy constraint: 32-bit signless integer attribute whose value is non-negative}}
+dxsa.dcl_function_body -1

mlir/test/Target/DXSA/dcl_function_table.mlir

@@ -0,0 +1,7 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_function_table.bin | FileCheck %s
+
+// CHECK-LABEL: module
+// CHECK-NEXT: dxsa.dcl_function_table 1, <functions = [0, 1]>
+
+// Test extended instruction length:
+// CHECK-NEXT: dxsa.dcl_function_table 2, <functions = [0, 1]>

mlir/test/Target/DXSA/dcl_function_table_invalid.mlir

@@ -0,0 +1,9 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// expected-error@+1 {{'dxsa.dcl_function_table' op attribute 'index' failed to satisfy constraint: 32-bit signless integer attribute whose value is non-negative}}
+dxsa.dcl_function_table -1, <functions = [0, 1]>
+
+// -----
+
+// expected-error@+1 {{'dxsa.dcl_function_table' op function body index must not be negative, got -1}}
+dxsa.dcl_function_table 1, <functions = [0, -1]>