[mlir][dxsa] Add unknown instruction

Example:
  dxsa.dcl_temps 1
  dxsa.unknown <tokens = [0x030007FF, 0xDEADBEEF, 0x12345678]>
  dxsa.dcl_temps 2

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

Author: tagolog

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

@@ -493,6 +493,28 @@ def DXSA_Instruction : DXSA_Op<"instruction"> {
   let assemblyFormat = "$mnemonic $operands attr-dict";
 }
 
+def DXSA_Unknown : DXSA_Op<"unknown"> {
+  let summary = "raw tokens fallback for an undecodable instruction";
+  let description = [{
+    The `dxsa.unknown` operation represents one instruction whose raw
+    tokens could not be decoded into a structured op — unknown opcode,
+    undecodable payload, length mismatch, truncated tail, etc. It acts
+    as a disassembler-style fallback so the surrounding well-formed
+    instructions still appear in IR.
+
+    Example:
+
+    ```mlir
+    dxsa.unknown <tokens = [0x030000FF, 0xDEADBEEF, 0x12345678]>
+    ```
+  }];
+
+  let arguments = (ins DenseI32ArrayAttr:$tokens);
+  let results = (outs);
+  let hasVerifier = 1;
+  let assemblyFormat = "custom<HexTokens>($tokens) attr-dict";
+}
+
 def DXSA_InlineOperandType_Temp                          : I32EnumAttrCase<"temp",                                  0>;
 def DXSA_InlineOperandType_Input                         : I32EnumAttrCase<"input",                                 1>;
 def DXSA_InlineOperandType_Output                        : I32EnumAttrCase<"output",                                2>;

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

@@ -13,6 +13,7 @@
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Format.h"
 
 using namespace mlir;
 using namespace mlir::dxsa;
@@ -35,6 +36,12 @@ void DXSADialect::initialize() {
       >();
 }
 
+/// Declarations for custom-directive helpers used by the
+/// TableGen-generated print/parse methods.
+static ParseResult parseHexTokens(OpAsmParser &parser, DenseI32ArrayAttr &attr);
+static void printHexTokens(OpAsmPrinter &printer, Operation *,
+                           DenseI32ArrayAttr attr);
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
@@ -162,6 +169,50 @@ LogicalResult DclConstantBuffer::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// UnknownOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult Unknown::verify() {
+  if (getTokens().empty())
+    return emitOpError("tokens must not be empty");
+  return success();
+}
+
+/// Parse `<tokens = [0x..., ...]>` for the unknown op.
+static ParseResult parseHexTokens(OpAsmParser &parser,
+                                  DenseI32ArrayAttr &attr) {
+  SmallVector<int32_t> tokens;
+  auto parseOneToken = [&]() -> ParseResult {
+    uint32_t value;
+    if (parser.parseInteger(value))
+      return failure();
+    tokens.push_back(static_cast<int32_t>(value));
+    return success();
+  };
+
+  if (parser.parseLess() || parser.parseKeyword("tokens") ||
+      parser.parseEqual() ||
+      parser.parseCommaSeparatedList(OpAsmParser::Delimiter::Square,
+                                     parseOneToken) ||
+      parser.parseGreater())
+    return failure();
+
+  attr = DenseI32ArrayAttr::get(parser.getContext(), tokens);
+  return success();
+}
+
+/// Print the tokens array as uppercase, 8-digit, 0x-prefixed hex.
+static void printHexTokens(OpAsmPrinter &printer, Operation *,
+                           DenseI32ArrayAttr attr) {
+  printer << "<tokens = [";
+  llvm::interleaveComma(attr.asArrayRef(), printer.getStream(), [&](int32_t t) {
+    printer.getStream() << llvm::format_hex(static_cast<uint32_t>(t),
+                                            /*Width=*/10, /*Upper=*/true);
+  });
+  printer << "]>";
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd attribute method definitions
 //===----------------------------------------------------------------------===//

mlir/lib/Target/DXSA/BinaryParser.cpp

@@ -9,6 +9,7 @@
 #include "mlir/Target/DXSA/BinaryParser.h"
 #include "mlir/Dialect/DXSA/IR/DXSA.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Location.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
@@ -531,6 +532,43 @@ class DXBuilder {
                                      builder.getStringAttr(name));
   }
 
