[Mlir-commits] [mlir] [mlir][spirv] Add Pooling, Fourier Transform, and MatMul operations t… (PR #177585)

[Jakub Kuderski]

================
@@ -84,6 +85,66 @@ def SPIRV_TosaArgMaxOp : SPIRV_TosaOp<"ArgMax", 0, [Pure]> {
 }
 
 
+def SPIRV_TosaAvgPool2DOp : SPIRV_TosaOp<"AvgPool2D", 1, [Pure,
+  AllElementTypesMatch<["input", "input_zp", "output", "output_zp"]>]> {
+  let summary = "Performs average pooling on the input.";
+
+  let description = [{
+    Performs an average pooling over the given input tensor. A sliding
+    window of size given by <kernel size> is passed over the input tensor, with
+    the mean value being placed in the output tensor. When calculating the
+    average, only the number of valid input tensor values, but not padding, are
+    used to calculate the divisor.
+
+    References:
+      * https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_avg_pool2d
+      * https://www.mlplatform.org/tosa/tosa_spec_1_0_1.html#_avg_pool2d
+
+    #### Example:
+    ```mlir
+    %6 = spirv.Tosa.AvgPool2D kernel = [3, 3], stride = [1, 2], pad = [0, 1, 0, 0], acc_type = <INT32>, %arg0, %4, %5 : !spirv.arm.tensor<1x3x65537x1xi8>, !spirv.arm.tensor<1xi8>, !spirv.arm.tensor<1xi8> -> !spirv.arm.tensor<1x2x32768x1xi8>
+    %6 = spirv.Tosa.AvgPool2D kernel = [2, 2], stride = [1, 1], pad = [1, 0, 0, 0], acc_type = <FP32>, %arg0, %4, %5 : !spirv.arm.tensor<1x2x65533x2xf32>, !spirv.arm.tensor<1xf32>, !spirv.arm.tensor<1xf32> -> !spirv.arm.tensor<1x2x65532x2xf32>
+    ```
+  }];
+
+  let arguments = (ins
+    SPIRV_Int32_1DTensorArmOfLength2Attr: $kernel,
+    SPIRV_Int32_1DTensorArmOfLength2Attr: $stride,
+    SPIRV_Int32_1DTensorArmOfLength4Attr: $pad,
+    SPIRV_TosaExtAccTypeAttr: $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 = [{
+    `kernel` `=` custom<SPIRV_I32_1DArmTensor>($kernel) `,`
----------------
kuhar wrote:

Can we put this custom assembly format in the attribute definition, so that we don't have to repeat it everywhere?

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