From 6b59659ce5fbe38fd75d39c750a3f71b24ffa3cd Mon Sep 17 00:00:00 2001 From: Vladimir Shiryaev Date: Sun, 31 May 2026 00:56:39 -0700 Subject: [PATCH 1/2] [mlir][dxsa] Add root dxsa.module Wraps the program into a module with optional attributes program type and shader version. When the binary has no header both attributes are omitted. Example: dxsa.module pixel_shader 5 0 { dxsa.dcl_global_flags } dxsa.module { dxsa.dcl_global_flags } Signed-off-by: Vladimir Shiryaev --- mlir/include/mlir/Dialect/DXSA/IR/DXSAOps.td | 110 +++++++++++++++++- mlir/include/mlir/Target/DXSA/BinaryParser.h | 12 +- mlir/lib/Dialect/DXSA/IR/DXSA.cpp | 69 +++++++++++ mlir/lib/Target/DXSA/BinaryParser.cpp | 80 ++++++++++--- mlir/lib/Target/DXSA/BinaryWriter.cpp | 2 +- .../lib/Target/DXSA/TranslateRegistration.cpp | 4 +- .../test/Target/DXSA/dcl_constant_buffer.mlir | 2 +- mlir/test/Target/DXSA/dcl_global_flags.mlir | 2 +- .../Target/DXSA/dcl_gs_instance_count.mlir | 2 +- .../dcl_hs_fork_phase_instance_count.mlir | 2 +- .../dcl_hs_join_phase_instance_count.mlir | 2 +- .../Target/DXSA/dcl_hs_max_tessfactor.mlir | 2 +- mlir/test/Target/DXSA/dcl_index_range.mlir | 2 +- .../DXSA/dcl_input_control_point_count.mlir | 2 +- .../test/Target/DXSA/dcl_input_primitive.mlir | 2 +- .../DXSA/dcl_max_output_vertex_count.mlir | 2 +- .../DXSA/dcl_output_control_point_count.mlir | 2 +- mlir/test/Target/DXSA/dcl_output_sgv.mlir | 2 +- mlir/test/Target/DXSA/dcl_output_siv.mlir | 2 +- .../test/Target/DXSA/dcl_output_topology.mlir | 2 +- mlir/test/Target/DXSA/dcl_temps.mlir | 2 +- .../Target/DXSA/dcl_tessellator_domain.mlir | 2 +- .../dcl_tessellator_output_primitive.mlir | 2 +- .../DXSA/dcl_tessellator_partitioning.mlir | 2 +- mlir/test/Target/DXSA/dcl_tgsm_raw.mlir | 2 +- .../test/Target/DXSA/dcl_tgsm_structured.mlir | 2 +- mlir/test/Target/DXSA/empty.mlir | 4 +- .../DXSA/inputs/program_header_absent.bin | Bin 0 -> 12 bytes .../DXSA/inputs/program_header_present.bin | Bin 0 -> 20 bytes mlir/test/Target/DXSA/mov-index.mlir | 2 +- mlir/test/Target/DXSA/mov.mlir | 2 +- .../Target/DXSA/program_header_absent.mlir | 6 + .../Target/DXSA/program_header_present.mlir | 6 + mlir/test/Target/DXSA/ret.mlir | 2 +- mlir/test/Target/DXSA/udiv.mlir | 2 +- 35 files changed, 287 insertions(+), 54 deletions(-) create mode 100644 mlir/test/Target/DXSA/inputs/program_header_absent.bin create mode 100644 mlir/test/Target/DXSA/inputs/program_header_present.bin create mode 100644 mlir/test/Target/DXSA/program_header_absent.mlir create mode 100644 mlir/test/Target/DXSA/program_header_present.mlir diff --git a/mlir/include/mlir/Dialect/DXSA/IR/DXSAOps.td b/mlir/include/mlir/Dialect/DXSA/IR/DXSAOps.td index d718879c353d..857db9954226 100644 --- a/mlir/include/mlir/Dialect/DXSA/IR/DXSAOps.td +++ b/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 traits = []> : + Op; + +//===----------------------------------------------------------------------===// +// 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 { + let assemblyFormat = "$value"; +} + +def DXSA_ShaderVersionAttr : AttrDef { + 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 + } + + // Binary content without a SHEX header + dxsa.module { + dxsa.dcl_global_flags + } + ``` + }]; + + let arguments = (ins + OptionalAttr:$program_type, + OptionalAttr:$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 traits = []> : - Op; - def DXSA_Operand : DXSA_Op<"operand"> { let summary = "defines an operand of an instruction"; let description = [{ diff --git a/mlir/include/mlir/Target/DXSA/BinaryParser.h b/mlir/include/mlir/Target/DXSA/BinaryParser.h index 3b81b8772a06..957ad82d1155 100644 --- a/mlir/include/mlir/Target/DXSA/BinaryParser.h +++ b/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 deserialize(llvm::SourceMgr &source, + MLIRContext *context); -/// Decode DXSA binary \p source and return an MLIR module. -OwningOpRef 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 diff --git a/mlir/lib/Dialect/DXSA/IR/DXSA.cpp b/mlir/lib/Dialect/DXSA/IR/DXSA.cpp index 5e4e2abfde24..365ee1dab82a 100644 --- a/mlir/lib/Dialect/DXSA/IR/DXSA.cpp +++ b/mlir/lib/Dialect/DXSA/IR/DXSA.cpp @@ -10,6 +10,7 @@ #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" @@ -41,6 +42,74 @@ void DXSADialect::initialize() { #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(version.getMajor()) << ' ' + << static_cast(version.getMinor()); + } + printer.printOptionalAttrDictWithKeyword((*this)->getAttrs(), + {"program_type", "shader_version"}); + printer << ' '; + printer.printRegion(getBody()); +} + +LogicalResult ModuleOp::verify() { + bool hasType = static_cast(getProgramTypeAttr()); + bool hasVersion = static_cast(getShaderVersionAttr()); + if (hasType != hasVersion) + return emitOpError( + "program_type and shader_version must both be present or both absent"); + return success(); +} + //===----------------------------------------------------------------------===// // Op verifiers //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Target/DXSA/BinaryParser.cpp b/mlir/lib/Target/DXSA/BinaryParser.cpp index 9227e71d4e8a..20c81da4711e 100644 --- a/mlir/lib/Target/DXSA/BinaryParser.cpp +++ b/mlir/lib/Target/DXSA/BinaryParser.cpp @@ -420,15 +420,23 @@ static dxsa::ComponentMask decodeComponentMask(uint32_t rawComponentMask) { class DXBuilder { public: - DXBuilder(MLIRContext *context, StringAttr name) - : context(context), - module(ModuleOp::create(builder, FileLineColLoc::get(name, 0, 0))), - builder(module.getRegion()) {} + explicit DXBuilder(MLIRContext *context) + : context(context), builder(context) {} using Index = mlir::Value; using Operand = mlir::Value; using Instruction = mlir::Operation *; - using Module = mlir::ModuleOp; + using Module = mlir::dxsa::ModuleOp; + + Module createModule(mlir::dxsa::ProgramTypeAttr programType, + mlir::dxsa::ShaderVersionAttr shaderVersion, + Location loc) { + OperationState state(loc, Module::getOperationName()); + Module::build(builder, state, programType, shaderVersion); + auto module = cast(Operation::create(state)); + builder.setInsertionPointToStart(&module.getBody().front()); + return module; + } Index buildIndexImm32(int32_t imm, FileLineColLoc loc) { Operation *op = @@ -523,10 +531,6 @@ class DXBuilder { builder.getStringAttr(name)); } - Module buildModule(ArrayRef 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); @@ -706,7 +710,6 @@ class DXBuilder { private: MLIRContext *context; - ModuleOp module; OpBuilder builder; }; @@ -1530,15 +1533,60 @@ class Parser { FailureOr parseModule() { FileLineColLoc loc = getLocation(0); - std::vector instructions; + auto header = parseProgramHeader(); + FAILURE_IF_FAILED(header); + mlir::dxsa::ProgramTypeAttr programType; + mlir::dxsa::ShaderVersionAttr shaderVersion; + if (*header) { + programType = + mlir::dxsa::ProgramTypeAttr::get(name.getContext(), (*header)->type); + shaderVersion = mlir::dxsa::ShaderVersionAttr::get( + name.getContext(), (*header)->major, (*header)->minor); + } + auto module = builder.createModule(programType, shaderVersion, loc); while (currentTokenOffset < buffer.size()) { FailureOr inst = parseInstruction(); if (failed(inst)) { return failure(); } - instructions.push_back(*inst); } - return builder.buildModule(instructions, loc); + return module; + } + + struct ProgramHeader { + mlir::dxsa::ProgramType type; + uint8_t major; + uint8_t minor; + }; + + /// If the buffer begins with a tokenized-program header (VersionToken + + /// LengthToken), decode and consume both tokens and return the program type + /// and shader model. Otherwise return without touching the parser current + /// position. + FailureOr> parseProgramHeader() { + constexpr size_t headerSize = 2 * sizeof(uint32_t); + if (currentTokenOffset + headerSize > buffer.size()) + return std::optional{}; + + auto versionToken = support::endian::read( + buffer.begin() + currentTokenOffset, endianness::little); + if (DECODE_D3D10_SB_TOKENIZED_INSTRUCTION_LENGTH(versionToken) != 0) + return std::optional{}; + + auto rawType = static_cast( + DECODE_D3D10_SB_TOKENIZED_PROGRAM_TYPE(versionToken)); + auto programType = dxsa::symbolizeProgramType(rawType); + if (!programType) + return std::optional{}; + + auto major = static_cast( + DECODE_D3D10_SB_TOKENIZED_PROGRAM_MAJOR_VERSION(versionToken)); + auto minor = static_cast( + DECODE_D3D10_SB_TOKENIZED_PROGRAM_MINOR_VERSION(versionToken)); + + FAILURE_IF_FAILED(parseToken()); // VersionToken + FAILURE_IF_FAILED(parseToken()); // LengthToken + return std::optional{{*programType, major, minor}}; } LogicalResult verifyInstructionLength(size_t beginOffset, uint32_t length) { @@ -1558,8 +1606,8 @@ class Parser { }; namespace mlir::dxsa { -OwningOpRef importDxsaBinaryToModule(llvm::SourceMgr &source, - MLIRContext *context) { +OwningOpRef deserialize(llvm::SourceMgr &source, + MLIRContext *context) { if (source.getNumBuffers() != 1) { emitError(UnknownLoc::get(context), "one source file should be provided"); @@ -1575,7 +1623,7 @@ OwningOpRef importDxsaBinaryToModule(llvm::SourceMgr &source, context->allowUnregisteredDialects(); context->loadAllAvailableDialects(); - DXBuilder builder(context, name); + DXBuilder builder(context); Parser parser(builder, name, buffer); FailureOr mod = parser.parseModule(); if (failed(mod)) diff --git a/mlir/lib/Target/DXSA/BinaryWriter.cpp b/mlir/lib/Target/DXSA/BinaryWriter.cpp index 4c19d41a7790..d5d83a927653 100644 --- a/mlir/lib/Target/DXSA/BinaryWriter.cpp +++ b/mlir/lib/Target/DXSA/BinaryWriter.cpp @@ -18,7 +18,7 @@ using namespace mlir; using namespace llvm; namespace mlir::dxsa { -LogicalResult exportModuleToDxsaBinary(ModuleOp source, raw_ostream &output) { +LogicalResult serialize(mlir::ModuleOp source, raw_ostream &output) { Region ®ion = source.getRegion(); assert(region.hasOneBlock() && "invalid module"); return failure(); diff --git a/mlir/lib/Target/DXSA/TranslateRegistration.cpp b/mlir/lib/Target/DXSA/TranslateRegistration.cpp index 5dcf9307668a..2b04fbe0954d 100644 --- a/mlir/lib/Target/DXSA/TranslateRegistration.cpp +++ b/mlir/lib/Target/DXSA/TranslateRegistration.cpp @@ -19,7 +19,7 @@ void registerFromDxsaBinTranslation() { "import-dxsa-bin", "Translate DXSA binary to MLIR", [](llvm::SourceMgr &sourceMgr, MLIRContext *context) -> OwningOpRef { - return dxsa::importDxsaBinaryToModule(sourceMgr, context); + return dxsa::deserialize(sourceMgr, context); }, [](DialectRegistry ®istry) { registry.insert(); }}; } @@ -28,7 +28,7 @@ void registerToDxsaBinTranslation() { TranslateFromMLIRRegistration registration{ "export-dxsa-bin", "Translate MLIR to DXSA binary", [](ModuleOp source, raw_ostream &output) { - return dxsa::exportModuleToDxsaBinary(source, output); + return dxsa::serialize(source, output); }, [](DialectRegistry ®istry) { registry.insert(); }}; } diff --git a/mlir/test/Target/DXSA/dcl_constant_buffer.mlir b/mlir/test/Target/DXSA/dcl_constant_buffer.mlir index 2e2cf2ec2598..5115eda138ce 100644 --- a/mlir/test/Target/DXSA/dcl_constant_buffer.mlir +++ b/mlir/test/Target/DXSA/dcl_constant_buffer.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_constant_buffer.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_constant_buffer , // CHECK-NEXT: dxsa.dcl_constant_buffer , // CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/dcl_global_flags.mlir b/mlir/test/Target/DXSA/dcl_global_flags.mlir index 9b33da52b3ba..21bf8d409902 100644 --- a/mlir/test/Target/DXSA/dcl_global_flags.mlir +++ b/mlir/test/Target/DXSA/dcl_global_flags.mlir @@ -2,7 +2,7 @@ // 9 individual flags (one per instruction), then multi-flag combo, then all flags. -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_global_flags // CHECK-NEXT: dxsa.dcl_global_flags // CHECK-NEXT: dxsa.dcl_global_flags diff --git a/mlir/test/Target/DXSA/dcl_gs_instance_count.mlir b/mlir/test/Target/DXSA/dcl_gs_instance_count.mlir index cafa2a29de41..1e7094720dfa 100644 --- a/mlir/test/Target/DXSA/dcl_gs_instance_count.mlir +++ b/mlir/test/Target/DXSA/dcl_gs_instance_count.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_gs_instance_count.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_gs_instance_count 1 // CHECK-NEXT: dxsa.dcl_gs_instance_count 32 // CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/dcl_hs_fork_phase_instance_count.mlir b/mlir/test/Target/DXSA/dcl_hs_fork_phase_instance_count.mlir index cc16e49cd0d6..c7ff78721ddd 100644 --- a/mlir/test/Target/DXSA/dcl_hs_fork_phase_instance_count.mlir +++ b/mlir/test/Target/DXSA/dcl_hs_fork_phase_instance_count.mlir @@ -1,5 +1,5 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_hs_fork_phase_instance_count.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_hs_fork_phase_instance_count 42 // CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/dcl_hs_join_phase_instance_count.mlir b/mlir/test/Target/DXSA/dcl_hs_join_phase_instance_count.mlir index 03e92f30645c..7890d278015e 100644 --- a/mlir/test/Target/DXSA/dcl_hs_join_phase_instance_count.mlir +++ b/mlir/test/Target/DXSA/dcl_hs_join_phase_instance_count.mlir @@ -1,5 +1,5 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_hs_join_phase_instance_count.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_hs_join_phase_instance_count 42 // CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/dcl_hs_max_tessfactor.mlir b/mlir/test/Target/DXSA/dcl_hs_max_tessfactor.mlir index 545574e059af..887c729d72c5 100644 --- a/mlir/test/Target/DXSA/dcl_hs_max_tessfactor.mlir +++ b/mlir/test/Target/DXSA/dcl_hs_max_tessfactor.mlir @@ -1,5 +1,5 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_hs_max_tessfactor.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_hs_max_tessfactor 4.200000e+01 // CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/dcl_index_range.mlir b/mlir/test/Target/DXSA/dcl_index_range.mlir index 1f43055a3264..7ca6ac792ea9 100644 --- a/mlir/test/Target/DXSA/dcl_index_range.mlir +++ b/mlir/test/Target/DXSA/dcl_index_range.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_index_range.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_index_range , index = [4]>, 6 // CHECK-NEXT: dxsa.dcl_index_range , index = [0]>, 4 // CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/dcl_input_control_point_count.mlir b/mlir/test/Target/DXSA/dcl_input_control_point_count.mlir index 0021300a7110..a8a0d5d7433c 100644 --- a/mlir/test/Target/DXSA/dcl_input_control_point_count.mlir +++ b/mlir/test/Target/DXSA/dcl_input_control_point_count.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_input_control_point_count.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_input_control_point_count 1 // CHECK-NEXT: dxsa.dcl_input_control_point_count 3 // CHECK-NEXT: dxsa.dcl_input_control_point_count 16 diff --git a/mlir/test/Target/DXSA/dcl_input_primitive.mlir b/mlir/test/Target/DXSA/dcl_input_primitive.mlir index 886157b1fc76..035045b6f361 100644 --- a/mlir/test/Target/DXSA/dcl_input_primitive.mlir +++ b/mlir/test/Target/DXSA/dcl_input_primitive.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_input_primitive.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_input_primitive point // CHECK-NEXT: dxsa.dcl_input_primitive line // CHECK-NEXT: dxsa.dcl_input_primitive triangle diff --git a/mlir/test/Target/DXSA/dcl_max_output_vertex_count.mlir b/mlir/test/Target/DXSA/dcl_max_output_vertex_count.mlir index d2ff559e0715..934b7ead8ab3 100644 --- a/mlir/test/Target/DXSA/dcl_max_output_vertex_count.mlir +++ b/mlir/test/Target/DXSA/dcl_max_output_vertex_count.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_max_output_vertex_count.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_max_output_vertex_count 1 // CHECK-NEXT: dxsa.dcl_max_output_vertex_count 1024 // CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/dcl_output_control_point_count.mlir b/mlir/test/Target/DXSA/dcl_output_control_point_count.mlir index 039b866dd75c..1f232bd17787 100644 --- a/mlir/test/Target/DXSA/dcl_output_control_point_count.mlir +++ b/mlir/test/Target/DXSA/dcl_output_control_point_count.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_output_control_point_count.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_output_control_point_count 0 // CHECK-NEXT: dxsa.dcl_output_control_point_count 1 // CHECK-NEXT: dxsa.dcl_output_control_point_count 4 diff --git a/mlir/test/Target/DXSA/dcl_output_sgv.mlir b/mlir/test/Target/DXSA/dcl_output_sgv.mlir index 2691bb5e502e..aad849a5efab 100644 --- a/mlir/test/Target/DXSA/dcl_output_sgv.mlir +++ b/mlir/test/Target/DXSA/dcl_output_sgv.mlir @@ -1,5 +1,5 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_output_sgv.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_output_sgv , index = [0]>, // CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/dcl_output_siv.mlir b/mlir/test/Target/DXSA/dcl_output_siv.mlir index e48c2b5e4b25..50c47895180f 100644 --- a/mlir/test/Target/DXSA/dcl_output_siv.mlir +++ b/mlir/test/Target/DXSA/dcl_output_siv.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_output_siv.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_output_siv , index = [0]>, // CHECK-NEXT: dxsa.dcl_output_siv , index = [1]>, // CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/dcl_output_topology.mlir b/mlir/test/Target/DXSA/dcl_output_topology.mlir index 62ff373d0974..538f81ead0cd 100644 --- a/mlir/test/Target/DXSA/dcl_output_topology.mlir +++ b/mlir/test/Target/DXSA/dcl_output_topology.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_output_topology.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_output_topology pointlist // CHECK-NEXT: dxsa.dcl_output_topology linestrip // CHECK-NEXT: dxsa.dcl_output_topology trianglestrip diff --git a/mlir/test/Target/DXSA/dcl_temps.mlir b/mlir/test/Target/DXSA/dcl_temps.mlir index 686700499217..a9acc1140bd9 100644 --- a/mlir/test/Target/DXSA/dcl_temps.mlir +++ b/mlir/test/Target/DXSA/dcl_temps.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_temps.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_temps 1 // CHECK-NEXT: dxsa.dcl_temps 16 // CHECK-NEXT: dxsa.dcl_temps 4096 diff --git a/mlir/test/Target/DXSA/dcl_tessellator_domain.mlir b/mlir/test/Target/DXSA/dcl_tessellator_domain.mlir index 6007a9f8cc72..4c9d988afe19 100644 --- a/mlir/test/Target/DXSA/dcl_tessellator_domain.mlir +++ b/mlir/test/Target/DXSA/dcl_tessellator_domain.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_tessellator_domain.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_tessellator_domain domain_isoline // CHECK-NEXT: dxsa.dcl_tessellator_domain domain_tri // CHECK-NEXT: dxsa.dcl_tessellator_domain domain_quad diff --git a/mlir/test/Target/DXSA/dcl_tessellator_output_primitive.mlir b/mlir/test/Target/DXSA/dcl_tessellator_output_primitive.mlir index 23f0e2e78e0c..15d9b4feb456 100644 --- a/mlir/test/Target/DXSA/dcl_tessellator_output_primitive.mlir +++ b/mlir/test/Target/DXSA/dcl_tessellator_output_primitive.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_tessellator_output_primitive.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_tessellator_output_primitive output_point // CHECK-NEXT: dxsa.dcl_tessellator_output_primitive output_line // CHECK-NEXT: dxsa.dcl_tessellator_output_primitive output_triangle_cw diff --git a/mlir/test/Target/DXSA/dcl_tessellator_partitioning.mlir b/mlir/test/Target/DXSA/dcl_tessellator_partitioning.mlir index f51030ec2869..810b5dac196d 100644 --- a/mlir/test/Target/DXSA/dcl_tessellator_partitioning.mlir +++ b/mlir/test/Target/DXSA/dcl_tessellator_partitioning.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_tessellator_partitioning.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_tessellator_partitioning partitioning_integer // CHECK-NEXT: dxsa.dcl_tessellator_partitioning partitioning_pow2 // CHECK-NEXT: dxsa.dcl_tessellator_partitioning partitioning_fractional_odd diff --git a/mlir/test/Target/DXSA/dcl_tgsm_raw.mlir b/mlir/test/Target/DXSA/dcl_tgsm_raw.mlir index 231622666a1c..7673bac89b72 100644 --- a/mlir/test/Target/DXSA/dcl_tgsm_raw.mlir +++ b/mlir/test/Target/DXSA/dcl_tgsm_raw.mlir @@ -1,5 +1,5 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_tgsm_raw.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_tgsm_raw , 40 // CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/dcl_tgsm_structured.mlir b/mlir/test/Target/DXSA/dcl_tgsm_structured.mlir index 3119f1ec7948..cbc2d80de21a 100644 --- a/mlir/test/Target/DXSA/dcl_tgsm_structured.mlir +++ b/mlir/test/Target/DXSA/dcl_tgsm_structured.mlir @@ -1,5 +1,5 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/dcl_tgsm_structured.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.dcl_tgsm_structured , 16, 64 // CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/empty.mlir b/mlir/test/Target/DXSA/empty.mlir index 19d6cd55d36e..2702863cd0d6 100644 --- a/mlir/test/Target/DXSA/empty.mlir +++ b/mlir/test/Target/DXSA/empty.mlir @@ -1,4 +1,4 @@ // RUN: mlir-translate --import-dxsa-bin | FileCheck %s -// CHECK: module { -// CHECK-NEXT } +// CHECK: dxsa.module { +// CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/inputs/program_header_absent.bin b/mlir/test/Target/DXSA/inputs/program_header_absent.bin new file mode 100644 index 0000000000000000000000000000000000000000..a70ed0e01e9aecc0cce4ff6b9f2ee394056ede3c GIT binary patch literal 12 Tcmc~`U|?clU|_IgU|<9Q1{VOX literal 0 HcmV?d00001 diff --git a/mlir/test/Target/DXSA/inputs/program_header_present.bin b/mlir/test/Target/DXSA/inputs/program_header_present.bin new file mode 100644 index 0000000000000000000000000000000000000000..e4cccfc0236d789869c48a1018d1334f30f986a4 GIT binary patch literal 20 YcmWGwU|?Vc;tU1`CKe!V$H2e{01U4I0{{R3 literal 0 HcmV?d00001 diff --git a/mlir/test/Target/DXSA/mov-index.mlir b/mlir/test/Target/DXSA/mov-index.mlir index ff3ac88aa078..13036c9f615d 100644 --- a/mlir/test/Target/DXSA/mov-index.mlir +++ b/mlir/test/Target/DXSA/mov-index.mlir @@ -1,7 +1,7 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/mov-index.bin | FileCheck %s // mov o0.xyzw, v[r0.x][0].xyzw -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: %0 = dxsa.index.imm {imm = 0 : i32} // CHECK-NEXT: %1 = dxsa.operand %0 {mask = 240 : i32, num_components = 4 : i32, type = 2 : i32} // CHECK-NEXT: %2 = dxsa.index.imm {imm = 0 : i32} diff --git a/mlir/test/Target/DXSA/mov.mlir b/mlir/test/Target/DXSA/mov.mlir index 4e4bd9989df0..47cfcc4dfa15 100644 --- a/mlir/test/Target/DXSA/mov.mlir +++ b/mlir/test/Target/DXSA/mov.mlir @@ -1,7 +1,7 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/mov.bin | FileCheck %s // mov r0.x, l(3.000000) -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: %0 = dxsa.index.imm {imm = 0 : i32} // CHECK-NEXT: %1 = dxsa.operand %0 {mask = 16 : i32, num_components = 4 : i32, type = 0 : i32} // CHECK-NEXT: %2 = dxsa.operand.imm {imm = dense<1077936128> : vector<1xi32>} diff --git a/mlir/test/Target/DXSA/program_header_absent.mlir b/mlir/test/Target/DXSA/program_header_absent.mlir new file mode 100644 index 000000000000..cc220d7e168d --- /dev/null +++ b/mlir/test/Target/DXSA/program_header_absent.mlir @@ -0,0 +1,6 @@ +// RUN: mlir-translate --import-dxsa-bin %S/inputs/program_header_absent.bin | FileCheck %s + +// CHECK: dxsa.module { +// CHECK-NEXT: dxsa.dcl_temps 4 +// CHECK-NEXT: dxsa.instruction "ret" +// CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/program_header_present.mlir b/mlir/test/Target/DXSA/program_header_present.mlir new file mode 100644 index 000000000000..ce98dcc67b4b --- /dev/null +++ b/mlir/test/Target/DXSA/program_header_present.mlir @@ -0,0 +1,6 @@ +// RUN: mlir-translate --import-dxsa-bin %S/inputs/program_header_present.bin | FileCheck %s + +// CHECK: dxsa.module pixel_shader 5 0 { +// CHECK-NEXT: dxsa.dcl_temps 4 +// CHECK-NEXT: dxsa.instruction "ret" +// CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/ret.mlir b/mlir/test/Target/DXSA/ret.mlir index 24e0ec711ffc..3ab5e5b82806 100644 --- a/mlir/test/Target/DXSA/ret.mlir +++ b/mlir/test/Target/DXSA/ret.mlir @@ -1,6 +1,6 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/ret.bin | FileCheck %s -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: dxsa.instruction "ret" // CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/udiv.mlir b/mlir/test/Target/DXSA/udiv.mlir index cbd535c7df8e..d0683fc1d0f2 100644 --- a/mlir/test/Target/DXSA/udiv.mlir +++ b/mlir/test/Target/DXSA/udiv.mlir @@ -1,7 +1,7 @@ // RUN: mlir-translate --import-dxsa-bin %S/inputs/udiv.bin | FileCheck %s // udiv NULL, r0.x, vOutputControlPointID, 4 -// CHECK: module { +// CHECK: dxsa.module { // CHECK-NEXT: %0 = dxsa.operand {num_components = 0 : i32, type = 13 : i32} // CHECK-NEXT: %1 = dxsa.index.imm {imm = 0 : i32} // CHECK-NEXT: %2 = dxsa.operand %1 {mask = 16 : i32, num_components = 4 : i32, type = 0 : i32} From 7162bffe40625bcd97b87d131339b34b31e2e35c Mon Sep 17 00:00:00 2001 From: Vladimir Shiryaev Date: Sun, 31 May 2026 17:46:27 -0700 Subject: [PATCH 2/2] [mlir][dxsa] Add unknown instruction Example: dxsa.dcl_temps 1 dxsa.unknown dxsa.dcl_temps 2 Signed-off-by: Vladimir Shiryaev --- mlir/include/mlir/Dialect/DXSA/IR/DXSAOps.td | 22 +++ mlir/lib/Dialect/DXSA/IR/DXSA.cpp | 51 ++++++ mlir/lib/Target/DXSA/BinaryParser.cpp | 150 +++++++++++++++++- mlir/test/Target/DXSA/inputs/unknown.bin | Bin 0 -> 64 bytes .../Target/DXSA/inputs/unknown_past_eof.bin | Bin 0 -> 20 bytes mlir/test/Target/DXSA/unknown.mlir | 10 ++ mlir/test/Target/DXSA/unknown_invalid.mlir | 4 + mlir/test/Target/DXSA/unknown_past_eof.mlir | 6 + 8 files changed, 239 insertions(+), 4 deletions(-) create mode 100644 mlir/test/Target/DXSA/inputs/unknown.bin create mode 100644 mlir/test/Target/DXSA/inputs/unknown_past_eof.bin create mode 100644 mlir/test/Target/DXSA/unknown.mlir create mode 100644 mlir/test/Target/DXSA/unknown_invalid.mlir create mode 100644 mlir/test/Target/DXSA/unknown_past_eof.mlir diff --git a/mlir/include/mlir/Dialect/DXSA/IR/DXSAOps.td b/mlir/include/mlir/Dialect/DXSA/IR/DXSAOps.td index 857db9954226..f9c11c8dc590 100644 --- a/mlir/include/mlir/Dialect/DXSA/IR/DXSAOps.td +++ b/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 + ``` + }]; + + let arguments = (ins DenseI32ArrayAttr:$tokens); + let results = (outs); + let hasVerifier = 1; + let assemblyFormat = "custom($tokens) attr-dict"; +} + def DXSA_InlineOperandType_Temp : I32EnumAttrCase<"temp", 0>; def DXSA_InlineOperandType_Input : I32EnumAttrCase<"input", 1>; def DXSA_InlineOperandType_Output : I32EnumAttrCase<"output", 2>; diff --git a/mlir/lib/Dialect/DXSA/IR/DXSA.cpp b/mlir/lib/Dialect/DXSA/IR/DXSA.cpp index 365ee1dab82a..fb2afc61bd37 100644 --- a/mlir/lib/Dialect/DXSA/IR/DXSA.cpp +++ b/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 `` for the unknown op. +static ParseResult parseHexTokens(OpAsmParser &parser, + DenseI32ArrayAttr &attr) { + SmallVector tokens; + auto parseOneToken = [&]() -> ParseResult { + uint32_t value; + if (parser.parseInteger(value)) + return failure(); + tokens.push_back(static_cast(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 << "(t), + /*Width=*/10, /*Upper=*/true); + }); + printer << "]>"; +} + //===----------------------------------------------------------------------===// // TableGen'd attribute method definitions //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Target/DXSA/BinaryParser.cpp b/mlir/lib/Target/DXSA/BinaryParser.cpp index 20c81da4711e..875f4ace0c9f 100644 --- a/mlir/lib/Target/DXSA/BinaryParser.cpp +++ b/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 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 sliceInstruction() { + auto beginOffset = currentTokenOffset; + auto remainingBytes = buffer.size() - beginOffset; + if (remainingBytes < tokenSize) + return std::nullopt; + auto token0 = support::endian::read(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( + buffer.begin() + beginOffset + tokenSize, endianness::little); + if (token1 < 2) + return InstructionSlice{beginOffset, 2, "customdata length < 2"}; + auto instructionSize = static_cast(token1) * tokenSize; + if (instructionSize > remainingBytes) + return InstructionSlice{ + beginOffset, static_cast(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(length) * tokenSize; + if (instructionSize > remainingBytes) + return InstructionSlice{beginOffset, + static_cast(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 snapshotTokens(size_t beginOffset, uint32_t numTokens) { + SmallVector out(numTokens); + for (uint32_t i = 0; i < numTokens; ++i) + out[i] = static_cast(support::endian::read( + 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 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 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 inst = parseInstruction(); - if (failed(inst)) { + if (failed(parseInstructionOrUnknown())) return failure(); - } } return module; } diff --git a/mlir/test/Target/DXSA/inputs/unknown.bin b/mlir/test/Target/DXSA/inputs/unknown.bin new file mode 100644 index 0000000000000000000000000000000000000000..367d481a02341895f901aa46ceeb649b10e04712 GIT binary patch literal 64 zcmc~`U|?coU|{&q&cOVB-`aZ>VJ1QuKrtqum?;pm0I?ttC;p FrvPnS3r_$5 literal 0 HcmV?d00001 diff --git a/mlir/test/Target/DXSA/inputs/unknown_past_eof.bin b/mlir/test/Target/DXSA/inputs/unknown_past_eof.bin new file mode 100644 index 0000000000000000000000000000000000000000..468a58fc674f906158366789414173782aa67466 GIT binary patch literal 20 Ycmc~`U|?coU|{&q&cM0~2zCPj04-(*3jhEB literal 0 HcmV?d00001 diff --git a/mlir/test/Target/DXSA/unknown.mlir b/mlir/test/Target/DXSA/unknown.mlir new file mode 100644 index 000000000000..80728a72d2fe --- /dev/null +++ b/mlir/test/Target/DXSA/unknown.mlir @@ -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 +// CHECK-NEXT: dxsa.dcl_temps 2 +// CHECK-NEXT: dxsa.unknown +// CHECK-NEXT: dxsa.dcl_temps 3 +// CHECK-NEXT: dxsa.unknown +// CHECK-NEXT: } diff --git a/mlir/test/Target/DXSA/unknown_invalid.mlir b/mlir/test/Target/DXSA/unknown_invalid.mlir new file mode 100644 index 000000000000..d8ad07aae906 --- /dev/null +++ b/mlir/test/Target/DXSA/unknown_invalid.mlir @@ -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 diff --git a/mlir/test/Target/DXSA/unknown_past_eof.mlir b/mlir/test/Target/DXSA/unknown_past_eof.mlir new file mode 100644 index 000000000000..5df161337fba --- /dev/null +++ b/mlir/test/Target/DXSA/unknown_past_eof.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 +// CHECK-NEXT: }