BinaryParser.cpp

//===----------- BinaryParser.cpp - Parse DXSA binary to MLIR -------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "mlir/Target/DXSA/BinaryParser.h"
#include "mlir/Dialect/DXSA/IR/DXSA.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Location.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/bit.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/Endian.h"
#include "llvm/Support/LogicalResult.h"

#include <optional>

// d3d12TokenizedProgramFormat.hpp references the `UINT` type in some DECODE_*
// macros. Mirror the Windows SDK alias (`typedef unsigned int UINT`) to use the
// header without modification.
using UINT = unsigned int;
#include "d3d12TokenizedProgramFormat.hpp" // NOLINT

#define DEBUG_TYPE "import-dxsa-bin"

using namespace mlir;
using namespace llvm;

#define FAILURE_IF_FAILED(RES)                                                 \
  if (failed(RES))                                                             \
    return failure();

enum OpcodeClass {
  D3D10_SB_FLOAT_OP,
  D3D10_SB_INT_OP,
  D3D10_SB_UINT_OP,
  D3D10_SB_BIT_OP,
  D3D10_SB_FLOW_OP,
  D3D10_SB_TEX_OP,
  D3D10_SB_DCL_OP,
  D3D11_SB_ATOMIC_OP,
  D3D11_SB_MEM_OP,
  D3D11_SB_DOUBLE_OP,
  D3D11_SB_FLOAT_TO_DOUBLE_OP,
  D3D11_SB_DOUBLE_TO_FLOAT_OP,
  D3D11_SB_DEBUG_OP,
};

struct InstructionInfo {
  unsigned numOperands;
  StringRef name;
  OpcodeClass opClass;
  uint32_t precisionFromOutMask;
};

