[mlir][dxsa] Add dcl_function_table

Author: asavonic

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

@@ -734,4 +734,24 @@ def DXSA_DclFunctionBody : DXSA_Op<"dcl_function_body"> {
   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 I32Attr:$index, DenseI32ArrayAttr:$functions);
+  let assemblyFormat = "$index `,` `<` `functions` `=` $functions `>` attr-dict";
+}
+
 #endif // DXSA_OPS

mlir/lib/Target/DXSA/BinaryParser.cpp

@@ -639,6 +639,13 @@ class DXBuilder {
     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;
@@ -1184,6 +1191,24 @@ class Parser {
     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;
@@ -1233,6 +1258,9 @@ class Parser {
     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;
     }

mlir/test/Target/DXSA/dcl_function_table.mlir

@@ -0,0 +1,4 @@
+// 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]>