[RFC][MLIR][SPIRV] Integrate TOSA Extended Instruction Set (001000.1) into MLIR SPIR-V Dialect

MLIR
,
1

Abstract

This RFC proposes adding first-class support for the TOSA Extended Instruction Set 001000.1 to MLIR’s SPIR-V dialect for use within spirv.ARM.Graph operations (corresponding to OpGraphARM in SPV_ARM_graph) and typed with !spirv.arm.tensor<...> (corresponding to OpTypeTensorARM in SPV_ARM_tensor). The change enables semantically faithful modeling of TOSA extended instructions inside SPIR-V graph operations, completing the connection between the existing SPV_ARM_graph/SPV_ARM_tensor support in MLIR and the Khronos SPIR-V TOSA specification.

Motivation

  • MLIR already supports the SPV_ARM_graph execution model and the SPV_ARM_tensor type system in the SPIR-V dialect.
  • The TOSA extended instruction set (001000.1) is designed to express ML/TOSA semantics in SPIR-V, intended to be invoked inside spirv.ARM.Graph (corresponding to OpGraphARM) over !spirv.arm.tensor<...> (corresponding to OpTypeTensorARM) values.
  • Without modeling the TOSA extended instructions in MLIR, we cannot represent TOSA ops directly and semantically in SPIR-V graphs.

This RFC enables: structured spirv.Tosa.* ops in MLIR, schema-checked typing via !spirv.arm.tensor, and graph-scoped verification.

Prior Art / Background

  • SPV_ARM_graph and SPV_ARM_tensor are already available in MLIR’s SPIR-V dialect.
  • SPIR-V supports extended instruction sets via OpExtInstImport/OpExtInst.
  • The TOSA 001000.1 extended set enumerates ML ops (e.g., pooling, elementwise, shape ops) to be used in graph contexts.

Proposed Changes (High-Level)

  1. Register the TOSA Extended Instruction Set in the SPIR-V dialect.
    • Add a new enumerant to SPIRVExtendedInstructionSet corresponding to TOSA 001000.1.
    • Dialect plumbing for import/serialization/deserialization of the TOSA set.
  2. Add spirv.Tosa.* operations that model each TOSA extended instruction.
    • Each op lowers to OpExtInst targeting the TOSA set with the appropriate opcode.
    • Ops are graph-scoped and must appear only within spirv.ARM.Graph regions.
  3. Type Integration with SPV_ARM_tensor.
    • All spirv.Tosa.* ops operate on/result in !spirv.arm.tensor<...> types.
    • Verifiers enforce tensor shape/rank/element-type constraints aligned with the TOSA spec.

Detailed Design

Dialect and Enumeration Extensions

  • Extend SPIRVExtendedInstructionSet with an entry for TOSA 001000.1.
  • Provide a stable symbolic name, e.g., TOSA.001000.1.
  • Teach the serializer/deserializer to import/export the TOSA set via OpExtInstImport.

New spirv.Tosa.* Ops

  • One MLIR op per TOSA extended instruction.
  • Each op takes operands that map directly to the operands defined in the TOSA extended instruction set.
  • Values that correspond to TOSA “attributes” (e.g., strides, padding, kernel size) are passed as spirv.Constant results or other SSA values, not as MLIR op attributes.
  • Operand/result types are restricted to !spirv.arm.tensor<...> where applicable.

Verification Rules

  • Context: spirv.Tosa.* ops must reside inside a spirv.ARM.Graph region.
  • Typing: All tensor operands/results must be !spirv.arm.tensor<...>; shapes and element types obey TOSA constraints.
  • Attributes/Immediates: Immediate scalars (e.g., integers) must follow the TOSA specification and are verified at build/parse time.

Assembly/Disassembly

  • Print/parse spirv.Tosa.* with clear positional operands, mirroring existing SPIR-V dialect style.
  • Round-trip to/from SPIR-V binary: use MLIR’s SPIR-V target serialization/deserialization infrastructure.
  • Ensure conformance with the Khronos specification via spirv-val.