static void initInstructionInfo(MutableArrayRef<InstructionInfo> instructions) {
#define SET(OpCode, Name, NumOperands, PrecMask, OpClass)                      \
  instructions[OpCode] = InstructionInfo{NumOperands, Name, OpClass, PrecMask};
  // clang-format off
  SET(D3D10_SB_OPCODE_ADD, "add", 3, 0x06, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_AND, "and", 3, 0x06, D3D10_SB_BIT_OP);
  SET(D3D10_SB_OPCODE_BREAK, "break", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_BREAKC, "breakc", 1, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_CALL, "call", 1, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_CALLC, "callc", 2, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_CONTINUE, "continue", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_CONTINUEC, "continuec", 1, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_CASE, "case", 1, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_CUT, "cut", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_DEFAULT, "default", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_DISCARD, "discard", 1, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_DIV, "div", 3, 0x06, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_DP2, "dp2", 3, 0x06, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_DP3, "dp3", 3, 0x06, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_DP4, "dp4", 3, 0x06, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_ELSE, "else", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_EMIT, "emit", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_EMITTHENCUT, "emit_then_cut", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_ENDIF, "endif", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_ENDLOOP, "endloop", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_ENDSWITCH, "endswitch", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_EQ, "eq", 3, 0x00, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_EXP, "exp", 2, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_FRC, "frc", 2, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_FTOI, "ftoi", 2, 0x00, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_FTOU, "ftou", 2, 0x00, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_GE, "ge", 3, 0x00, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_DERIV_RTX, "deriv_rtx", 2, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_DERIV_RTY, "deriv_rty", 2, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_IADD, "iadd", 3, 0x06, D3D10_SB_INT_OP);
  SET(D3D10_SB_OPCODE_IF, "if", 1, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_IEQ, "ieq", 3, 0x00, D3D10_SB_INT_OP);
  SET(D3D10_SB_OPCODE_IGE, "ige", 3, 0x00, D3D10_SB_INT_OP);
  SET(D3D10_SB_OPCODE_ILT, "ilt", 3, 0x00, D3D10_SB_INT_OP);
  SET(D3D10_SB_OPCODE_IMAD, "imad", 4, 0x0e, D3D10_SB_INT_OP);
  SET(D3D10_SB_OPCODE_IMAX, "imax", 3, 0x06, D3D10_SB_INT_OP);
  SET(D3D10_SB_OPCODE_IMIN, "imin", 3, 0x06, D3D10_SB_INT_OP);
  SET(D3D10_SB_OPCODE_IMUL, "imul", 4, 0x0c, D3D10_SB_INT_OP);
  SET(D3D10_SB_OPCODE_INE, "ine", 3, 0x00, D3D10_SB_INT_OP);
  SET(D3D10_SB_OPCODE_INEG, "ineg", 2, 0x02, D3D10_SB_INT_OP);
  SET(D3D10_SB_OPCODE_ISHL, "ishl", 3, 0x02, D3D10_SB_INT_OP);
  SET(D3D10_SB_OPCODE_ISHR, "ishr", 3, 0x02, D3D10_SB_INT_OP);
  SET(D3D10_SB_OPCODE_ITOF, "itof", 2, 0x00, D3D10_SB_INT_OP);
  SET(D3D10_SB_OPCODE_LABEL, "label", 1, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_LD, "ld", 3, 0x00, D3D10_SB_TEX_OP);
  SET(D3D10_SB_OPCODE_LD_MS, "ldms", 4, 0x00, D3D10_SB_TEX_OP);
  SET(D3D10_SB_OPCODE_LOG, "log", 2, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_LOOP, "loop", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_LT, "lt", 3, 0x00, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_MAD, "mad", 4, 0x0e, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_MAX, "max", 3, 0x06, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_MIN, "min", 3, 0x06, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_MOV, "mov", 2, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_MOVC, "movc", 4, 0x0c, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_MUL, "mul", 3, 0x06, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_NE, "ne", 3, 0x00, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_NOP, "nop", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_NOT, "not", 2, 0x02, D3D10_SB_BIT_OP);
  SET(D3D10_SB_OPCODE_OR, "or", 3, 0x06, D3D10_SB_BIT_OP);
  SET(D3D10_SB_OPCODE_RESINFO, "resinfo", 3, 0x00, D3D10_SB_TEX_OP);
  SET(D3D10_SB_OPCODE_RET, "ret", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_RETC, "retc", 1, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_ROUND_NE, "round_ne", 2, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_ROUND_NI, "round_ni", 2, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_ROUND_PI, "round_pi", 2, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_ROUND_Z, "round_z", 2, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_RSQ, "rsq", 2, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_SAMPLE, "sample", 4, 0x00, D3D10_SB_TEX_OP);
  SET(D3D10_SB_OPCODE_SAMPLE_B, "sample_b", 5, 0x00, D3D10_SB_TEX_OP);
  SET(D3D10_SB_OPCODE_SAMPLE_L, "sample_l", 5, 0x00, D3D10_SB_TEX_OP);
  SET(D3D10_SB_OPCODE_SAMPLE_D, "sample_d", 6, 0x00, D3D10_SB_TEX_OP);
  SET(D3D10_SB_OPCODE_SAMPLE_C, "sample_c", 5, 0x00, D3D10_SB_TEX_OP);
  SET(D3D10_SB_OPCODE_SAMPLE_C_LZ, "sample_c_lz", 5, 0x00, D3D10_SB_TEX_OP);
  SET(D3D10_SB_OPCODE_SQRT, "sqrt", 2, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_SWITCH, "switch", 1, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_SINCOS, "sincos", 3, 0x04, D3D10_SB_FLOAT_OP);
  SET(D3D10_SB_OPCODE_UDIV, "udiv", 4, 0x0c, D3D10_SB_UINT_OP);
  SET(D3D10_SB_OPCODE_ULT, "ult", 3, 0x00, D3D10_SB_UINT_OP);
  SET(D3D10_SB_OPCODE_UGE, "uge", 3, 0x00, D3D10_SB_UINT_OP);
  SET(D3D10_SB_OPCODE_UMAX, "umax", 3, 0x06, D3D10_SB_UINT_OP);
  SET(D3D10_SB_OPCODE_UMIN, "umin", 3, 0x06, D3D10_SB_UINT_OP);
  SET(D3D10_SB_OPCODE_UMUL, "umul", 4, 0x0c, D3D10_SB_UINT_OP);
  SET(D3D10_SB_OPCODE_UMAD, "umad", 4, 0x0e, D3D10_SB_UINT_OP);
  SET(D3D10_SB_OPCODE_USHR, "ushr", 3, 0x02, D3D10_SB_UINT_OP);
  SET(D3D10_SB_OPCODE_UTOF, "utof", 2, 0x00, D3D10_SB_UINT_OP);
  SET(D3D10_SB_OPCODE_XOR, "xor", 3, 0x06, D3D10_SB_BIT_OP);
  SET(D3D10_SB_OPCODE_RESERVED0, "jmp", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D10_SB_OPCODE_DCL_INPUT, "dcl_input", 1, 0x00, D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_OUTPUT, "dcl_output", 1, 0x00, D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_INPUT_SGV, "dcl_input_sgv", 1, 0x00, D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_INPUT_PS_SGV, "dcl_input_ps_sgv", 1, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_GS_INPUT_PRIMITIVE, "dcl_inputprimitive", 0, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_GS_OUTPUT_PRIMITIVE_TOPOLOGY, "dcl_outputtopology", 0,
      0x00, D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_MAX_OUTPUT_VERTEX_COUNT,
      "dcl_max_output_vertex_count", 0, 0x00, D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_INPUT_PS, "dcl_input_ps", 1, 0x00, D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_CONSTANT_BUFFER, "dcl_constantbuffer", 1, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_SAMPLER, "dcl_sampler", 1, 0x00, D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_RESOURCE, "dcl_resource", 1, 0x00, D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_INPUT_SIV, "dcl_input_siv", 1, 0x00, D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_INPUT_PS_SIV, "dcl_input_ps_siv", 1, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_OUTPUT_SIV, "dcl_output_siv", 1, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_OUTPUT_SGV, "dcl_output_sgv", 1, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_TEMPS, "dcl_temps", 0, 0x00, D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_INDEXABLE_TEMP, "dcl_indexableTemp", 0, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_INDEX_RANGE, "dcl_indexrange", 1, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D10_SB_OPCODE_DCL_GLOBAL_FLAGS, "dcl_globalFlags", 0, 0x00,
      D3D10_SB_DCL_OP);

  SET(D3D10_1_SB_OPCODE_SAMPLE_INFO, "sampleinfo", 2, 0x00, D3D10_SB_TEX_OP);
  SET(D3D10_1_SB_OPCODE_SAMPLE_POS, "samplepos", 3, 0x00, D3D10_SB_TEX_OP);
  SET(D3D10_1_SB_OPCODE_GATHER4, "gather4", 4, 0x00, D3D10_SB_TEX_OP);
  SET(D3D10_1_SB_OPCODE_LOD, "lod", 4, 0x00, D3D10_SB_TEX_OP);

  SET(D3D11_SB_OPCODE_EMIT_STREAM, "emit_stream", 1, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D11_SB_OPCODE_CUT_STREAM, "cut_stream", 1, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D11_SB_OPCODE_EMITTHENCUT_STREAM, "emit_then_cut_stream", 1, 0x00,
      D3D10_SB_FLOW_OP);
  SET(D3D11_SB_OPCODE_INTERFACE_CALL, "fcall", 1, 0x00, D3D10_SB_FLOW_OP);

  SET(D3D11_SB_OPCODE_DCL_STREAM, "dcl_stream", 1, 0x00, D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_FUNCTION_BODY, "dcl_function_body", 0, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_FUNCTION_TABLE, "dcl_function_table", 0, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_INTERFACE, "dcl_interface", 0, 0x00, D3D10_SB_DCL_OP);

  SET(D3D11_SB_OPCODE_BUFINFO, "bufinfo", 2, 0x00, D3D10_SB_TEX_OP);
  SET(D3D11_SB_OPCODE_DERIV_RTX_COARSE, "deriv_rtx_coarse", 2, 0x02,
      D3D10_SB_FLOAT_OP);
  SET(D3D11_SB_OPCODE_DERIV_RTX_FINE, "deriv_rtx_fine", 2, 0x02,
      D3D10_SB_FLOAT_OP);
  SET(D3D11_SB_OPCODE_DERIV_RTY_COARSE, "deriv_rty_coarse", 2, 0x02,
      D3D10_SB_FLOAT_OP);
  SET(D3D11_SB_OPCODE_DERIV_RTY_FINE, "deriv_rty_fine", 2, 0x02,
      D3D10_SB_FLOAT_OP);
  SET(D3D11_SB_OPCODE_GATHER4_C, "gather4_c", 5, 0x00, D3D10_SB_TEX_OP);
  SET(D3D11_SB_OPCODE_GATHER4_PO, "gather4_po", 5, 0x00, D3D10_SB_TEX_OP);
  SET(D3D11_SB_OPCODE_GATHER4_PO_C, "gather4_po_c", 6, 0x00, D3D10_SB_TEX_OP);
  SET(D3D11_SB_OPCODE_RCP, "rcp", 2, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D11_SB_OPCODE_F32TOF16, "f32tof16", 2, 0x00, D3D10_SB_FLOAT_OP);
  SET(D3D11_SB_OPCODE_F16TOF32, "f16tof32", 2, 0x00, D3D10_SB_FLOAT_OP);
  SET(D3D11_SB_OPCODE_UADDC, "uaddc", 4, 0x0c, D3D10_SB_UINT_OP);
  SET(D3D11_SB_OPCODE_USUBB, "usubb", 4, 0x0c, D3D10_SB_UINT_OP);
  SET(D3D11_SB_OPCODE_COUNTBITS, "countbits", 2, 0x02, D3D10_SB_BIT_OP);
  SET(D3D11_SB_OPCODE_FIRSTBIT_HI, "firstbit_hi", 2, 0x02, D3D10_SB_BIT_OP);
  SET(D3D11_SB_OPCODE_FIRSTBIT_LO, "firstbit_lo", 2, 0x02, D3D10_SB_BIT_OP);
  SET(D3D11_SB_OPCODE_FIRSTBIT_SHI, "firstbit_shi", 2, 0x02, D3D10_SB_BIT_OP);
  SET(D3D11_SB_OPCODE_UBFE, "ubfe", 4, 0x02, D3D10_SB_BIT_OP);
  SET(D3D11_SB_OPCODE_IBFE, "ibfe", 4, 0x02, D3D10_SB_BIT_OP);
  SET(D3D11_SB_OPCODE_BFI, "bfi", 5, 0x02, D3D10_SB_BIT_OP);
  SET(D3D11_SB_OPCODE_BFREV, "bfrev", 2, 0x02, D3D10_SB_BIT_OP);
  SET(D3D11_SB_OPCODE_SWAPC, "swapc", 5, 0x02, D3D10_SB_FLOAT_OP);

  SET(D3D11_SB_OPCODE_HS_DECLS, "hs_decls", 0, 0x00, D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_HS_CONTROL_POINT_PHASE, "hs_control_point_phase", 0, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_HS_FORK_PHASE, "hs_fork_phase", 0, 0x00, D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_HS_JOIN_PHASE, "hs_join_phase", 0, 0x00, D3D10_SB_DCL_OP);

  SET(D3D11_SB_OPCODE_DCL_INPUT_CONTROL_POINT_COUNT,
      "dcl_input_control_point_count", 0, 0x00, D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_OUTPUT_CONTROL_POINT_COUNT,
      "dcl_output_control_point_count", 0, 0x00, D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_TESS_DOMAIN, "dcl_tessellator_domain", 0, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_TESS_PARTITIONING, "dcl_tessellator_partitioning", 0,
      0x00, D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_TESS_OUTPUT_PRIMITIVE,
      "dcl_tessellator_output_primitive", 0, 0x00, D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_HS_MAX_TESSFACTOR, "dcl_hs_max_tessfactor", 0, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_HS_FORK_PHASE_INSTANCE_COUNT,
      "dcl_hs_fork_phase_instance_count", 0, 0x00, D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_HS_JOIN_PHASE_INSTANCE_COUNT,
      "dcl_hs_join_phase_instance_count", 0, 0x00, D3D10_SB_DCL_OP);

  SET(D3D11_SB_OPCODE_DCL_THREAD_GROUP, "dcl_thread_group", 0, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_UNORDERED_ACCESS_VIEW_TYPED, "dcl_uav_typed", 1, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_UNORDERED_ACCESS_VIEW_RAW, "dcl_uav_raw", 1, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_UNORDERED_ACCESS_VIEW_STRUCTURED,
      "dcl_uav_structured", 1, 0x00, D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_THREAD_GROUP_SHARED_MEMORY_RAW, "dcl_tgsm_raw", 1,
      0x00, D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_THREAD_GROUP_SHARED_MEMORY_STRUCTURED,
      "dcl_tgsm_structured", 1, 0x00, D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_RESOURCE_RAW, "dcl_resource_raw", 1, 0x00,
      D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_DCL_RESOURCE_STRUCTURED, "dcl_resource_structured", 1,
      0x00, D3D10_SB_DCL_OP);
  SET(D3D11_SB_OPCODE_LD_UAV_TYPED, "ld_uav_typed", 3, 0x00, D3D11_SB_MEM_OP);
  SET(D3D11_SB_OPCODE_STORE_UAV_TYPED, "store_uav_typed", 3, 0x00,
      D3D11_SB_MEM_OP);
  SET(D3D11_SB_OPCODE_LD_RAW, "ld_raw", 3, 0x00, D3D11_SB_MEM_OP);
  SET(D3D11_SB_OPCODE_STORE_RAW, "store_raw", 3, 0x00, D3D11_SB_MEM_OP);
  SET(D3D11_SB_OPCODE_LD_STRUCTURED, "ld_structured", 4, 0x00, D3D11_SB_MEM_OP);
  SET(D3D11_SB_OPCODE_STORE_STRUCTURED, "store_structured", 4, 0x00,
      D3D11_SB_MEM_OP);
  SET(D3D11_SB_OPCODE_ATOMIC_AND, "atomic_and", 3, 0x00, D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_ATOMIC_OR, "atomic_or", 3, 0x00, D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_ATOMIC_XOR, "atomic_xor", 3, 0x00, D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_ATOMIC_CMP_STORE, "atomic_cmp_store", 4, 0x00,
      D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_ATOMIC_IADD, "atomic_iadd", 3, 0x00, D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_ATOMIC_IMAX, "atomic_imax", 3, 0x00, D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_ATOMIC_IMIN, "atomic_imin", 3, 0x00, D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_ATOMIC_UMAX, "atomic_umax", 3, 0x00, D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_ATOMIC_UMIN, "atomic_umin", 3, 0x00, D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_IMM_ATOMIC_ALLOC, "imm_atomic_alloc", 2, 0x00,
      D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_IMM_ATOMIC_CONSUME, "imm_atomic_consume", 2, 0x00,
      D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_IMM_ATOMIC_IADD, "imm_atomic_iadd", 4, 0x00,
      D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_IMM_ATOMIC_AND, "imm_atomic_and", 4, 0x00,
      D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_IMM_ATOMIC_OR, "imm_atomic_or", 4, 0x00,
      D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_IMM_ATOMIC_XOR, "imm_atomic_xor", 4, 0x00,
      D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_IMM_ATOMIC_EXCH, "imm_atomic_exch", 4, 0x00,
      D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_IMM_ATOMIC_CMP_EXCH, "imm_atomic_cmp_exch", 5, 0x00,
      D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_IMM_ATOMIC_IMAX, "imm_atomic_imax", 4, 0x00,
      D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_IMM_ATOMIC_IMIN, "imm_atomic_imin", 4, 0x00,
      D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_IMM_ATOMIC_UMAX, "imm_atomic_umax", 4, 0x00,
      D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_IMM_ATOMIC_UMIN, "imm_atomic_umin", 4, 0x00,
      D3D11_SB_ATOMIC_OP);
  SET(D3D11_SB_OPCODE_SYNC, "sync", 0, 0x00, D3D10_SB_FLOW_OP);
  SET(D3D11_SB_OPCODE_EVAL_SNAPPED, "eval_snapped", 3, 0x02, D3D10_SB_FLOAT_OP);
  SET(D3D11_SB_OPCODE_EVAL_SAMPLE_INDEX, "eval_sample_index", 3, 0x02,
      D3D10_SB_FLOAT_OP);
  SET(D3D11_SB_OPCODE_EVAL_CENTROID, "eval_centroid", 2, 0x02,
      D3D10_SB_FLOAT_OP);

  SET(D3D11_SB_OPCODE_DCL_GS_INSTANCE_COUNT, "dcl_gsinstances", 0, 0x00,
      D3D10_SB_DCL_OP);

  SET(D3D11_SB_OPCODE_DADD, "dadd", 3, 0x06, D3D11_SB_DOUBLE_OP);
  SET(D3D11_SB_OPCODE_DMAX, "dmax", 3, 0x06, D3D11_SB_DOUBLE_OP);
  SET(D3D11_SB_OPCODE_DMIN, "dmin", 3, 0x06, D3D11_SB_DOUBLE_OP);
  SET(D3D11_SB_OPCODE_DMUL, "dmul", 3, 0x06, D3D11_SB_DOUBLE_OP);
  SET(D3D11_SB_OPCODE_DEQ, "deq", 3, 0x00, D3D11_SB_DOUBLE_OP);
  SET(D3D11_SB_OPCODE_DGE, "dge", 3, 0x00, D3D11_SB_DOUBLE_OP);
  SET(D3D11_SB_OPCODE_DLT, "dlt", 3, 0x00, D3D11_SB_DOUBLE_OP);
  SET(D3D11_SB_OPCODE_DNE, "dne", 3, 0x00, D3D11_SB_DOUBLE_OP);
  SET(D3D11_SB_OPCODE_DMOV, "dmov", 2, 0x02, D3D11_SB_DOUBLE_OP);
  SET(D3D11_SB_OPCODE_DMOVC, "dmovc", 4, 0x0c, D3D11_SB_DOUBLE_OP);
  SET(D3D11_SB_OPCODE_DTOF, "dtof", 2, 0x02, D3D11_SB_DOUBLE_TO_FLOAT_OP);
  SET(D3D11_SB_OPCODE_FTOD, "ftod", 2, 0x00, D3D11_SB_FLOAT_TO_DOUBLE_OP);

  SET(D3D11_SB_OPCODE_ABORT, "abort", 0, 0x00, D3D11_SB_DEBUG_OP);
  SET(D3D11_SB_OPCODE_DEBUG_BREAK, "debug_break", 0, 0x00, D3D11_SB_DEBUG_OP);

  SET(D3D11_1_SB_OPCODE_DDIV, "ddiv", 3, 0x06, D3D11_SB_DOUBLE_OP);
  SET(D3D11_1_SB_OPCODE_DFMA, "dfma", 4, 0x0e, D3D11_SB_DOUBLE_OP);
  SET(D3D11_1_SB_OPCODE_DRCP, "drcp", 2, 0x02, D3D11_SB_DOUBLE_OP);

  SET(D3D11_1_SB_OPCODE_MSAD, "msad", 4, 0x0e, D3D10_SB_UINT_OP);

  SET(D3D11_1_SB_OPCODE_DTOI, "dtoi", 2, 0x00, D3D11_SB_DOUBLE_OP);
  SET(D3D11_1_SB_OPCODE_DTOU, "dtou", 2, 0x00, D3D11_SB_DOUBLE_OP);
  SET(D3D11_1_SB_OPCODE_ITOD, "itod", 2, 0x00, D3D10_SB_INT_OP);
  SET(D3D11_1_SB_OPCODE_UTOD, "utod", 2, 0x00, D3D10_SB_UINT_OP);

  SET(D3DWDDM1_3_SB_OPCODE_GATHER4_FEEDBACK, "gather4_s", 5, 0x00,
      D3D10_SB_TEX_OP);
  SET(D3DWDDM1_3_SB_OPCODE_GATHER4_C_FEEDBACK, "gather4_c_s", 6, 0x00,
      D3D10_SB_TEX_OP);
  SET(D3DWDDM1_3_SB_OPCODE_GATHER4_PO_FEEDBACK, "gather4_po_s", 6, 0x00,
      D3D10_SB_TEX_OP);
  SET(D3DWDDM1_3_SB_OPCODE_GATHER4_PO_C_FEEDBACK, "gather4_po_c_s", 7, 0x00,
      D3D10_SB_TEX_OP);
  SET(D3DWDDM1_3_SB_OPCODE_LD_FEEDBACK, "ld_s", 4, 0x00, D3D10_SB_TEX_OP);
  SET(D3DWDDM1_3_SB_OPCODE_LD_MS_FEEDBACK, "ldms_s", 5, 0x00, D3D10_SB_TEX_OP);
  SET(D3DWDDM1_3_SB_OPCODE_LD_UAV_TYPED_FEEDBACK, "ld_uav_typed_s", 4, 0x00,
      D3D11_SB_MEM_OP);
  SET(D3DWDDM1_3_SB_OPCODE_LD_RAW_FEEDBACK, "ld_raw_s", 4, 0x00,
      D3D11_SB_MEM_OP);
  SET(D3DWDDM1_3_SB_OPCODE_LD_STRUCTURED_FEEDBACK, "ld_structured_s", 5, 0x00,
      D3D11_SB_MEM_OP);
  SET(D3DWDDM1_3_SB_OPCODE_SAMPLE_L_FEEDBACK, "sample_l_s", 6, 0x00,
      D3D10_SB_TEX_OP);
  SET(D3DWDDM1_3_SB_OPCODE_SAMPLE_C_LZ_FEEDBACK, "sample_c_lz_s", 6, 0x00,
      D3D10_SB_TEX_OP);
  SET(D3DWDDM1_3_SB_OPCODE_SAMPLE_CLAMP_FEEDBACK, "sample_cl_s", 6, 0x00,
      D3D10_SB_TEX_OP);
  SET(D3DWDDM1_3_SB_OPCODE_SAMPLE_B_CLAMP_FEEDBACK, "sample_b_cl_s", 7, 0x00,
      D3D10_SB_TEX_OP);
  SET(D3DWDDM1_3_SB_OPCODE_SAMPLE_D_CLAMP_FEEDBACK, "sample_d_cl_s", 8, 0x00,
      D3D10_SB_TEX_OP);
  SET(D3DWDDM1_3_SB_OPCODE_SAMPLE_C_CLAMP_FEEDBACK, "sample_c_cl_s", 7, 0x00,
      D3D10_SB_TEX_OP);
  SET(D3DWDDM1_3_SB_OPCODE_CHECK_ACCESS_FULLY_MAPPED,
      "check_access_fully_mapped", 2, 0x00, D3D10_SB_TEX_OP);
  // clang-format on
}

