[Mlir-commits] [mlir] [mlir][spirv] Initial support for TOSA Extended Instruction Set (0010… (PR #174402)

[Jakub Kuderski]

================
@@ -0,0 +1,74 @@
+//===- 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 = "Perform argmax on the input.";
+
+  let description = [{
+    Returns the index with the largest value across the given axis of the
+    input tensor. If multiple locations have equal values, returns the first
+    match along the search axis.
+
+    References:
+    https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_argmax
+    https://www.mlplatform.org/tosa/tosa_spec_1_0_1.html#_argmax
+
+    #### Example:
+    ```mlir
+    %2 = spirv.Tosa.ArgMax %0, %1 {nan_mode = #spirv.tosa_ext_nan_propagation_mode_type<Propagate>} : i32, !spirv.arm.tensor<3x28x17x17xi8> -> !spirv.arm.tensor<3x28x17xi32>
----------------
kuhar wrote:

Can we use an abbreviated format that looks like this? This is similar to how we handle attributes in integer dot prod and cooperative matrix ops, for example

```suggestion
    %2 = spirv.Tosa.ArgMax %0, %1, <Propagate> : i32, !spirv.arm.tensor<3x28x17x17xi8> -> !spirv.arm.tensor<3x28x17xi32>
```

I suggest to add this attribute to the assembly format.

https://github.com/llvm/llvm-project/pull/174402