[mlir][spirv] Add support for TOSA Extended Instruction Set (001000.1) #168519
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This patch adds support for the TOSA Extended Instruction Set (001000.1) to the SPIR-V dialect in MLIR. The TOSA extended instruction set provides a standardized set of machine learning operations designed to be used 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 introduces: * Dialect plumbing for import, serialization, and deserialization of the TOSA extended instruction set. * One `spirv.Tosa.*` operation per TOSA extended instruction, each lowering to the corresponding `OpExtInst`. * Verification enforcing that all `spirv.Tosa.*` ops appear only within `spirv.ARM.Graph` regions, operate on `!spirv.arm.tensor<...>` types, and are well-formed according to the TOSA 001000.1 specification. All TOSA 001000.1 extended instructions are covered. Parser, printer, verifier, and round-trip tests using MLIR’s SPIR-V serialization/deserialization infrastructure are included. This work completes support for expressing TOSA extended instructions inside SPIR-V graphs in MLIR, aligning with Khronos SPIR-V TOSA specifications. Specification: https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html Change-Id: I11642bc7a7982e86026d9d6a8af3a070fee2217f Signed-off-by: Davide Grohmann <[email protected]>
Member
|
@llvm/pr-subscribers-mlir-spirv @llvm/pr-subscribers-mlir Author: Davide Grohmann (davidegrohmann) ChangesThis patch adds support for the TOSA Extended Instruction Set (001000.1) to the SPIR-V dialect in MLIR. The TOSA extended instruction set provides a standardized set of machine learning operations designed to be used within The change introduces:
All TOSA 001000.1 extended instructions are covered. Parser, printer, verifier, and round-trip tests using MLIR’s SPIR-V serialization/deserialization infrastructure are included. This work completes support for expressing TOSA extended instructions inside SPIR-V graphs in MLIR, aligning with Khronos SPIR-V TOSA specifications. Specification: Patch is 357.82 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/168519.diff 9 Files Affected:
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
index b628f1a3f7b20..d9a132c1fc792 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
@@ -4233,6 +4233,7 @@ def SPIRV_IsTensorArmType : CPred<"::llvm::isa<::mlir::spirv::TensorArmType>($_s
def SPIRV_Void : TypeAlias<NoneType, "void">;
def SPIRV_Bool : TypeAlias<I1, "bool">;
def SPIRV_Integer : AnyIntOfWidths<[8, 16, 32, 64]>;
+def SPIRV_Int8 : TypeAlias<I8, "Int8">;
def SPIRV_Int16 : TypeAlias<I16, "Int16">;
def SPIRV_Int32 : TypeAlias<I32, "Int32">;
def SPIRV_Float32 : TypeAlias<F32, "Float32">;
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.td
index 96ef035eda37a..3ef9699154cd1 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.td
@@ -45,6 +45,7 @@ include "mlir/Dialect/SPIRV/IR/SPIRVNonUniformOps.td"
include "mlir/Dialect/SPIRV/IR/SPIRVPrimitiveOps.td"
include "mlir/Dialect/SPIRV/IR/SPIRVCLOps.td"
include "mlir/Dialect/SPIRV/IR/SPIRVStructureOps.td"
+include "mlir/Dialect/SPIRV/IR/SPIRVTosaOps.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
#endif // MLIR_DIALECT_SPIRV_IR_OPS
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTosaOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTosaOps.td
new file mode 100644
index 0000000000000..551c0827dcd5d
--- /dev/null
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTosaOps.td
@@ -0,0 +1,1806 @@
+//===- SPIRVTosaOps.td - TOSA extended insts spec file -----*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the op definition spec of TOSA extension ops.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPIRV_IR_TOSA_OPS
+#define MLIR_DIALECT_SPIRV_IR_TOSA_OPS
+
+include "mlir/Dialect/SPIRV/IR/SPIRVBase.td"
+include "mlir/Dialect/SPIRV/IR/SPIRVGraphOps.td"
+include "mlir/Dialect/SPIRV/IR/SPIRVTosaTypes.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+//===----------------------------------------------------------------------===//
+// SPIR-V TOSA opcode specification.
+//===----------------------------------------------------------------------===//
+
+// Base class for all TOSA ops.
+class SPIRV_TosaOp<string mnemonic, int opcode, list<Trait> traits = []> :
+ SPIRV_ExtInstOp<mnemonic, "Tosa", "TOSA.001000.1", opcode, !listconcat(traits, [InGraphScope])> {
+
+ let availability = [
+ MinVersion<SPIRV_V_1_5>,
+ MaxVersion<SPIRV_V_1_6>,
+ Extension<[SPV_ARM_graph, SPV_ARM_tensors]>,
+ Capability<[SPIRV_C_GraphARM]>
+ ];
+}
+
+
+def SPIRV_TosaArgMaxOp : SPIRV_TosaOp<"ArgMax", 0, [Pure]> {
+ let summary = "ArgMax - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_argmax
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Int32: $axis,
+ SPIRV_Int32: $nan_mode,
+ SPIRV_TosaNumerical_TensorArm1DTo6D: $input
+ );
+
+
+ let results = (outs
+ SPIRV_TosaInteger_TensorArmUpTo5D: $output
+ );
+
+ let hasVerifier = 1;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaAvgPool2DOp : SPIRV_TosaOp<"AvgPool2D", 1, [Pure]> {
+ let summary = "AvgPool2D - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_avg_pool2d
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Int32_1DTensorArmOfLength2: $kernel,
+ SPIRV_Int32_1DTensorArmOfLength2: $stride,
+ SPIRV_Int32_1DTensorArmOfLength4: $pad,
+ SPIRV_Int32: $acc_type,
+ SPIRV_TosaNumerical_TensorArm4D: $input,
+ SPIRV_TosaNumerical_1DTensorArmOfLength1: $input_zp,
+ SPIRV_TosaNumerical_1DTensorArmOfLength1: $output_zp
+ );
+
+
+ let results = (outs
+ SPIRV_TosaNumerical_TensorArm4D: $output
+ );
+
+ let hasVerifier = 1;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaConv2DOp : SPIRV_TosaOp<"Conv2D", 2, [Pure]> {
+ let summary = "Conv2D - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_conv2d
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Int32_1DTensorArmOfLength4: $pad,
+ SPIRV_Int32_1DTensorArmOfLength2: $stride,
+ SPIRV_Int32_1DTensorArmOfLength2: $dilation,
+ SPIRV_Int32: $acc_type,
+ SPIRV_Bool: $local_bound,
+ SPIRV_TosaNumerical_TensorArm4D: $input,
+ SPIRV_TosaNumerical_TensorArm4D: $weight,
+ SPIRV_TosaNumerical_TensorArm1D: $bias,
+ SPIRV_TosaNumerical_1DTensorArmOfLength1: $input_zp,
+ SPIRV_TosaNumerical_1DTensorArmOfLength1: $weight_zp
+ );
+
+
+ let results = (outs
+ SPIRV_TosaNumerical_TensorArm4D: $output
+ );
+
+ let hasVerifier = 1;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaConv3DOp : SPIRV_TosaOp<"Conv3D", 3, [Pure]> {
+ let summary = "Conv3D - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_conv3d
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Int32_1DTensorArmOfLength6: $pad,
+ SPIRV_Int32_1DTensorArmOfLength3: $stride,
+ SPIRV_Int32_1DTensorArmOfLength3: $dilation,
+ SPIRV_Int32: $acc_type,
+ SPIRV_Bool: $local_bound,
+ SPIRV_TosaNumerical_TensorArm5D: $input,
+ SPIRV_TosaNumerical_TensorArm5D: $weight,
+ SPIRV_TosaNumerical_TensorArm1D: $bias,
+ SPIRV_TosaNumerical_1DTensorArmOfLength1: $input_zp,
+ SPIRV_TosaNumerical_1DTensorArmOfLength1: $weight_zp
+ );
+
+
+ let results = (outs
+ SPIRV_TosaNumerical_TensorArm5D: $output
+ );
+
+ let hasVerifier = 1;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaDepthwiseConv2DOp : SPIRV_TosaOp<"DepthwiseConv2D", 4, [Pure]> {
+ let summary = "DepthwiseConv2D - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_depthwise_conv2d
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Int32_1DTensorArmOfLength4: $pad,
+ SPIRV_Int32_1DTensorArmOfLength2: $stride,
+ SPIRV_Int32_1DTensorArmOfLength2: $dilation,
+ SPIRV_Int32: $acc_type,
+ SPIRV_Bool: $local_bound,
+ SPIRV_TosaNumerical_TensorArm4D: $input,
+ SPIRV_TosaNumerical_TensorArm4D: $weight,
+ SPIRV_TosaNumerical_TensorArm1D: $bias,
+ SPIRV_TosaNumerical_1DTensorArmOfLength1: $input_zp,
+ SPIRV_TosaNumerical_1DTensorArmOfLength1: $weight_zp
+ );
+
+
+ let results = (outs
+ SPIRV_TosaNumerical_TensorArm4D: $output
+ );
+
+ let hasVerifier = 1;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaFFT2DOp : SPIRV_TosaOp<"FFT2D", 5, [Pure]> {
+ let summary = "FFT2D - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_fft2d
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Bool: $inverse,
+ SPIRV_Bool: $local_bound,
+ SPIRV_TosaFloat_TensorArm3D: $input_real,
+ SPIRV_TosaFloat_TensorArm3D: $input_imag
+ );
+
+
+ let results = (outs
+ SPIRV_Struct_2_TosaFloat_TensorArm3D: $output
+ );
+
+ let hasVerifier = 1;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaMatMulOp : SPIRV_TosaOp<"MatMul", 6, [Pure]> {
+ let summary = "MatMul - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_matmul
+ }];
+
+
+ let arguments = (ins
+ SPIRV_TosaNumerical_TensorArm3D: $A,
+ SPIRV_TosaNumerical_TensorArm3D: $B,
+ SPIRV_TosaNumerical_1DTensorArmOfLength1: $A_zp,
+ SPIRV_TosaNumerical_1DTensorArmOfLength1: $B_zp
+ );
+
+
+ let results = (outs
+ SPIRV_TosaNumerical_TensorArm3D: $output
+ );
+
+ let hasVerifier = 1;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaMaxPool2DOp : SPIRV_TosaOp<"MaxPool2D", 7, [Pure]> {
+ let summary = "MaxPool2D - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_max_pool2d
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Int32_1DTensorArmOfLength2: $kernel,
+ SPIRV_Int32_1DTensorArmOfLength2: $stride,
+ SPIRV_Int32_1DTensorArmOfLength4: $pad,
+ SPIRV_Int32: $nan_mode,
+ SPIRV_TosaNumerical_TensorArm4D: $input
+ );
+
+
+ let results = (outs
+ SPIRV_TosaNumerical_TensorArm4D: $output
+ );
+
+ let hasVerifier = 1;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaRFFT2DOp : SPIRV_TosaOp<"RFFT2D", 8, [Pure]> {
+ let summary = "RFFT2D - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_rfft2d
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Bool: $local_bound,
+ SPIRV_TosaFloat_TensorArm3D: $input_real
+ );
+
+
+ let results = (outs
+ SPIRV_Struct_2_TosaFloat_TensorArm3D: $output
+ );
+
+ let hasVerifier = 1;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaTransposeConv2DOp : SPIRV_TosaOp<"TransposeConv2D", 9, [Pure]> {
+ let summary = "TransposeConv2D - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_transpose_conv2d
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Int32_1DTensorArmOfLength4: $out_pad,
+ SPIRV_Int32_1DTensorArmOfLength2: $stride,
+ SPIRV_Int32: $acc_type,
+ SPIRV_Bool: $local_bound,
+ SPIRV_TosaNumerical_TensorArm4D: $input,
+ SPIRV_TosaNumerical_TensorArm4D: $weight,
+ SPIRV_TosaNumerical_TensorArm1D: $bias,
+ SPIRV_TosaNumerical_1DTensorArmOfLength1: $input_zp,
+ SPIRV_TosaNumerical_1DTensorArmOfLength1: $weight_zp
+ );
+
+
+ let results = (outs
+ SPIRV_TosaNumerical_TensorArm4D: $output
+ );
+
+ let hasVerifier = 1;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaClampOp : SPIRV_TosaOp<"Clamp", 10, [Pure]> {
+ let summary = "Clamp - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_clamp
+ }];
+
+
+ let arguments = (ins
+ SPIRV_TosaNumerical: $min_val,
+ SPIRV_TosaNumerical: $max_val,
+ SPIRV_Int32: $nan_mode,
+ SPIRV_TosaNumerical_TensorArmUpTo6D: $input
+ );
+
+
+ let results = (outs
+ SPIRV_TosaNumerical_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 1;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaErfOp : SPIRV_TosaOp<"Erf", 11, [Pure]> {
+ let summary = "Erf - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_erf
+ }];
+
+
+ let arguments = (ins
+ SPIRV_TosaFloat_TensorArmUpTo6D: $input
+ );
+
+
+ let results = (outs
+ SPIRV_TosaFloat_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaSigmoidOp : SPIRV_TosaOp<"Sigmoid", 12, [Pure]> {
+ let summary = "Sigmoid - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_sigmoid
+ }];
+
+
+ let arguments = (ins
+ SPIRV_TosaFloat_TensorArmUpTo6D: $input
+ );
+
+
+ let results = (outs
+ SPIRV_TosaFloat_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaTanhOp : SPIRV_TosaOp<"Tanh", 13, [Pure]> {
+ let summary = "Tanh - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_tanh
+ }];
+
+
+ let arguments = (ins
+ SPIRV_TosaFloat_TensorArmUpTo6D: $input
+ );
+
+
+ let results = (outs
+ SPIRV_TosaFloat_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaAddOp : SPIRV_TosaOp<"Add", 14, [Pure]> {
+ let summary = "Add - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_add
+ }];
+
+
+ let arguments = (ins
+ SPIRV_TosaNumerical_TensorArmUpTo6D: $input1,
+ SPIRV_TosaNumerical_TensorArmUpTo6D: $input2
+ );
+
+
+ let results = (outs
+ SPIRV_TosaNumerical_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaArithmeticRightShiftOp : SPIRV_TosaOp<"ArithmeticRightShift", 15, [Pure]> {
+ let summary = "ArithmeticRightShift - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_arithmetic_right_shift
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Bool: $round,
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input1,
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input2
+ );
+
+
+ let results = (outs
+ SPIRV_TosaInteger_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaBitwiseAndOp : SPIRV_TosaOp<"BitwiseAnd", 16, [Pure]> {
+ let summary = "BitwiseAnd - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_bitwise_and
+ }];
+
+
+ let arguments = (ins
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input1,
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input2
+ );
+
+
+ let results = (outs
+ SPIRV_TosaInteger_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaBitwiseOrOp : SPIRV_TosaOp<"BitwiseOr", 17, [Pure]> {
+ let summary = "BitwiseOr - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_bitwise_or
+ }];
+
+
+ let arguments = (ins
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input1,
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input2
+ );
+
+
+ let results = (outs
+ SPIRV_TosaInteger_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaBitwiseXorOp : SPIRV_TosaOp<"BitwiseXor", 18, [Pure]> {
+ let summary = "BitwiseXor - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_bitwise_xor
+ }];
+
+
+ let arguments = (ins
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input1,
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input2
+ );
+
+
+ let results = (outs
+ SPIRV_TosaInteger_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaIntDivOp : SPIRV_TosaOp<"IntDiv", 19, [Pure]> {
+ let summary = "IntDiv - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_intdiv
+ }];
+
+
+ let arguments = (ins
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input1,
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input2
+ );
+
+
+ let results = (outs
+ SPIRV_TosaInteger_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaLogicalAndOp : SPIRV_TosaOp<"LogicalAnd", 20, [Pure]> {
+ let summary = "LogicalAnd - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_logical_and
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Bool_TensorArmUpTo6D: $input1,
+ SPIRV_Bool_TensorArmUpTo6D: $input2
+ );
+
+
+ let results = (outs
+ SPIRV_Bool_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaLogicalLeftShiftOp : SPIRV_TosaOp<"LogicalLeftShift", 21, [Pure]> {
+ let summary = "LogicalLeftShift - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_logical_left_shift
+ }];
+
+
+ let arguments = (ins
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input1,
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input2
+ );
+
+
+ let results = (outs
+ SPIRV_TosaInteger_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaLogicalRightShiftOp : SPIRV_TosaOp<"LogicalRightShift", 22, [Pure]> {
+ let summary = "LogicalRightShift - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_logical_right_shift
+ }];
+
+
+ let arguments = (ins
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input1,
+ SPIRV_TosaInteger_TensorArmUpTo6D: $input2
+ );
+
+
+ let results = (outs
+ SPIRV_TosaInteger_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaLogicalOrOp : SPIRV_TosaOp<"LogicalOr", 23, [Pure]> {
+ let summary = "LogicalOr - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_logical_or
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Bool_TensorArmUpTo6D: $input1,
+ SPIRV_Bool_TensorArmUpTo6D: $input2
+ );
+
+
+ let results = (outs
+ SPIRV_Bool_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaLogicalXorOp : SPIRV_TosaOp<"LogicalXor", 24, [Pure]> {
+ let summary = "LogicalXor - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_logical_xor
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Bool_TensorArmUpTo6D: $input1,
+ SPIRV_Bool_TensorArmUpTo6D: $input2
+ );
+
+
+ let results = (outs
+ SPIRV_Bool_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{
+ operands attr-dict `:` `(` type(operands) `)` `->` type(results)
+ }];
+}
+
+
+def SPIRV_TosaMaximumOp : SPIRV_TosaOp<"Maximum", 25, [Pure]> {
+ let summary = "Maximum - TOSA extended instruction set 001000.1";
+
+ let description = [{
+ https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_maximum
+ }];
+
+
+ let arguments = (ins
+ SPIRV_Int32: $nan_mode,
+ SPIRV_TosaNumerical_TensorArmUpTo6D: $input1,
+ SPIRV_TosaNumerical_TensorArmUpTo6D: $input2
+ );
+
+
+ let results = (outs
+ SPIRV_TosaNumerical_TensorArmUpTo6D: $output
+ );
+
+ let hasVerifier = 0;
+
+ let assemblyFormat = [{...
[truncated]
|
🐧 Linux x64 Test Results
|
Contributor
|
I've just skimmed through the changes, and I'll try to review it in more details in the coming days, but my first thought: is it possible to push some more verification into ODS. Immediately I see a lot of checks with the pattern: if (xxx.hasRank() && yyy.getRank() != N)Maybe a mix of existing and custom predicates could deal with that nicely? The image operations may give you some inspiration for that: https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVImageOps.td |
kuhar
requested changes
Nov 18, 2025
Member
kuhar
left a comment
kuhar
left a comment
There was a problem hiding this comment.
This is cool! Could you please break it down into a dozen or two smaller PRs so that we can review?
Contributor
Author
|
@kuhar I can split it as the following:
|
Member
|
Code and tests should be landed together, so that each incremental step is self-contained and tested. For example, you can start with the initial scaffolding and then have a series of PRs with ops related to each other. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This patch adds support for the TOSA Extended Instruction Set (001000.1) to the SPIR-V dialect in MLIR. The TOSA extended instruction set provides a standardized set of machine learning operations designed to be used within
spirv.ARM.Graphoperations (corresponding to OpGraphARM in SPV_ARM_graph) and typed with!spirv.arm.tensor<...>(corresponding to OpTypeTensorARM in SPV_ARM_tensor).The change introduces:
spirv.Tosa.*operation per TOSA extended instruction, each lowering to the correspondingOpExtInst.spirv.Tosa.*ops appear only withinspirv.ARM.Graphregions, operate on!spirv.arm.tensor<...>types, and are well-formed according to the TOSA 001000.1 specification.All TOSA 001000.1 extended instructions are covered. Parser, printer, verifier, and round-trip tests using MLIR’s SPIR-V serialization/deserialization infrastructure are included.
This work completes support for expressing TOSA extended instructions inside SPIR-V graphs in MLIR, aligning with Khronos SPIR-V TOSA specifications.
Specification:
https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html