struct InstructionModifier {
  uint32_t preciseMask{0};
  uint32_t saturate{0};
};

struct OperandModifier {
  uint32_t modifier{0};
  uint32_t minPrecision{0};
  uint32_t nonUniform{0};
};

enum class OperandComponentsKind {
  None,
  Mask,
  Swizzle,
  One,
};

struct OperandComponents {
  unsigned num;
  OperandComponentsKind kind;
  union {
    uint32_t mask;
    uint32_t swizzle[4];
    uint32_t one;
  };
};

static dxsa::ComponentMask decodeComponentMask(uint32_t rawComponentMask) {
  auto componentMask = static_cast<dxsa::ComponentMask>(0);
  if (rawComponentMask & D3D10_SB_OPERAND_4_COMPONENT_MASK_X)
    componentMask |= dxsa::ComponentMask::x;
  if (rawComponentMask & D3D10_SB_OPERAND_4_COMPONENT_MASK_Y)
    componentMask |= dxsa::ComponentMask::y;
  if (rawComponentMask & D3D10_SB_OPERAND_4_COMPONENT_MASK_Z)
    componentMask |= dxsa::ComponentMask::z;
  if (rawComponentMask & D3D10_SB_OPERAND_4_COMPONENT_MASK_W)
    componentMask |= dxsa::ComponentMask::w;
  return componentMask;
}

class DXBuilder {
public:
  DXBuilder(MLIRContext *context, StringAttr name)
      : context(context),
        module(ModuleOp::create(builder, FileLineColLoc::get(name, 0, 0))),
        builder(module.getRegion()) {}

  using Index = mlir::Value;
  using Operand = mlir::Value;
  using Instruction = mlir::Operation *;
  using Module = mlir::ModuleOp;

  Index buildIndexImm32(int32_t imm, FileLineColLoc loc) {
    Operation *op =
        dxsa::IndexImm::create(builder, loc, builder.getType<dxsa::IndexType>(),
                               builder.getI32IntegerAttr(imm));
    return op->getResults()[0];
  }

  Index buildIndexImm64(int64_t imm, FileLineColLoc loc) {
    Operation *op =
        dxsa::IndexImm::create(builder, loc, builder.getType<dxsa::IndexType>(),
                               builder.getI64IntegerAttr(imm));
    return op->getResults()[0];
  }

  Index buildIndexRelative(Operand operand, FileLineColLoc loc) {
    Operation *op = dxsa::IndexRel::create(
        builder, loc, builder.getType<dxsa::IndexType>(), operand);
    return op->getResults()[0];
  }

  Index buildIndexImm32PlusRelative(int32_t imm, Operand operand,
                                    FileLineColLoc loc) {
    Operation *op = dxsa::IndexRelImm::create(
        builder, loc, builder.getType<dxsa::IndexType>(), operand,
        builder.getStringAttr("add"), builder.getI32IntegerAttr(imm));
    return op->getResults()[0];
  }

  Operand buildOperandImm32(ArrayRef<int32_t> values, FileLineColLoc loc) {
    Operation *op = dxsa::OperandImm::create(
        builder, loc, builder.getType<dxsa::OperandType>(),
        builder.getI32VectorAttr(values));
    return op->getResults()[0];
  }

  Operand buildOperandImm64(ArrayRef<int64_t> values, FileLineColLoc loc) {
    Operation *op = dxsa::OperandImm::create(
        builder, loc, builder.getType<dxsa::OperandType>(),
        builder.getI64VectorAttr(values));
    return op->getResults()[0];
  }

