[mlir][dxsa] Add unknown instruction

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

@@ -15,6 +15,112 @@ include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
 include "mlir/IR/EnumAttr.td"
 
+//===----------------------------------------------------------------------===//
+// DXSA op base class
+//===----------------------------------------------------------------------===//
+
+// Base class for all operations in this dialect.
+class DXSA_Op<string mnemonic, list<Trait> traits = []> :
+    Op<DXSADialect, mnemonic, traits>;
+
+//===----------------------------------------------------------------------===//
+// DXSA module — top-level container op for a DXBC tokenized program
+//===----------------------------------------------------------------------===//
+
+def DXSA_ProgramType_PixelShader         : I32EnumAttrCase<"pixel_shader",         0>;
+def DXSA_ProgramType_VertexShader        : I32EnumAttrCase<"vertex_shader",        1>;
+def DXSA_ProgramType_GeometryShader      : I32EnumAttrCase<"geometry_shader",      2>;
+def DXSA_ProgramType_HullShader          : I32EnumAttrCase<"hull_shader",          3>;
+def DXSA_ProgramType_DomainShader        : I32EnumAttrCase<"domain_shader",        4>;
+def DXSA_ProgramType_ComputeShader       : I32EnumAttrCase<"compute_shader",       5>;
+def DXSA_ProgramType_MeshShader          : I32EnumAttrCase<"mesh_shader",         13>;
+def DXSA_ProgramType_AmplificationShader : I32EnumAttrCase<"amplification_shader",14>;
+
+def DXSA_ProgramType : I32EnumAttr<
+    "ProgramType", "DXBC tokenized program type", [
+      DXSA_ProgramType_PixelShader,
+      DXSA_ProgramType_VertexShader,
+      DXSA_ProgramType_GeometryShader,
+      DXSA_ProgramType_HullShader,
+      DXSA_ProgramType_DomainShader,
+      DXSA_ProgramType_ComputeShader,
+      DXSA_ProgramType_MeshShader,
+      DXSA_ProgramType_AmplificationShader
+    ]> {
+  let cppNamespace = "::mlir::dxsa";
+  let genSpecializedAttr = 0;
+}
+
+def DXSA_ProgramTypeAttr :
+    EnumAttr<DXSADialect, DXSA_ProgramType, "program_type"> {
+  let assemblyFormat = "$value";
+}
+
+def DXSA_ShaderVersionAttr : AttrDef<DXSADialect, "ShaderVersion"> {
+  let mnemonic = "shader_version";
+  let summary = "DXBC shader version (major.minor)";
+  let description = [{
+    The `#dxsa.shader_version` attribute holds the major and minor
+    components of shader model version.
+
+    Example:
+
+    ```mlir
+    #dxsa.shader_version<5, 0>
+    ```
+  }];
+  let parameters = (ins "uint8_t":$major, "uint8_t":$minor);
+  let assemblyFormat = "`<` $major `,` $minor `>`";
+}
+
+def DXSA_ModuleOp : DXSA_Op<"module", [
+    IsolatedFromAbove, NoRegionArguments, NoTerminator, SingleBlock]> {
+  let summary = "the top-level container for a shader program";
+  let description = [{
+    The `dxsa.module` operation models the top-level container of a single
+    shader tokenized program (one SHEX section of the DXBC binary).
+
+    The two optional attributes are shader program type and version.
+    Both attributes are either both present (real binary with a SHEX
+    header) or both absent (header-less raw token streams).
+
+    Example:
+
+    ```mlir
+    // Binary content with a SHEX header
+    dxsa.module pixel_shader 5 0 {
+      dxsa.dcl_global_flags <refactoringAllowed>
+    }
+
+    // Binary content without a SHEX header
+    dxsa.module {
+      dxsa.dcl_global_flags <refactoringAllowed>
+    }
+    ```
+  }];
+
+  let arguments = (ins
+    OptionalAttr<DXSA_ProgramTypeAttr>:$program_type,
+    OptionalAttr<DXSA_ShaderVersionAttr>:$shader_version);
+  let regions = (region SizedRegion<1>:$body);
+
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+
+  let skipDefaultBuilders = 1;
+  let builders = [
+    OpBuilder<(ins
+      CArg<"::mlir::dxsa::ProgramTypeAttr",
+           "::mlir::dxsa::ProgramTypeAttr()">:$programType,
+      CArg<"::mlir::dxsa::ShaderVersionAttr",
+           "::mlir::dxsa::ShaderVersionAttr()">:$shaderVersion)>
+  ];
+
+  let extraClassDeclaration = [{
+    ::mlir::Block *getBodyBlock() { return &getBody().front(); }
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // DXSA enum definitions
 //===----------------------------------------------------------------------===//
@@ -313,10 +419,6 @@ def DXSA_UIntNonZero : AttrConstraint<
 // DXSA op definitions
 //===----------------------------------------------------------------------===//
 
-// Base class for the operation in this dialect
-class DXSA_Op<string mnemonic, list<Trait> traits = []> :
-    Op<DXSADialect, mnemonic, traits>;
-
 def DXSA_Operand : DXSA_Op<"operand"> {
   let summary = "defines an operand of an instruction";
   let description = [{
@@ -391,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/include/mlir/Target/DXSA/BinaryParser.h

@@ -9,18 +9,20 @@
 #ifndef MLIR_TARGET_DXSA_BINARYPARSER_H
 #define MLIR_TARGET_DXSA_BINARYPARSER_H
 
+#include "mlir/Dialect/DXSA/IR/DXSA.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OwningOpRef.h"
 #include "llvm/Support/SourceMgr.h"
 
 namespace mlir::dxsa {
+/// Deserializes the given binary \p source and creates a MLIR ModuleOp in the
+/// given \p context.
+OwningOpRef<dxsa::ModuleOp> deserialize(llvm::SourceMgr &source,
+                                        MLIRContext *context);
 
-/// Decode DXSA binary \p source and return an MLIR module.
-OwningOpRef<ModuleOp> importDxsaBinaryToModule(llvm::SourceMgr &source,
-                                               MLIRContext *context);
-/// Encode \p source to DXSA binary.
-LogicalResult exportModuleToDxsaBinary(ModuleOp source, raw_ostream &output);
+/// Serializes the given MLIR \p moduleOp and writes to \p output.
+LogicalResult serialize(mlir::ModuleOp moduleOp, raw_ostream &output);
 } // namespace mlir::dxsa
 
 #endif // MLIR_TARGET_DXSA_BINARYPARSER_H

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

@@ -10,8 +10,10 @@
 
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
+#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;
@@ -34,13 +36,87 @@ 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
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/DXSA/IR/DXSAOps.cpp.inc"
 
+//===----------------------------------------------------------------------===//
+// ModuleOp
+//===----------------------------------------------------------------------===//
+
+void ModuleOp::build(OpBuilder &builder, OperationState &state,
+                     ProgramTypeAttr programType,
+                     ShaderVersionAttr shaderVersion) {
+  if (programType)
+    state.addAttribute("program_type", programType);
+  if (shaderVersion)
+    state.addAttribute("shader_version", shaderVersion);
+  OpBuilder::InsertionGuard guard(builder);
+  builder.createBlock(state.addRegion());
+}
+
+ParseResult ModuleOp::parse(OpAsmParser &parser, OperationState &result) {
+  // Parse optional shader information like `pixel_shader 5 0`.
+  StringRef typeKeyword;
+  auto typeLoc = parser.getCurrentLocation();
+  if (succeeded(parser.parseOptionalKeyword(&typeKeyword))) {
+    auto programType = symbolizeProgramType(typeKeyword);
+    if (!programType)
+      return parser.emitError(typeLoc)
+             << "unknown program type: " << typeKeyword;
+    result.addAttribute("program_type", ProgramTypeAttr::get(
+                                            parser.getContext(), *programType));
+
+    uint8_t major = 0, minor = 0;
+    if (parser.parseInteger(major) || parser.parseInteger(minor))
+      return failure();
+    result.addAttribute(
+        "shader_version",
+        ShaderVersionAttr::get(parser.getContext(), major, minor));
+  }
+
+  Region *body = result.addRegion();
+  if (parser.parseOptionalAttrDictWithKeyword(result.attributes) ||
+      parser.parseRegion(*body, /*arguments=*/{}))
+    return failure();
+
+  if (body->empty())
+    body->push_back(new Block());
+
+  return success();
+}
+
+void ModuleOp::print(OpAsmPrinter &printer) {
+  if (auto programType = getProgramType()) {
+    printer << ' ' << stringifyProgramType(*programType);
+    auto version = getShaderVersionAttr();
+    printer << ' ' << static_cast<unsigned>(version.getMajor()) << ' '
+            << static_cast<unsigned>(version.getMinor());
+  }
+  printer.printOptionalAttrDictWithKeyword((*this)->getAttrs(),
+                                           {"program_type", "shader_version"});
+  printer << ' ';
+  printer.printRegion(getBody());
+}
+
+LogicalResult ModuleOp::verify() {
+  bool hasType = static_cast<bool>(getProgramTypeAttr());
+  bool hasVersion = static_cast<bool>(getShaderVersionAttr());
+  if (hasType != hasVersion)
+    return emitOpError(
+        "program_type and shader_version must both be present or both absent");
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Op verifiers
 //===----------------------------------------------------------------------===//
@@ -93,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
 //===----------------------------------------------------------------------===//