[Mlir-commits] [mlir] [mlir][spirv] Add reduction ops in TOSA Ext Inst Set (PR #187278)

[Davide Grohmann]

================
@@ -1902,4 +1915,211 @@ def SPIRV_TosaGreaterEqualOp : SPIRV_TosaComparisonOp<"GreaterEqual", 47> {
 }
 
 
+def SPIRV_TosaReduceAllOp : SPIRV_TosaReductionOp<"ReduceAll", 48> {
+  let summary = "Reduce All operator.";
+
+  let description = [{
+    Reduces a tensor along the given axis with a Logical AND operation.
+
+    References:
+      * https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_reduce_all
+      * https://www.mlplatform.org/tosa/tosa_spec_1_0_1.html#_reduce_all
+
+    #### Example:
+    ```mlir
+    %1 = spirv.Tosa.ReduceAll axis = 2, %arg0 : !spirv.arm.tensor<18x22x23x12xi1> -> !spirv.arm.tensor<18x22x1x12xi1>
+    ```
+  }];
+
+  let arguments = (ins
+    SPIRV_TensorArmAxisAttr: $axis,
+    SPIRV_Bool_TensorArm: $input
+  );
+
+  let results = (outs
+    SPIRV_Bool_TensorArm: $output
+  );
+
+  let assemblyFormat = [{
+    `axis` `=` $axis `,`
+    $input
+    attr-dict `:` type(operands) `->` type(results)
+  }];
+}
+
+
+def SPIRV_TosaReduceAnyOp : SPIRV_TosaReductionOp<"ReduceAny", 49> {
+  let summary = "Reduce Any operator.";
+
+  let description = [{
+    Reduces a tensor along the given axis with a Logical OR operation.
+
+    References:
+      * https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_reduce_any
+      * https://www.mlplatform.org/tosa/tosa_spec_1_0_1.html#_reduce_any
+
+    #### Example:
+    ```mlir
+    %1 = spirv.Tosa.ReduceAny axis = 2, %arg0 : !spirv.arm.tensor<25x13x30x8xi1> -> !spirv.arm.tensor<25x13x1x8xi1>
+    ```
+  }];
+
+  let arguments = (ins
+    SPIRV_TensorArmAxisAttr: $axis,
+    SPIRV_Bool_TensorArm: $input
+  );
+
+  let results = (outs
+    SPIRV_Bool_TensorArm: $output
+  );
+
+  let assemblyFormat = [{
+    `axis` `=` $axis `,`
+    $input
+    attr-dict `:` type(operands) `->` type(results)
+  }];
+}
+
+
+def SPIRV_TosaReduceMaxOp : SPIRV_TosaReductionOp<"ReduceMax", 50> {
+  let summary = "Reduce Max operator.";
+
+  let description = [{
+    Reduces a tensor along the given axis with a Maximum operation.
+    NaN Propagation Mode is ignored for inputs using integer element types.
+
+    References:
+      * https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_reduce_max
+      * https://www.mlplatform.org/tosa/tosa_spec_1_0_1.html#_reduce_max
+
+    #### Example:
+    ```mlir
+    %2 = spirv.Tosa.ReduceMax axis = 2, nan_mode = <Propagate>, %arg0 : !spirv.arm.tensor<8x30x12x3xi8> -> !spirv.arm.tensor<8x30x1x3xi8>
+    %2 = spirv.Tosa.ReduceMax axis = 2, nan_mode = <Propagate>, %arg0 : !spirv.arm.tensor<16x20x10xf16> -> !spirv.arm.tensor<16x20x1xf16>
+    ```
+  }];
+
+  let arguments = (ins
+    SPIRV_TensorArmAxisAttr: $axis,
+    SPIRV_TosaExtNaNPropagationModeAttr: $nan_mode,
+    SPIRV_TosaNumerical_TensorArm: $input
+  );
+
+  let results = (outs
+    SPIRV_TosaNumerical_TensorArm: $output
+  );
+
+  let assemblyFormat = [{
+    `axis` `=` $axis `,`
+    `nan_mode` `=` $nan_mode `,`
+    $input
+    attr-dict `:` type(operands) `->` type(results)
+  }];
+}
+
+
+def SPIRV_TosaReduceMinOp : SPIRV_TosaReductionOp<"ReduceMin", 51> {
+  let summary = "Reduce Min operator.";
+
+  let description = [{
+    Reduces a tensor along the given axis with a Minimum operation.
+    NaN Propagation Mode is ignored for inputs using integer element types.
+
+    References:
+      * https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_reduce_min
+      * https://www.mlplatform.org/tosa/tosa_spec_1_0_1.html#_reduce_min
+
+    #### Example:
+    ```mlir
+    %2 = spirv.Tosa.ReduceMin axis = 2, nan_mode = <Propagate>, %arg0 : !spirv.arm.tensor<2x5x5x1xi8> -> !spirv.arm.tensor<2x5x1x1xi8>
+    %2 = spirv.Tosa.ReduceMin axis = 2, nan_mode = <Propagate>, %arg0 : !spirv.arm.tensor<27x10x25x9xf16> -> !spirv.arm.tensor<27x10x1x9xf16>
+    ```
+  }];
+
+  let arguments = (ins
+    SPIRV_TensorArmAxisAttr: $axis,
+    SPIRV_TosaExtNaNPropagationModeAttr: $nan_mode,
+    SPIRV_TosaNumerical_TensorArm: $input
+  );
+
+  let results = (outs
+    SPIRV_TosaNumerical_TensorArm: $output
+  );
+
+  let assemblyFormat = [{
+    `axis` `=` $axis `,`
+    `nan_mode` `=` $nan_mode `,`
+    $input
+    attr-dict `:` type(operands) `->` type(results)
+  }];
+}
+
+
+def SPIRV_TosaReduceProductOp : SPIRV_TosaReductionOp<"ReduceProduct", 52> {
+  let summary = "Reduce Product operator.";
+
+  let description = [{
+    Reduces a tensor along the given axis by computing the Product of the axis.
+
+    References:
+      * https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_reduce_product
+      * https://www.mlplatform.org/tosa/tosa_spec_1_0_1.html#_reduce_product
+
+    #### Example:
+    ```mlir
+    %1 = spirv.Tosa.ReduceProduct axis = 2, %arg0 : !spirv.arm.tensor<2x16x25xf16> -> !spirv.arm.tensor<2x16x1xf16>
+    ```
+  }];
+
+  let arguments = (ins
+    SPIRV_TensorArmAxisAttr: $axis,
+    SPIRV_TosaFloat_TensorArm: $input
+  );
+
+  let results = (outs
+    SPIRV_TosaFloat_TensorArm: $output
+  );
+
+  let assemblyFormat = [{
+    `axis` `=` $axis `,`
+    $input
+    attr-dict `:` type(operands) `->` type(results)
+  }];
+}
+
+
+def SPIRV_TosaReduceSumOp : SPIRV_TosaReductionOp<"ReduceSum", 53> {
+  let summary = "Reduce Sum operator.";
+
+  let description = [{
+    Reduces a tensor along the given axis by computing the Sum of the axis.
----------------
davidegrohmann wrote:

Sorry my mistake.
Sum is actually defined for integer types. So yes it can overflow/underflows and cannot be pure.
The product is defined only on floats, so it can be pure.

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