  Operand buildOperand(uint32_t opType, const OperandComponents &components,
                       ArrayRef<Index> indices,
                       const std::optional<OperandModifier> &modifier,
                       FileLineColLoc loc) {
    NamedAttrList attrs;
    attrs.append("type", builder.getI32IntegerAttr(opType));
    if (modifier) {
      const OperandModifier &mod = modifier.value();
      if (uint32_t modifier = mod.modifier) {
        attrs.append("modifier", builder.getI32IntegerAttr(modifier));
      }
      if (uint32_t minPrecision = mod.minPrecision) {
        attrs.append("min_precision", builder.getI32IntegerAttr(minPrecision));
      }
      if (uint32_t nonUniform = mod.nonUniform) {
        attrs.append("non_uniform", builder.getI32IntegerAttr(nonUniform));
      }
    }
    attrs.append("num_components", builder.getI32IntegerAttr(components.num));
    switch (components.kind) {
    case OperandComponentsKind::Mask: {
      attrs.append("mask", builder.getI32IntegerAttr(components.mask));
      break;
    }
    case OperandComponentsKind::Swizzle: {
      SmallVector<int32_t, 4> values;
      for (uint32_t i = 0; i < components.num; ++i) {
        values.push_back(components.swizzle[i]);
      }
      attrs.append("swizzle", builder.getI32VectorAttr(values));
      break;
    }
    case OperandComponentsKind::One: {
      attrs.append("one", builder.getI32IntegerAttr(components.one));
      break;
    }
    case OperandComponentsKind::None:
      break;
    }
    Operation *op = dxsa::Operand::create(
        builder, loc, builder.getType<dxsa::OperandType>(), indices, attrs);
    return op->getResults()[0];
  }

  Instruction buildInstruction(StringRef name, ArrayRef<Operand> operands,
                               const InstructionModifier &modifier,
                               FileLineColLoc loc) {
    return dxsa::Instruction::create(builder, loc, operands,
                                     builder.getStringAttr(name));
  }

  Module buildModule(ArrayRef<Instruction> instructions, FileLineColLoc loc) {
    return module;
  }

  Instruction buildDclGlobalFlags(dxsa::GlobalFlags flags, Location loc) {
    auto flagsAttr = dxsa::GlobalFlagsAttr::get(builder.getContext(), flags);
    return dxsa::DclGlobalFlags::create(builder, loc, flagsAttr);
  }

  Instruction buildDclTemps(uint32_t count, Location loc) {
    return dxsa::DclTemps::create(builder, loc,
                                  builder.getI32IntegerAttr(count));
  }

  Instruction buildDclInputControlPointCount(uint32_t count, Location loc) {
    return dxsa::DclInputControlPointCount::create(
        builder, loc, builder.getI32IntegerAttr(count));
  }

  Instruction buildDclOutputControlPointCount(uint32_t count, Location loc) {
    return dxsa::DclOutputControlPointCount::create(
        builder, loc, builder.getI32IntegerAttr(count));
  }

  Instruction buildDclTessellatorDomain(dxsa::TessellatorDomain domain,
                                        Location loc) {
    auto domainAttr =
        dxsa::TessellatorDomainAttr::get(builder.getContext(), domain);
    return dxsa::DclTessellatorDomain::create(builder, loc, domainAttr);
  }

  Instruction buildDclTessellatorOutputPrimitive(
      dxsa::TessellatorOutputPrimitiveType outputPrimitiveType, Location loc) {
    auto outputPrimitiveTypeAttr =
        dxsa::TessellatorOutputPrimitiveTypeAttr::get(builder.getContext(),
                                                      outputPrimitiveType);
    return dxsa::DclTessellatorOutputPrimitive::create(builder, loc,
                                                       outputPrimitiveTypeAttr);
  }

  Instruction
  buildDclOutputTopology(dxsa::OutputPrimitiveTopology outputTopology,
                         Location loc) {
    auto outputTopologyAttr = dxsa::OutputPrimitiveTopologyAttr::get(
        builder.getContext(), outputTopology);
    return dxsa::DclOutputTopology::create(builder, loc, outputTopologyAttr);
  }

  Instruction buildDclTessellatorPartitioning(
      dxsa::TessellatorPartitioningMode partitioningMode, Location loc) {
    auto partitioningModeAttr = dxsa::TessellatorPartitioningModeAttr::get(
        builder.getContext(), partitioningMode);
    return dxsa::DclTessellatorPartitioning::create(builder, loc,
                                                    partitioningModeAttr);
  }

  Instruction buildDclInputPrimitive(dxsa::InputPrimitive inputPrimitive,
                                     Location loc) {
    auto inputPrimitiveAttr =
        dxsa::InputPrimitiveAttr::get(builder.getContext(), inputPrimitive);
    return dxsa::DclInputPrimitive::create(builder, loc, inputPrimitiveAttr);
  }

  Instruction buildDclGsInstanceCount(uint32_t count, Location loc) {
    return dxsa::DclGsInstanceCount::create(builder, loc,
                                            builder.getI32IntegerAttr(count));
  }

  Instruction buildDclMaxOutputVertexCount(uint32_t count, Location loc) {
    return dxsa::DclMaxOutputVertexCount::create(builder, loc, count);
  }

  Instruction buildDclStream(uint32_t index, Location loc) {
    return dxsa::DclStream::create(builder, loc,
                                   builder.getI32IntegerAttr(index));
  }

  Instruction buildDclInputPs(dxsa::InterpolationMode interpolationMode,
                              Operand operand, Location loc) {
    auto interpolationModeAttr = dxsa::InterpolationModeAttr::get(
        builder.getContext(), interpolationMode);
    return dxsa::DclInputPs::create(builder, loc, interpolationModeAttr,
                                    operand);
  }

  Instruction buildDclInputPsSiv(dxsa::InterpolationMode interpolationMode,
                                 Operand operand,
                                 dxsa::SystemValueName systemValueName,
                                 Location loc) {
    auto interpolationModeAttr = dxsa::InterpolationModeAttr::get(
        builder.getContext(), interpolationMode);
    auto systemValueNameAttr =
        dxsa::SystemValueNameAttr::get(builder.getContext(), systemValueName);
    return dxsa::DclInputPsSiv::create(builder, loc, interpolationModeAttr,
                                       operand, systemValueNameAttr);
  }

  Instruction buildDclInputPsSgv(Operand operand,
                                 dxsa::SystemValueName systemValueName,
                                 Location loc) {
    auto systemValueNameAttr =
        dxsa::SystemValueNameAttr::get(builder.getContext(), systemValueName);
    return dxsa::DclInputPsSgv::create(builder, loc, operand,
                                       systemValueNameAttr);
  }

  dxsa::InlineOperandAttr buildInlineOperandAttr(
      dxsa::InlineOperandType operandType, uint32_t components,
      std::optional<dxsa::ComponentMask> mask, ArrayRef<int64_t> indexArray) {
    auto *ctx = builder.getContext();
    auto maskAttr = mask ? dxsa::ComponentMaskAttr::get(ctx, *mask)
                         : dxsa::ComponentMaskAttr();
    auto indexAttr = indexArray.empty()
                         ? DenseI64ArrayAttr()
                         : DenseI64ArrayAttr::get(ctx, indexArray);
    return dxsa::InlineOperandAttr::get(ctx, operandType, components, maskAttr,
                                        indexAttr);
  }

  Instruction buildDclInput(dxsa::InlineOperandAttr operand, Location loc) {
    return dxsa::DclInput::create(builder, loc, operand);
  }

  Instruction buildDclOutput(dxsa::InlineOperandAttr operand, Location loc) {
    return dxsa::DclOutput::create(builder, loc, operand);
  }

  Instruction buildDclIndexRange(dxsa::InlineOperandAttr operand,
                                 uint32_t count, Location loc) {
    return dxsa::DclIndexRange::create(builder, loc, operand,
                                       builder.getI32IntegerAttr(count));
  }

  Instruction buildDclInputSgv(dxsa::InlineOperandAttr operand,
                               dxsa::SystemValueName name, Location loc) {
    auto nameAttr = dxsa::SystemValueNameAttr::get(builder.getContext(), name);
    return dxsa::DclInputSgv::create(builder, loc, operand, nameAttr);
  }

  Instruction buildDclInputSiv(dxsa::InlineOperandAttr operand,
                               dxsa::SystemValueName name, Location loc) {
    auto nameAttr = dxsa::SystemValueNameAttr::get(builder.getContext(), name);
    return dxsa::DclInputSiv::create(builder, loc, operand, nameAttr);
  }

  Instruction buildDclOutputSgv(dxsa::InlineOperandAttr operand,
                                dxsa::SystemValueName name, Location loc) {
    auto nameAttr = dxsa::SystemValueNameAttr::get(builder.getContext(), name);
    return dxsa::DclOutputSgv::create(builder, loc, operand, nameAttr);
  }

  Instruction buildDclOutputSiv(dxsa::InlineOperandAttr operand,
                                dxsa::SystemValueName name, Location loc) {
    auto nameAttr = dxsa::SystemValueNameAttr::get(builder.getContext(), name);
    return dxsa::DclOutputSiv::create(builder, loc, operand, nameAttr);
  }

  Instruction buildDclHsMaxTessFactor(float maxTessFactor, Location loc) {
    return dxsa::DclHsMaxTessFactor::create(
        builder, loc, builder.getF32FloatAttr(maxTessFactor));
  }

  Instruction buildDclHsJoinPhaseInstanceCount(uint32_t count, Location loc) {
    return dxsa::DclHsJoinPhaseInstanceCount::create(
        builder, loc, builder.getUI32IntegerAttr(count));
  }

  Instruction buildDclHsForkPhaseInstanceCount(uint32_t count, Location loc) {
    return dxsa::DclHsForkPhaseInstanceCount::create(
        builder, loc, builder.getUI32IntegerAttr(count));
  }

  Instruction buildDclTgsmRaw(dxsa::InlineOperandAttr operand,
                              uint32_t byteCount, Location loc) {
    return dxsa::DclTgsmRaw::create(builder, loc, operand,
                                    builder.getI32IntegerAttr(byteCount));
  }

