[Mlir-commits] [mlir] [mlir][spirv] Expand support for TOSA Extended Instruction Set (00100… (PR #176908)

[Jakub Kuderski]

================
@@ -83,4 +83,254 @@ def SPIRV_TosaArgMaxOp : SPIRV_TosaOp<"ArgMax", 0, [Pure]> {
   }];
 }
 
+
+def SPIRV_TosaConv2DOp : SPIRV_TosaOp<"Conv2D", 2, [Pure]> {
+  let summary = "2D Convolution operator.";
+
+  let description = [{
+    Performs a 2D convolution over the given tensor input, using the weight
+    tensor. Implementations may choose to skip calculation of multiplies in
+    the padding area.
+
+    References:
+      * https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_conv2d
+      * https://www.mlplatform.org/tosa/tosa_spec_1_0_1.html#_conv2d
+
+    #### Example:
+    ```mlir
+    %7 = spirv.Tosa.Conv2D pad = dense<[1, 0, 0, 0]> : !spirv.arm.tensor<4xi32>, stride = dense<[1, 2]> : !spirv.arm.tensor<2xi32>, dilation = dense<[7, 1]> : !spirv.arm.tensor<2xi32>, acc_type = <INT32>, local_bound = false, %arg0, %arg1, %arg2, %5, %6 : !spirv.arm.tensor<1x65535x3x1xi8>, !spirv.arm.tensor<7x1x1x1xi8>, !spirv.arm.tensor<1xi32>, !spirv.arm.tensor<1xi8>, !spirv.arm.tensor<1xi8> -> !spirv.arm.tensor<1x65536x2x7xi32>
+    %7 = spirv.Tosa.Conv2D pad = dense<0> : !spirv.arm.tensor<4xi32>, stride = dense<1> : !spirv.arm.tensor<2xi32>, dilation = dense<1> : !spirv.arm.tensor<2xi32>, acc_type = <FP16>, local_bound = true, %arg0, %arg1, %arg2, %5, %6 : !spirv.arm.tensor<1x34x18x27xf16>, !spirv.arm.tensor<11x1x1x27xf16>, !spirv.arm.tensor<11xf16>, !spirv.arm.tensor<1xf16>, !spirv.arm.tensor<1xf16> -> !spirv.arm.tensor<1x34x18x11xf16>
+    ```
+  }];
+
+  let arguments = (ins
+    SPIRV_Int32_1DTensorArmOfLength4Attr: $pad,
+    SPIRV_Int32_1DTensorArmOfLength2Attr: $stride,
+    SPIRV_Int32_1DTensorArmOfLength2Attr: $dilation,
----------------
kuhar wrote:

If these must be Int32 attrs, can we skip the types in the assembly format? You can find examples in other dialects like `vector.insert_strided_slice`

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