+  Instruction buildUnknown(ArrayRef<int32_t> tokens, Location loc) {
+    return dxsa::Unknown::create(
+        builder, loc, DenseI32ArrayAttr::get(builder.getContext(), tokens));
+  }
+
+  /// RAII guard for speculative IR construction. On destruction, erases
+  /// every op appended to the active insertion block after the guard
+  /// was created, unless release() was called first.
+  class RewindGuard {
+  public:
+    explicit RewindGuard(DXBuilder &dxBuilder)
+        : builder(dxBuilder.builder),
+          block(dxBuilder.builder.getInsertionBlock()),
+          numOps(block->getOperations().size()) {}
+
+    ~RewindGuard() {
+      if (released)
+        return;
+      while (block->getOperations().size() > numOps)
+        block->back().erase();
+      builder.setInsertionPointToEnd(block);
+    }
+
+    RewindGuard(const RewindGuard &) = delete;
+    RewindGuard &operator=(const RewindGuard &) = delete;
+
+    /// Accept the speculative work: keep all ops added since
+    /// construction and disarm the destructor.
+    void release() { released = true; }
+
+  private:
+    OpBuilder &builder;
+    Block *block;
+    size_t numOps;
+    bool released = false;
+  };
+
   Instruction buildDclGlobalFlags(dxsa::GlobalFlags flags, Location loc) {
     auto flagsAttr = dxsa::GlobalFlagsAttr::get(builder.getContext(), flags);
     return dxsa::DclGlobalFlags::create(builder, loc, flagsAttr);
@@ -726,9 +764,11 @@ class Parser {
   using Instruction = DXBuilder::Instruction;
   using Module = DXBuilder::Module;
 
+  /// Width of the token in the program binary stream.
+  static constexpr size_t tokenSize = sizeof(uint32_t);
+
   /// 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();
@@ -1531,6 +1571,110 @@ class Parser {
                                     loc);
   }
 
+  struct InstructionSlice {
+    size_t beginOffset; // absolute byte offset in `buffer`
+    uint32_t numTokens; // number of tokens to consume (>= 1)
+    StringRef hint;     // short reason, included in the warning
+  };
+
+  std::optional<InstructionSlice> sliceInstruction() {
+    auto beginOffset = currentTokenOffset;
+    auto remainingBytes = buffer.size() - beginOffset;
+    if (remainingBytes < tokenSize)
+      return std::nullopt;
+    auto token0 = support::endian::read<uint32_t>(buffer.begin() + beginOffset,
+                                                  endianness::little);
+    auto opcode = DECODE_D3D10_SB_OPCODE_TYPE(token0);
+    // CUSTOMDATA: total token count lives in the second token.
+    // The minimum legal value is 2 (the header pair itself).
+    if (opcode == D3D10_SB_OPCODE_CUSTOMDATA) {
+      if (remainingBytes < 2 * tokenSize)
+        return InstructionSlice{beginOffset, 1, "truncated customdata header"};
+      auto token1 = support::endian::read<uint32_t>(
+          buffer.begin() + beginOffset + tokenSize, endianness::little);
+      if (token1 < 2)
+        return InstructionSlice{beginOffset, 2, "customdata length < 2"};
+      auto instructionSize = static_cast<uint64_t>(token1) * tokenSize;
+      if (instructionSize > remainingBytes)
+        return InstructionSlice{
+            beginOffset, static_cast<uint32_t>(remainingBytes / tokenSize),
+            "customdata length past EOF"};
+      return InstructionSlice{beginOffset, token1, "customdata block"};
+    }
+
+    // Regular instruction
+    auto length = DECODE_D3D10_SB_TOKENIZED_INSTRUCTION_LENGTH(token0);
+    if (length == 0)
+      return InstructionSlice{beginOffset, 1,
+                              "instruction length encoded as 0"};
+    auto instructionSize = static_cast<uint64_t>(length) * tokenSize;
+    if (instructionSize > remainingBytes)
+      return InstructionSlice{beginOffset,
+                              static_cast<uint32_t>(remainingBytes / tokenSize),
+                              "instruction length past EOF"};
+    return InstructionSlice{beginOffset, length, "structured decode failed"};
+  }
+
+  /// Read `numTokens` tokens from `beginOffset` and return them as
+  /// i32s in little-endian order. Caller must guarantee the range is
+  /// in-bounds.
+  SmallVector<int32_t> snapshotTokens(size_t beginOffset, uint32_t numTokens) {
+    SmallVector<int32_t> out(numTokens);
+    for (uint32_t i = 0; i < numTokens; ++i)
+      out[i] = static_cast<int32_t>(support::endian::read<uint32_t>(
+          buffer.begin() + beginOffset + i * tokenSize, endianness::little));
+    return out;
+  }
+
+  /// Run the structured parser over the slice.
+  /// On success the produced ops are kept.
+  /// On failure any partial IR is rewound by guard.
+  bool tryParseAsStructured(const InstructionSlice &slice) {
+    auto expectedEnd = slice.beginOffset + slice.numTokens * tokenSize;
+    DXBuilder::RewindGuard guard(builder);
+    ScopedDiagnosticHandler suppress(name.getContext(),
+                                     [](Diagnostic &) { return success(); });
+    if (succeeded(parseInstruction()) && currentTokenOffset == expectedEnd) {
+      guard.release();
+      return true;
+    }
+    return false;
+  }
+
+  /// Emit an unknown op for the slice.
+  LogicalResult parseAsUnknown(const InstructionSlice &slice) {
+    currentTokenOffset = slice.beginOffset + slice.numTokens * tokenSize;
+    auto loc = FileLineColLoc::get(name, 0, slice.beginOffset);
+    SmallVector<int32_t> rawTokens =
+        snapshotTokens(slice.beginOffset, slice.numTokens);
+    builder.buildUnknown(rawTokens, loc);
+    emitWarning(loc) << "could not decode instruction (" << slice.hint
+                     << "); emitted dxsa.unknown with " << slice.numTokens
+                     << " token(s)";
+    return success();
+  }
+
+  /// Try structured parse; on any failure (or partial consumption)
+  /// rewind any IR built during the attempt, advance the byte cursor
+  /// past the slice, and emit `dxsa.unknown` covering the same range.
+  LogicalResult parseInstructionOrUnknown() {
+    auto slice = sliceInstruction();
+    if (!slice) {
+      // Nothing to wrap in an op.
+      auto trailingBytes = buffer.size() - currentTokenOffset;
+      if (trailingBytes > 0)
+        emitWarning(getLocation(0))
+            << "ignoring " << trailingBytes << " trailing byte(s)";
+      currentTokenOffset = buffer.size();
+      return success();
+    }
+
+    if (tryParseAsStructured(*slice))
+      return success();
+
+    return parseAsUnknown(*slice);
+  }
+
   FailureOr<Module> parseModule() {
     FileLineColLoc loc = getLocation(0);
     auto header = parseProgramHeader();
@@ -1545,10 +1689,8 @@ class Parser {
     }
     auto module = builder.createModule(programType, shaderVersion, loc);
     while (currentTokenOffset < buffer.size()) {
-      FailureOr<Instruction> inst = parseInstruction();
-      if (failed(inst)) {
+      if (failed(parseInstructionOrUnknown()))
         return failure();
-      }
     }
     return module;
   }

mlir/test/Target/DXSA/inputs/unknown.bin

@@ -0,0 +1,10 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/unknown.bin | FileCheck %s
+
+// CHECK:      dxsa.module {
+// CHECK-NEXT:   dxsa.dcl_temps 1
+// CHECK-NEXT:   dxsa.unknown <tokens = [0x030007FF, 0xDEADBEEF, 0x12345678]>
+// CHECK-NEXT:   dxsa.dcl_temps 2
+// CHECK-NEXT:   dxsa.unknown <tokens = [0x00000035, 0x00000004, 0x11111111, 0x22222222]>
+// CHECK-NEXT:   dxsa.dcl_temps 3
+// CHECK-NEXT:   dxsa.unknown <tokens = [0x03000068, 0x00000005, 0xCAFEBABE]>
+// CHECK-NEXT: }

mlir/test/Target/DXSA/inputs/unknown_past_eof.bin

@@ -0,0 +1,4 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// expected-error@+1 {{'dxsa.unknown' op tokens must not be empty}}
+dxsa.unknown <tokens = []>

mlir/test/Target/DXSA/unknown.mlir

@@ -0,0 +1,6 @@
+// RUN: mlir-translate --import-dxsa-bin %S/inputs/unknown_past_eof.bin | FileCheck %s
+
+// CHECK:      dxsa.module {
+// CHECK-NEXT:   dxsa.dcl_temps 1
+// CHECK-NEXT:   dxsa.unknown <tokens = [0x050007FF, 0xAAAAAAAA, 0xBBBBBBBB]>
+// CHECK-NEXT: }