  Instruction buildDclTgsmStructured(dxsa::InlineOperandAttr operand,
                                     uint32_t structByteStride,
                                     uint32_t structCount, Location loc) {
    return dxsa::DclTgsmStructured::create(
        builder, loc, operand, builder.getI32IntegerAttr(structByteStride),
        builder.getI32IntegerAttr(structCount));
  }

  Instruction buildDclConstantBuffer(
      uint32_t id, uint32_t size, std::optional<uint32_t> lbound,
      std::optional<uint32_t> ubound, std::optional<uint32_t> space,
      dxsa::ConstantBufferAccessPattern accessPattern, Location loc) {
    auto optionalToAttr = [&](std::optional<uint32_t> v) -> IntegerAttr {
      return v ? builder.getI32IntegerAttr(*v) : IntegerAttr();
    };
    return dxsa::DclConstantBuffer::create(
        builder, loc, id, size, optionalToAttr(lbound), optionalToAttr(ubound),
        optionalToAttr(space), accessPattern);
  }

  Instruction buildDclSampler(uint32_t id, std::optional<uint32_t> lbound,
                              std::optional<uint32_t> ubound,
                              std::optional<uint32_t> space,
                              dxsa::SamplerMode mode, Location loc) {
    auto optionalToAttr = [&](std::optional<uint32_t> v) -> IntegerAttr {
      return v ? builder.getI32IntegerAttr(*v) : IntegerAttr();
    };
    return dxsa::DclSampler::create(
        builder, loc, id, mode, optionalToAttr(lbound), optionalToAttr(ubound),
        optionalToAttr(space));
  }

  Instruction buildInterfaceCall(Operand operand, uint32_t callSite,
                                 Location loc) {
    return dxsa::InterfaceCall::create(builder, loc, operand, callSite);
  }

private:
  MLIRContext *context;
  ModuleOp module;
  OpBuilder builder;
};

class Parser {
public:
  Parser(DXBuilder &builder, StringAttr name, StringRef buffer)
      : builder(builder), name(name), buffer(buffer) {
    initInstructionInfo(instrInfo);
  }

  using Token = FailureOr<uint32_t>;
  using Index = DXBuilder::Index;
  using Operand = DXBuilder::Operand;
  using Instruction = DXBuilder::Instruction;
  using Module = DXBuilder::Module;

  /// Parse the current token and move the cursor to the next one.
  Token parseToken() {
    constexpr size_t tokenSize = sizeof(uint32_t);
    if ((currentTokenOffset + tokenSize) > buffer.size()) {
      emitError(getLocation(), "unexpected end of file");
      return failure();
    }

    uint32_t value = support::endian::read<uint32_t>(
        buffer.begin() + currentTokenOffset, endianness::little);
    currentTokenOffset += tokenSize;

    return value;
  }

  /// Returns location where the last parsed token begins (at offset
  /// -4 from the currentTokenOffset).
  FileLineColLoc getLocation(int offset = -4) const {
    return FileLineColLoc::get(name, 0, currentTokenOffset + offset);
  }

  bool isImmOperand(uint32_t token) {
    switch (DECODE_D3D10_SB_OPERAND_TYPE(token)) {
    case D3D10_SB_OPERAND_TYPE_IMMEDIATE32:
    case D3D10_SB_OPERAND_TYPE_IMMEDIATE64:
      return true;
    default:
      return false;
    }
  }

  FailureOr<OperandComponents> parseOperandComponents(uint32_t token) {
    OperandComponents components{};
    switch (DECODE_D3D10_SB_OPERAND_NUM_COMPONENTS(token)) {
    case D3D10_SB_OPERAND_0_COMPONENT: {
      components.num = 0;
      break;
    }
    case D3D10_SB_OPERAND_1_COMPONENT: {
      components.num = 1;
      break;
    }
    case D3D10_SB_OPERAND_4_COMPONENT: {
      components.num = 4;
      break;
    }
    default:
      emitError(getLocation(), "unexpected number of components");
      return failure();
    }

    if (components.num != 4 || isImmOperand(token))
      return components;

    switch (DECODE_D3D10_SB_OPERAND_4_COMPONENT_SELECTION_MODE(token)) {
    case D3D10_SB_OPERAND_4_COMPONENT_MASK_MODE: {
      components.kind = OperandComponentsKind::Mask;
      components.mask = DECODE_D3D10_SB_OPERAND_4_COMPONENT_MASK(token);
      break;
    }
    case D3D10_SB_OPERAND_4_COMPONENT_SWIZZLE_MODE: {
      components.kind = OperandComponentsKind::Swizzle;
      components.swizzle[0] =
          DECODE_D3D10_SB_OPERAND_4_COMPONENT_SWIZZLE_SOURCE(token, 0);
      components.swizzle[1] =
          DECODE_D3D10_SB_OPERAND_4_COMPONENT_SWIZZLE_SOURCE(token, 1);
      components.swizzle[2] =
          DECODE_D3D10_SB_OPERAND_4_COMPONENT_SWIZZLE_SOURCE(token, 2);
      components.swizzle[3] =
          DECODE_D3D10_SB_OPERAND_4_COMPONENT_SWIZZLE_SOURCE(token, 3);
      break;
    }
    case D3D10_SB_OPERAND_4_COMPONENT_SELECT_1_MODE: {
      components.kind = OperandComponentsKind::One;
      components.one = DECODE_D3D10_SB_OPERAND_4_COMPONENT_SELECT_1(token);
      break;
    }
    default:
      emitError(getLocation(), "unexpected component selection");
      return failure();
    }

    return components;
  }

  using OperandIndexTypes = SmallVector<uint32_t, 3>;

  FailureOr<OperandIndexTypes> parseOperandIndexTypes(uint32_t token) {
    SmallVector<uint32_t, 3> indexTypes;
    if (isImmOperand(token))
      return indexTypes; // none

    uint32_t indexDimension = DECODE_D3D10_SB_OPERAND_INDEX_DIMENSION(token);
    if (indexDimension > 3) {
      emitError(getLocation(),
                "invalid operand index dimension (must be <= 3)");
      return failure();
    }

    if (indexDimension == D3D10_SB_OPERAND_INDEX_0D)
      return indexTypes; // none

    indexTypes.resize(indexDimension);
    for (unsigned i = 0; i < indexTypes.size(); ++i) {
      indexTypes[i] = DECODE_D3D10_SB_OPERAND_INDEX_REPRESENTATION(i, token);
    }

    return indexTypes;
  }

  FailureOr<Index> parseIndex(uint32_t indexType) {
    switch (indexType) {
    case D3D10_SB_OPERAND_INDEX_IMMEDIATE32: {
      Token value = parseToken();
      if (failed(value)) {
        emitError(getLocation(), "expected an operand index imm32");
        return failure();
      }
      return builder.buildIndexImm32(*value, getLocation());
    }
    case D3D10_SB_OPERAND_INDEX_IMMEDIATE64: {
      Token value0 = parseToken();
      if (failed(value0)) {
        emitError(getLocation(), "expected an operand index imm64");
        return failure();
      }

      FileLineColLoc loc = getLocation();

      Token value1 = parseToken();
      if (failed(value1)) {
        emitError(getLocation(),
                  "expected an operand index imm64 (second token)");
        return failure();
      }

      // TODO: check the order of tokens (MSB or LSB?)
      return builder.buildIndexImm64((((uint64_t)*value0) << 32) | *value1,
                                     loc);
    }
    case D3D10_SB_OPERAND_INDEX_RELATIVE: {
      FailureOr<Operand> operand = parseOperand();
      if (failed(operand)) {
        emitError(getLocation(), "expected an index operand");
        return failure();
      }
      return builder.buildIndexRelative(*operand, getLocation());
    }
    case D3D10_SB_OPERAND_INDEX_IMMEDIATE32_PLUS_RELATIVE: {
      Token imm = parseToken();
      if (failed(imm)) {
        emitError(getLocation(), "expected an operand index relative (imm)");
        return failure();
      }

      FileLineColLoc loc = getLocation();

      FailureOr<Operand> operand = parseOperand();
      if (failed(operand)) {
        emitError(getLocation(),
                  "expected an operand index relative (operand)");
        return failure();
      }

      return builder.buildIndexImm32PlusRelative(*imm, *operand, loc);
    }
    default:
      emitError(getLocation(), "invalid operand index type");
      return failure();
    }
  }

  FailureOr<std::optional<OperandModifier>>
  parseOperandExtendedModifier(uint32_t extToken) {
    std::optional<OperandModifier> none;
    if (D3D10_SB_EXTENDED_OPERAND_MODIFIER !=
        DECODE_D3D10_SB_EXTENDED_OPERAND_TYPE(extToken))
      return none;

    OperandModifier modifier;
    modifier.modifier = DECODE_D3D10_SB_OPERAND_MODIFIER(extToken);
    modifier.minPrecision = DECODE_D3D11_SB_OPERAND_MIN_PRECISION(extToken);
    modifier.nonUniform = DECODE_D3D12_SB_OPERAND_NON_UNIFORM(extToken);
    return std::make_optional(modifier);
  }