Testing Strategy

  • Dialect Tests: Parser/printer, verifier, and builder coverage for each spirv.Tosa.* op; negative tests for context/type violations.
  • Round-Trip Tests: MLIR ↔ SPIR-V binary round-trips using MLIR’s SPIR-V serialization/deserialization, validated with spirv-val.

Backward Compatibility & Migration

  • No impact on existing users of SPV_ARM_graph/SPV_ARM_tensor.
  • New functionality is additive; existing pipelines remain valid.

Future Work

  • Conversion Pass: Introduce a lowering pipeline from MLIR TOSA dialect ops to spirv.Tosa.* ops inside spirv.ARM.Graph. This RFC only introduces the ops and dialect support; conversion can be staged later.
  • Verifier Enhancements: Add richer shape/padding/dilation consistency checks.

Implementation Plan

  1. Add TOSA extended set enumerant to SPIRVExtendedInstructionSet and implement import/export plumbing.
  2. Introduce all spirv.Tosa.* ops covering the complete TOSA 001000.1 extended instruction set.
  3. Verification and tests: parser/printer, verifier coverage, round-trip serialization/deserialization.

References

2

Having a new extended instruction set alongside GL and CL makes sense to me.

How many new ops will this import in total? Do you plan to add an import script or do this by hand? We have these for mass-importing GL/CL extensions: llvm-project/mlir/utils/spirv at main · llvm/llvm-project · GitHub , but the code has bitrotten a fair bit.

3

There are 65 operators.

We have a script that generated the SPIR-V dialect portions for the TOSA extension starting from the TOSA specification file (an xml file).

Is the plan to create the dialect on the fly when compiling or just use the script as a one-of import?

4

The existing scripts are meant to be invoked manually when we import new extensions etc., and often require manual cleanup. Having .td files generated during the build process with these scripts is a non-goal.

5

If you develop a new script to import the TOSA EIS, it would be nice to also check it in under the utils/spirv directory.

6

Happy to share the script as well as part of the contribution.

7

Sounds good to me.

Will you at any point need to support, in MLIR, multiple versions of the SPIR-V TOSA extension at the same time? If so, it’s probably something to think about ahead of time, since I don’t think we currently do a great job at versioning — mostly because I do not think there was a need. Ops seems to be added against the version of spec that is available at the time of adding and are updated as needed.

nit: Since it’s spirv.GL. and spirv.CL. for other extended instructions, should it be spirv.TOSA.*? It seems to me all references to TOSA use all caps.

8

Likely there will be multiple versions of the extension at some point. But I think you will have to define different sets of operators for each version of the extension so there is no risk to have incompatibilities or mixing versions.

I have no strong opinions about spirv.Tosa.* vs spirv.TOSA.*.

Should we add the version as part of the prefix? Like spirv.TOSA.001000.1.*?

9

I’d keep it simple since SPIR-V does a good job of versioning on their end – extensions can modify existing ops which composes well with how we handle things on the MLIR side.

Also, no offence, but this versioning scheme it not something I’d appreciate in my op names :smiley:

10

Fair enough, :slight_smile:

11

I’m also in favour of keeping it simple. Just to give a bit more context to why I raised this issue. When enabling validation, I came across a weird validation failure in spirv.Select [0] (that I expected should be caught by MLIR verification) and it turned out that the op was correct, but to support arguments it had, a higher version of SPIR-V was needed (1.0 vs 1.4). Since op definition doesn’t really have a concept of version, this problem only surfaced with spirv-val.

I’m not saying it’s a big problem (especially now we have better way to do validation) or something should fix anything, but I think that’s something that’s worth keeping in mind. I guess it’s a broader question of how much we want to handle in MLIR and how much can be handled externally.

[0] [mlir][spirv] Fix entry point, logical ops and sampled image Target tests by IgWod-IMG · Pull Request #159376 · llvm/llvm-project · GitHub

12

Initial patch here: [mlir][spirv] Add support for TOSA Extended Instruction Set (001000.1) by davidegrohmann · Pull Request #168519 · llvm/llvm-project · GitHub