[Mlir-commits] [mlir] [mlir][spirv] Add Gather/Scatter/Resize ops in TOSA Ext Inst Set (PR #188497)

[Igor Wodiany]

================
@@ -2462,4 +2462,187 @@ def SPIRV_TosaTransposeOp : SPIRV_TosaOpWithResult<"Transpose", 60, [Pure,
 }
 
 
+def SPIRV_TosaGatherOp : SPIRV_TosaOpWithResult<"Gather", 61, [NoMemoryEffect,
+  AllElementTypesMatch<["values", "output"]>]> {
+  let summary = "Gather operation.";
+
+  let description = [{
+    Generate a tensor for which each element in the output is a subtensor of the
+    values tensor based on the indices. Undefined behaviour may occur if the
+    specified indices are out of range.
+
+    References:
+      * https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_gather
+      * https://www.mlplatform.org/tosa/tosa_spec_1_0_1.html#_gather
+
+    #### Example:
+    ```mlir
+    %0 = spirv.Tosa.Gather %values, %indices : !spirv.arm.tensor<31x11x45xi32>, !spirv.arm.tensor<31x15xi32> -> !spirv.arm.tensor<31x15x45xi32>
+    %0 = spirv.Tosa.Gather %values, %indices : !spirv.arm.tensor<59x61x19xf32>, !spirv.arm.tensor<59x65xi32> -> !spirv.arm.tensor<59x65x19xf32>
+    ```
+  }];
+
+  let arguments = (ins
+    SPIRV_TosaNumerical_TensorArm3D: $values,
+    SPIRV_Int32_TensorArm2D: $indices
+  );
+
+  let results = (outs
+    SPIRV_TosaNumerical_TensorArm3D: $output
+  );
+
+  let assemblyFormat = [{
+    $values `,`
+    $indices
+    attr-dict `:` type(operands) `->` type(results)
+  }];
+
+  let extraClassDeclaration = extraBaseClassDeclaration#[{
+    ::mlir::spirv::TensorArmType getValuesType() {
+      return cast<::mlir::spirv::TensorArmType>(getValues().getType());
+    }
+    ::mlir::spirv::TensorArmType getIndicesType() {
+      return cast<::mlir::spirv::TensorArmType>(getIndices().getType());
+    }
+  }];
+}
+
+
+def SPIRV_TosaScatterOp : SPIRV_TosaOpWithResult<"Scatter", 62, [NoMemoryEffect,
+  AllElementTypesMatch<["values_in", "input", "values_out"]>]> {
+  let summary = "Scatter operation.";
+
+  let description = [{
+    The values_out tensor is set to the values_in tensor with data modified as
+    follows: data from the input tensor is inserted at the positions specified
+    by the indices tensor. In use cases that require multiple updates to the
+    same output position, these must be decomposed into multiple scatter
+    operations.  Undefined behaviour may occur if the specified indices are
+    out of range or duplicate indices are provided.
+
+    References:
+      * https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_scatter
+      * https://www.mlplatform.org/tosa/tosa_spec_1_0_1.html#_scatter
+
+    #### Example:
+    ```mlir
+    %0 = spirv.Tosa.Scatter %values_in, %indices, %arg0 : !spirv.arm.tensor<34x28x54xi32>, !spirv.arm.tensor<34x18xi32>, !spirv.arm.tensor<34x18x54xi32> -> !spirv.arm.tensor<34x28x54xi32>
----------------
IgWod wrote:

I guess similar questions to the one above. In case dims can be arbitrary we should probably add a less symmetric example to make is less misleading.

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