  FailureOr<Operand> parseOperand() {
    Token token = parseToken();
    if (failed(token))
      return failure();

    FileLineColLoc loc = getLocation();

    uint32_t opType = DECODE_D3D10_SB_OPERAND_TYPE(*token);
    bool isExtended = DECODE_IS_D3D10_SB_OPERAND_EXTENDED(*token);

    FailureOr<OperandComponents> components = parseOperandComponents(*token);
    if (failed(components))
      return failure();

    FailureOr<OperandIndexTypes> indexTypes = parseOperandIndexTypes(*token);
    if (failed(indexTypes))
      return failure();

    std::optional<OperandModifier> modifier;
    if (isExtended) {
      Token extToken = parseToken();
      if (failed(extToken)) {
        emitError(getLocation(), "unexpected an extended operand token");
        return failure();
      }
      auto failureOrModifier = parseOperandExtendedModifier(*extToken);
      if (failed(failureOrModifier))
        return failure();
      modifier = *failureOrModifier;
    }

    if (isImmOperand(*token)) {
      switch (opType) {
      case D3D10_SB_OPERAND_TYPE_IMMEDIATE32: {
        SmallVector<int32_t, 4> values;
        for (unsigned i = 0; i < components->num; ++i) {
          Token value = parseToken();
          if (failed(value)) {
            emitError(getLocation(), "expected an immediate operand (imm32)");
            return failure();
          }
          values.push_back(*value);
        }
        return builder.buildOperandImm32(values, loc);
      }
      case D3D10_SB_OPERAND_TYPE_IMMEDIATE64: {
        if (components->num != 4) {
          emitError(getLocation(), "imm64 operand must have 4 components");
          return failure();
        }
        SmallVector<int64_t, 2> values;
        for (unsigned i = 0; i < 2; ++i) {
          Token high = parseToken();
          if (failed(high)) {
            emitError(getLocation(),
                      "expected an immediate operand (imm64 high)");
            return failure();
          }

          Token low = parseToken();
          if (failed(low)) {
            emitError(getLocation(),
                      "expected an immediate operand (imm64 low)");
            return failure();
          }

          values.push_back((((int64_t)*high) << 32) | *low);
        }
        return builder.buildOperandImm64(values, loc);
      }
      }
      emitError(getLocation(), "unhandled immediate type");
      return failure();
    }

    // Operand indices
    SmallVector<Index, 3> indices;
    for (uint32_t indexType : *indexTypes) {
      FailureOr<Index> index = parseIndex(indexType);
      if (failed(index))
        return failure();
      indices.push_back(*index);
    }
    return builder.buildOperand(opType, *components, indices, modifier,
                                getLocation());
  }

  FailureOr<Instruction> parseDclGlobalFlags(uint32_t opcodeToken,
                                             Location loc) {
    auto raw = DECODE_D3D10_SB_GLOBAL_FLAGS(opcodeToken);
    if (raw == 0) {
      emitError(loc, "expected at least one global flag to be set");
      return failure();
    }
    auto flags = static_cast<dxsa::GlobalFlags>(0);
    if (raw & D3D10_SB_GLOBAL_FLAG_REFACTORING_ALLOWED)
      flags |= dxsa::GlobalFlags::refactoringAllowed;
    if (raw & D3D11_SB_GLOBAL_FLAG_ENABLE_DOUBLE_PRECISION_FLOAT_OPS)
      flags |= dxsa::GlobalFlags::enableDoublePrecisionFloatOps;
    if (raw & D3D11_SB_GLOBAL_FLAG_FORCE_EARLY_DEPTH_STENCIL)
      flags |= dxsa::GlobalFlags::forceEarlyDepthStencil;
    if (raw & D3D11_SB_GLOBAL_FLAG_ENABLE_RAW_AND_STRUCTURED_BUFFERS)
      flags |= dxsa::GlobalFlags::enableRawAndStructuredBuffers;
    if (raw & D3D11_1_SB_GLOBAL_FLAG_SKIP_OPTIMIZATION)
      flags |= dxsa::GlobalFlags::skipOptimization;
    if (raw & D3D11_1_SB_GLOBAL_FLAG_ENABLE_MINIMUM_PRECISION)
      flags |= dxsa::GlobalFlags::enableMinimumPrecision;
    if (raw & D3D11_1_SB_GLOBAL_FLAG_ENABLE_DOUBLE_EXTENSIONS)
      flags |= dxsa::GlobalFlags::enableDoubleExtensions;
    if (raw & D3D11_1_SB_GLOBAL_FLAG_ENABLE_SHADER_EXTENSIONS)
      flags |= dxsa::GlobalFlags::enableShaderExtensions;
    if (raw & D3D12_SB_GLOBAL_FLAG_ALL_RESOURCES_BOUND)
      flags |= dxsa::GlobalFlags::allResourcesBound;
    return builder.buildDclGlobalFlags(flags, loc);
  }

  FailureOr<Instruction> parseDclTemps(Location 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);
  }

  FailureOr<Instruction> parseDclInputControlPointCount(uint32_t opcodeToken,
                                                        Location loc) {
    auto count = DECODE_D3D11_SB_INPUT_CONTROL_POINT_COUNT(opcodeToken);
    if (count == 0) {
      emitError(loc, "input control point count cannot be zero");
      return failure();
    }
    if (count > 32) {
      emitError(loc, "input control point count must be <= 32, got ") << count;
      return failure();
    }
    return builder.buildDclInputControlPointCount(count, loc);
  }

  FailureOr<Instruction> parseDclOutputControlPointCount(uint32_t opcodeToken,
                                                         Location loc) {
    auto count = DECODE_D3D11_SB_OUTPUT_CONTROL_POINT_COUNT(opcodeToken);
    if (count > 32) {
      emitError(loc, "output control point count must be <= 32, got ") << count;
      return failure();
    }
    return builder.buildDclOutputControlPointCount(count, loc);
  }

  FailureOr<Instruction> parseDclTessellatorDomain(uint32_t opcodeToken,
                                                   Location loc) {
    auto rawDomain = DECODE_D3D11_SB_TESS_DOMAIN(opcodeToken);
    auto domain = dxsa::symbolizeTessellatorDomain(rawDomain);
    if (!domain)
      return emitError(loc, "unknown tessellator domain: ") << rawDomain;
    return builder.buildDclTessellatorDomain(*domain, loc);
  }

  FailureOr<Instruction>
  parseDclTessellatorOutputPrimitive(uint32_t opcodeToken, Location loc) {
    auto rawOutputPrimitiveType =
        DECODE_D3D11_SB_TESS_OUTPUT_PRIMITIVE(opcodeToken);
    auto outputPrimitiveType =
        dxsa::symbolizeTessellatorOutputPrimitiveType(rawOutputPrimitiveType);
    if (!outputPrimitiveType)
      return emitError(loc, "unknown tessellator output primitive type: ")
             << rawOutputPrimitiveType;
    return builder.buildDclTessellatorOutputPrimitive(*outputPrimitiveType,
                                                      loc);
  }

  FailureOr<Instruction> parseDclOutputTopology(uint32_t opcodeToken,
                                                Location loc) {
    auto rawOutputTopology =
        DECODE_D3D10_SB_GS_OUTPUT_PRIMITIVE_TOPOLOGY(opcodeToken);
    auto outputTopology =
        dxsa::symbolizeOutputPrimitiveTopology(rawOutputTopology);
    if (!outputTopology)
      return emitError(loc, "unknown output primitive topology: ")
             << rawOutputTopology;
    return builder.buildDclOutputTopology(*outputTopology, loc);
  }

  FailureOr<Instruction> parseDclTessellatorPartitioning(uint32_t opcodeToken,
                                                         Location loc) {
    auto rawPartitioningMode = DECODE_D3D11_SB_TESS_PARTITIONING(opcodeToken);
    auto partitioningMode =
        dxsa::symbolizeTessellatorPartitioningMode(rawPartitioningMode);
    if (!partitioningMode)
      return emitError(loc, "unknown tessellator partitioning mode: ")
             << rawPartitioningMode;
    return builder.buildDclTessellatorPartitioning(*partitioningMode, loc);
  }

  FailureOr<Instruction> parseDclInputPrimitive(uint32_t opcodeToken,
                                                Location loc) {
    auto rawInputPrimitive = DECODE_D3D10_SB_GS_INPUT_PRIMITIVE(opcodeToken);
    auto inputPrimitive = dxsa::symbolizeInputPrimitive(rawInputPrimitive);
    if (!inputPrimitive)
      return emitError(loc, "unknown input primitive: ") << rawInputPrimitive;
    return builder.buildDclInputPrimitive(*inputPrimitive, loc);
  }

  FailureOr<Instruction> parseDclGsInstanceCount(Location loc) {
    auto countToken = parseToken();
    FAILURE_IF_FAILED(countToken);
    auto count = *countToken;
    if (count == 0)
      return emitError(loc, "instance count cannot be zero");
    if (count > 32)
      return emitError(loc, "instance count must be <= 32, got ") << count;
    return builder.buildDclGsInstanceCount(count, loc);
  }

  FailureOr<Instruction> parseDclMaxOutputVertexCount(Location loc) {
    auto countToken = parseToken();
    FAILURE_IF_FAILED(countToken);
    auto count = *countToken;
    if (count == 0)
      return emitError(getLocation(), "max output vertex count cannot be zero");
    if (count > 1024)
      return emitError(getLocation(),
                       "max output vertex count must be <= 1024, got ")
             << count;
    return builder.buildDclMaxOutputVertexCount(count, loc);
  }

  FailureOr<Instruction> parseDclStream(Location loc) {
    auto operand = parseInlineOperand();
    FAILURE_IF_FAILED(operand);
    if (operand->getType() != dxsa::InlineOperandType::stream)
      return emitError(loc, "unexpected operand type: ")
             << dxsa::stringifyInlineOperandType(operand->getType());
    if (operand->getComponents() != 0)
      return emitError(loc, "unexpected number of components: ")
             << operand->getComponents();
    auto indexArray = operand->getIndex();
    if (!indexArray || indexArray.size() != 1)
      return emitError(loc, "unsupported index dimension: ")
             << (indexArray ? indexArray.size() : 0);
    return builder.buildDclStream(static_cast<uint32_t>(indexArray[0]), loc);
  }

  FailureOr<dxsa::InterpolationMode>
  parseInterpolationMode(uint32_t opcodeToken, Location loc) {
    auto rawInterpolationMode =
        DECODE_D3D10_SB_INPUT_INTERPOLATION_MODE(opcodeToken);
    auto interpolationMode =
        dxsa::symbolizeInterpolationMode(rawInterpolationMode);
    if (!interpolationMode)
      return emitError(loc, "unknown interpolation mode: ")
             << rawInterpolationMode;
    return *interpolationMode;
  }

  FailureOr<dxsa::SystemValueName> parseSystemValueName(Location loc) {
    Token nameToken = parseToken();
    if (failed(nameToken))
      return failure();
    auto rawSystemValueName = DECODE_D3D10_SB_NAME(*nameToken);
    auto systemValueName = dxsa::symbolizeSystemValueName(rawSystemValueName);
    if (!systemValueName)
      return emitError(loc, "unknown system value name: ")
             << rawSystemValueName;
    return *systemValueName;
  }

  FailureOr<Instruction> parseDclInputPs(uint32_t opcodeToken, Location loc) {
    FailureOr<dxsa::InterpolationMode> interpolationMode =
        parseInterpolationMode(opcodeToken, loc);
    if (failed(interpolationMode))
      return failure();
    FailureOr<Operand> operand = parseOperand();
    if (failed(operand))
      return failure();
    return builder.buildDclInputPs(*interpolationMode, *operand, loc);
  }

  FailureOr<Instruction> parseDclInputPsSiv(uint32_t opcodeToken,
                                            Location loc) {
    FailureOr<dxsa::InterpolationMode> interpolationMode =
        parseInterpolationMode(opcodeToken, loc);
    if (failed(interpolationMode))
      return failure();
    FailureOr<Operand> operand = parseOperand();
    if (failed(operand))
      return failure();
    FailureOr<dxsa::SystemValueName> systemValueName =
        parseSystemValueName(getLocation());
    if (failed(systemValueName))
      return failure();
    return builder.buildDclInputPsSiv(*interpolationMode, *operand,
                                      *systemValueName, loc);
  }

  FailureOr<Instruction> parseDclInputPsSgv(Location loc) {
    FailureOr<Operand> operand = parseOperand();
    if (failed(operand))
      return failure();
    FailureOr<dxsa::SystemValueName> systemValueName =
        parseSystemValueName(getLocation());
    if (failed(systemValueName))
      return failure();
    return builder.buildDclInputPsSgv(*operand, *systemValueName, loc);
  }

  FailureOr<dxsa::InlineOperandAttr> parseInlineOperand() {
    auto token = parseToken();
    FAILURE_IF_FAILED(token);

    auto rawOperandType = DECODE_D3D10_SB_OPERAND_TYPE(*token);
    auto isExtended = DECODE_IS_D3D10_SB_OPERAND_EXTENDED(*token);
    auto loc = getLocation();

    if (isImmOperand(*token))
      return emitError(loc, "immediate operand is not supported yet");

    auto type = dxsa::symbolizeInlineOperandType(rawOperandType);
    if (!type)
      return emitError(loc, "unknown operand type: ") << rawOperandType;

    auto components = parseOperandComponents(*token);
    FAILURE_IF_FAILED(components);

    auto indexTypes = parseOperandIndexTypes(*token);
    FAILURE_IF_FAILED(indexTypes);

    if (isExtended)
      return emitError(loc, "extended operand tokens are not yet supported in "
                            "inline operand attribute");

    if (components->kind == OperandComponentsKind::Swizzle ||
        components->kind == OperandComponentsKind::One)
      return emitError(loc, "swizzled / single-component operand selection is "
                            "not supported in inline operand attribute");

    std::optional<dxsa::ComponentMask> mask;
    if (components->kind == OperandComponentsKind::Mask)
      mask = decodeComponentMask(components->mask);

    SmallVector<int64_t, 3> indices;
    for (auto indexType : *indexTypes) {
      if (indexType != D3D10_SB_OPERAND_INDEX_IMMEDIATE32)
        return emitError(getLocation(), "unsupported index representation: ")
               << indexType;
      auto value = parseToken();
      FAILURE_IF_FAILED(value);
      indices.push_back(static_cast<int32_t>(*value));
    }

    return builder.buildInlineOperandAttr(*type, components->num, mask,
                                          indices);
  }

  FailureOr<Instruction> parseDclInput(Location loc) {
    auto operand = parseInlineOperand();
    FAILURE_IF_FAILED(operand);
    return builder.buildDclInput(*operand, loc);
  }

  FailureOr<Instruction> parseDclOutput(Location loc) {
    auto operand = parseInlineOperand();
    FAILURE_IF_FAILED(operand);
    return builder.buildDclOutput(*operand, loc);
  }

  FailureOr<Instruction> parseDclIndexRange(Location loc) {
    auto operand = parseInlineOperand();
    FAILURE_IF_FAILED(operand);
    auto count = parseToken();
    FAILURE_IF_FAILED(count);
    return builder.buildDclIndexRange(*operand, *count, loc);
  }

  FailureOr<Instruction> parseDclInputSiv(Location loc) {
    auto operand = parseInlineOperand();
    FAILURE_IF_FAILED(operand);
    auto name = parseSystemValueName(getLocation());
    FAILURE_IF_FAILED(name);
    return builder.buildDclInputSiv(*operand, *name, loc);
  }

  FailureOr<Instruction> parseDclOutputSgv(Location loc) {
    auto operand = parseInlineOperand();
    FAILURE_IF_FAILED(operand);
    auto name = parseSystemValueName(getLocation());
    FAILURE_IF_FAILED(name);
    return builder.buildDclOutputSgv(*operand, *name, loc);
  }

  FailureOr<Instruction> parseDclOutputSiv(Location loc) {
    auto operand = parseInlineOperand();
    FAILURE_IF_FAILED(operand);
    auto name = parseSystemValueName(getLocation());
    FAILURE_IF_FAILED(name);
    return builder.buildDclOutputSiv(*operand, *name, loc);
  }

  FailureOr<Instruction> parseDclInputSgv(Location loc) {
    auto operand = parseInlineOperand();
    FAILURE_IF_FAILED(operand);
    auto name = parseSystemValueName(getLocation());
    FAILURE_IF_FAILED(name);
    return builder.buildDclInputSgv(*operand, *name, loc);
  }

  FailureOr<Instruction> parseDclHsMaxTessFactor(Location loc) {
    auto token = parseToken();
    FAILURE_IF_FAILED(token);
    auto maxTessFactor = llvm::bit_cast<float>(*token);
    return builder.buildDclHsMaxTessFactor(maxTessFactor, loc);
  }

  FailureOr<Instruction> parseDclHsJoinPhaseInstanceCount(Location loc) {
    auto count = parseToken();
    FAILURE_IF_FAILED(count);
    return builder.buildDclHsJoinPhaseInstanceCount(*count, loc);
  }

  FailureOr<Instruction> parseDclHsForkPhaseInstanceCount(Location loc) {
    auto count = parseToken();
    FAILURE_IF_FAILED(count);
    return builder.buildDclHsForkPhaseInstanceCount(*count, loc);
  }

  FailureOr<Instruction> parseDclTgsmRaw(Location loc) {
    auto operand = parseInlineOperand();
    FAILURE_IF_FAILED(operand);
    auto byteCount = parseToken();
    FAILURE_IF_FAILED(byteCount);
    return builder.buildDclTgsmRaw(*operand, *byteCount, loc);
  }

  FailureOr<Instruction> parseDclTgsmStructured(Location loc) {
    auto operand = parseInlineOperand();
    FAILURE_IF_FAILED(operand);
    auto structByteStride = parseToken();
    FAILURE_IF_FAILED(structByteStride);
    auto structCount = parseToken();
    FAILURE_IF_FAILED(structCount);
    return builder.buildDclTgsmStructured(*operand, *structByteStride,
                                          *structCount, loc);
  }

  FailureOr<Instruction> parseDclConstantBuffer(uint32_t opcodeToken,
                                                Location loc) {
    auto rawAccessPattern =
        DECODE_D3D10_SB_CONSTANT_BUFFER_ACCESS_PATTERN(opcodeToken);
    auto accessPattern =
        dxsa::symbolizeConstantBufferAccessPattern(rawAccessPattern);
    if (!accessPattern)
      return emitError(loc, "unknown constant buffer access pattern: ")
             << rawAccessPattern;

    auto operandToken = parseToken();
    FAILURE_IF_FAILED(operandToken);

    auto operandType = DECODE_D3D10_SB_OPERAND_TYPE(*operandToken);
    if (operandType != D3D10_SB_OPERAND_TYPE_CONSTANT_BUFFER)
      return emitError(loc, "unexpected operand type: ") << operandType;

    if (DECODE_IS_D3D10_SB_OPERAND_EXTENDED(*operandToken))
      return emitError(loc, "extended operand tokens are not supported");

    auto indexDim = DECODE_D3D10_SB_OPERAND_INDEX_DIMENSION(*operandToken);
    if (indexDim != D3D10_SB_OPERAND_INDEX_2D &&
        indexDim != D3D10_SB_OPERAND_INDEX_3D)
      return emitError(loc, "unsupported index dimension: ") << indexDim;

    SmallVector<uint32_t, 3> indices;
    indices.reserve(indexDim);
    for (uint32_t i = 0; i < indexDim; ++i) {
      auto indexRepesentation =
          DECODE_D3D10_SB_OPERAND_INDEX_REPRESENTATION(i, *operandToken);
      if (indexRepesentation != D3D10_SB_OPERAND_INDEX_IMMEDIATE32)
        return emitError(loc, "unsupported index representation: ")
               << indexRepesentation;
      auto value = parseToken();
      FAILURE_IF_FAILED(value);
      indices.push_back(*value);
    }

    switch (indexDim) {
    case D3D10_SB_OPERAND_INDEX_2D:
      return builder.buildDclConstantBuffer(
          /*id=*/indices[0], /*size=*/indices[1], /*lbound=*/std::nullopt,
          /*ubound=*/std::nullopt, /*space=*/std::nullopt, *accessPattern, loc);
    case D3D10_SB_OPERAND_INDEX_3D: {
      auto sizeToken = parseToken();
      FAILURE_IF_FAILED(sizeToken);
      auto spaceToken = parseToken();
      FAILURE_IF_FAILED(spaceToken);
      return builder.buildDclConstantBuffer(
          /*id=*/indices[0], /*size=*/*sizeToken, /*lbound=*/indices[1],
          /*ubound=*/indices[2], /*space=*/*spaceToken, *accessPattern, loc);
    }
    default:
      llvm_unreachable("indexDim was validated above");
    }
  }

  FailureOr<Instruction> parseDclSampler(uint32_t opcodeToken, Location loc) {
    auto rawMode = DECODE_D3D10_SB_SAMPLER_MODE(opcodeToken);
    auto mode = dxsa::symbolizeSamplerMode(rawMode);
    if (!mode)
      return emitError(loc, "unknown sampler mode: ") << rawMode;

    auto operand = parseInlineOperand();
    FAILURE_IF_FAILED(operand);
    if (operand->getType() != dxsa::InlineOperandType::sampler)
      return emitError(loc, "operand must be a sampler register, got ")
             << dxsa::stringifyInlineOperandType(operand->getType());
    auto indexArray = operand->getIndex();
    auto indexDim = indexArray ? indexArray.size() : 0;
    if (indexDim != 1 && indexDim != 3)
      return emitError(loc, "operand must have a 1D or 3D index, got ")
             << indexDim;
    auto id = indexArray[0];
    std::optional<uint32_t> lbound, ubound, space;
    if (indexDim == 3) {
      lbound = indexArray[1];
      ubound = indexArray[2];
      auto spaceToken = parseToken();
      FAILURE_IF_FAILED(spaceToken);
      space = *spaceToken;
    }

    return builder.buildDclSampler(id, lbound, ubound, space, *mode, loc);
  }

  FailureOr<Instruction> parseInterfaceCall(uint32_t opcodeToken,
                                            Location loc) {
    if (DECODE_IS_D3D10_SB_OPCODE_EXTENDED(opcodeToken)) {
      return emitError(getLocation(),
                       "extended interface calls are not supported");
    }

    auto callSite = parseToken();
    FAILURE_IF_FAILED(callSite);

    auto operand = parseOperand();
    FAILURE_IF_FAILED(operand);

    return builder.buildInterfaceCall(*operand, *callSite, loc);
  }

  OptionalParseResult parseDclInstruction(uint32_t opcodeToken, Location loc,
                                          Instruction &out) {
    FailureOr<Instruction> result;
    switch (DECODE_D3D10_SB_OPCODE_TYPE(opcodeToken)) {
    case D3D10_SB_OPCODE_DCL_GLOBAL_FLAGS:
      result = parseDclGlobalFlags(opcodeToken, loc);
      break;
    case D3D10_SB_OPCODE_DCL_TEMPS:
      result = parseDclTemps(loc);
      break;
    case D3D11_SB_OPCODE_DCL_INPUT_CONTROL_POINT_COUNT:
      result = parseDclInputControlPointCount(opcodeToken, loc);
      break;
    case D3D11_SB_OPCODE_DCL_OUTPUT_CONTROL_POINT_COUNT:
      result = parseDclOutputControlPointCount(opcodeToken, loc);
      break;
    case D3D11_SB_OPCODE_DCL_TESS_DOMAIN:
      result = parseDclTessellatorDomain(opcodeToken, loc);
      break;
    case D3D11_SB_OPCODE_DCL_TESS_OUTPUT_PRIMITIVE:
      result = parseDclTessellatorOutputPrimitive(opcodeToken, loc);
      break;
    case D3D11_SB_OPCODE_DCL_TESS_PARTITIONING:
      result = parseDclTessellatorPartitioning(opcodeToken, loc);
      break;
    case D3D10_SB_OPCODE_DCL_GS_OUTPUT_PRIMITIVE_TOPOLOGY:
      result = parseDclOutputTopology(opcodeToken, loc);
      break;
    case D3D10_SB_OPCODE_DCL_GS_INPUT_PRIMITIVE:
      result = parseDclInputPrimitive(opcodeToken, loc);
      break;
    case D3D11_SB_OPCODE_DCL_GS_INSTANCE_COUNT:
      result = parseDclGsInstanceCount(loc);
      break;
    case D3D10_SB_OPCODE_DCL_MAX_OUTPUT_VERTEX_COUNT:
      result = parseDclMaxOutputVertexCount(loc);
      break;
    case D3D11_SB_OPCODE_DCL_STREAM:
      result = parseDclStream(loc);
      break;
    case D3D10_SB_OPCODE_DCL_INPUT_PS:
      result = parseDclInputPs(opcodeToken, loc);
      break;
    case D3D10_SB_OPCODE_DCL_INPUT_PS_SIV:
      result = parseDclInputPsSiv(opcodeToken, loc);
      break;
    case D3D10_SB_OPCODE_DCL_INPUT_PS_SGV:
      result = parseDclInputPsSgv(loc);
      break;
    case D3D10_SB_OPCODE_DCL_INPUT:
      result = parseDclInput(loc);
      break;
    case D3D10_SB_OPCODE_DCL_OUTPUT:
      result = parseDclOutput(loc);
      break;
    case D3D10_SB_OPCODE_DCL_INDEX_RANGE:
      result = parseDclIndexRange(loc);
      break;
    case D3D10_SB_OPCODE_DCL_INPUT_SGV:
      result = parseDclInputSgv(loc);
      break;
    case D3D10_SB_OPCODE_DCL_INPUT_SIV:
      result = parseDclInputSiv(loc);
      break;
    case D3D10_SB_OPCODE_DCL_OUTPUT_SGV:
      result = parseDclOutputSgv(loc);
      break;
    case D3D10_SB_OPCODE_DCL_OUTPUT_SIV:
      result = parseDclOutputSiv(loc);
      break;
    case D3D11_SB_OPCODE_DCL_HS_MAX_TESSFACTOR:
      result = parseDclHsMaxTessFactor(loc);
      break;
    case D3D11_SB_OPCODE_DCL_HS_JOIN_PHASE_INSTANCE_COUNT:
      result = parseDclHsJoinPhaseInstanceCount(loc);
      break;
    case D3D11_SB_OPCODE_DCL_HS_FORK_PHASE_INSTANCE_COUNT:
      result = parseDclHsForkPhaseInstanceCount(loc);
      break;
    case D3D11_SB_OPCODE_DCL_THREAD_GROUP_SHARED_MEMORY_RAW:
      result = parseDclTgsmRaw(loc);
      break;
    case D3D11_SB_OPCODE_DCL_THREAD_GROUP_SHARED_MEMORY_STRUCTURED:
      result = parseDclTgsmStructured(loc);
      break;
    case D3D10_SB_OPCODE_DCL_CONSTANT_BUFFER:
      result = parseDclConstantBuffer(opcodeToken, loc);
      break;
    case D3D10_SB_OPCODE_DCL_SAMPLER:
      result = parseDclSampler(opcodeToken, loc);
      break;
    case D3D11_SB_OPCODE_INTERFACE_CALL:
      result = parseInterfaceCall(opcodeToken, loc);
      break;
    default:
      return std::nullopt;
    }
    if (failed(result))
      return failure();
    out = *result;
    return success();
  }

  FailureOr<Instruction> parseInstruction() {
    size_t beginOffset = currentTokenOffset;
    Token token = parseToken();
    if (failed(token))
      return failure();

    FileLineColLoc loc = getLocation();

    uint32_t opcode = DECODE_D3D10_SB_OPCODE_TYPE(*token);
    InstructionModifier modifier;
    modifier.preciseMask = DECODE_D3D11_SB_INSTRUCTION_PRECISE_VALUES(*token);
    modifier.saturate = DECODE_IS_D3D10_SB_INSTRUCTION_SATURATE_ENABLED(*token);

    uint32_t length = DECODE_D3D10_SB_TOKENIZED_INSTRUCTION_LENGTH(*token);

    // TODO: extended instructions:
    // BOOL b51PlusShader =
    // BOOL bExtended = DECODE_IS_D3D10_SB_OPCODE_EXTENDED(Token)
    // ...

    if (opcode >= D3D10_SB_NUM_OPCODES) {
      emitError(getLocation(), "unknown opcode");
      return failure();
    }

    auto opcodeToken = *token;

    Instruction dclInstruction;
    auto parseResult = parseDclInstruction(opcodeToken, loc, dclInstruction);
    if (parseResult.has_value()) {
      if (failed(*parseResult))
        return failure();
      if (failed(verifyInstructionLength(beginOffset, length)))
        return failure();
      return dclInstruction;
    }

    unsigned numOperands = instrInfo[opcode].numOperands;

    SmallVector<Operand, 8> operands;
    for (unsigned i = 0; i < numOperands; ++i) {
      FailureOr<Operand> operand = parseOperand();
      if (failed(operand))
        return failure();
      operands.push_back(*operand);
    }

    if (failed(verifyInstructionLength(beginOffset, length)))
      return failure();

    return builder.buildInstruction(instrInfo[opcode].name, operands, modifier,
                                    loc);
  }

  FailureOr<Module> parseModule() {
    FileLineColLoc loc = getLocation(0);
    std::vector<Instruction> instructions;
    while (currentTokenOffset < buffer.size()) {
      FailureOr<Instruction> inst = parseInstruction();
      if (failed(inst)) {
        return failure();
      }
      instructions.push_back(*inst);
    }
    return builder.buildModule(instructions, loc);
  }

  LogicalResult verifyInstructionLength(size_t beginOffset, uint32_t length) {
    if (((currentTokenOffset - beginOffset) / 4) != length) {
      emitError(getLocation(), "instruction length mismatch");
      return failure();
    }
    return success();
  }

private:
  DXBuilder &builder;
  StringAttr name;
  StringRef buffer;
  size_t currentTokenOffset{0};
  InstructionInfo instrInfo[D3D10_SB_NUM_OPCODES];
};

namespace mlir::dxsa {
OwningOpRef<ModuleOp> importDxsaBinaryToModule(llvm::SourceMgr &source,
                                               MLIRContext *context) {

  if (source.getNumBuffers() != 1) {
    emitError(UnknownLoc::get(context), "one source file should be provided");
    return nullptr;
  }

  uint32_t sourceBufId = source.getMainFileID();
  StringRef buffer = source.getMemoryBuffer(sourceBufId)->getBuffer();
  StringAttr name = StringAttr::get(
      context, source.getMemoryBuffer(sourceBufId)->getBufferIdentifier());

  // FIXME:
  context->allowUnregisteredDialects();
  context->loadAllAvailableDialects();

  DXBuilder builder(context, name);
  Parser parser(builder, name, buffer);
  FailureOr<ModuleOp> mod = parser.parseModule();
  if (failed(mod))
    return {nullptr};
  return {*mod};
}
} // namespace mlir::dxsa