[mlir][linalg] Define a depthwise 2-D convolution op

[antiagainst]

Summary

This commit defines linalg.depthwise_conv_2d_nhwc for depthwise
2-D convolution with NHWC input/output data format.

This op right now only support channel multiplier == 1, which is
the most common case.

Depends on D96297

[mehdi_amini]

mlir/test/Dialect/Linalg/named-ops.mlir

26

Can you add failing tests here? (missing strides attribute, or invalid type/shape for the strides attributes for example)

I'd also look for a linalg-generalize-named-ops test to check the semantics.

[nicolasvasilache]

+1 on the negative tests but IIRC there is no size check/verifier automatically generated from the affine maps?
This is generally something we want for all linalg ops but it is missing atm.

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

77

The commit message mentions channel multiplier must be 1, this description should also be loud about it if people expect this parameter, as it does not appear here.

[antiagainst]

Address comments

[antiagainst]

mlir/test/Dialect/Linalg/named-ops.mlir

26

Good point. Actually this reveals a hidden issue. I've captured more in D96297 and proposed a solution there. Let me know if you think there are better ways to handle it.

[nicolasvasilache]

looks good post D96297, thanks!

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

O(n, d, h, w, f) = std_addf<kd, kh, kw>(std_mulf(
     I(n, d + kd, h + kh, w + kw, c), K(f, kd, kh, kw, c)));
 }
 
 ods_def<ConvNCDHWOp>:
 def conv_3d_ncdhw(I: f32(N, C, D, H, W), K: f32(F, C, KD, KH, KW)) -> (O: f32(N, F, D, H, W)) {
   O(n, f, d, h, w) = std_addf<kd, kh, kw>(std_mulf(
     I(n, c, d + kd, h + kh, w + kw), K(f, c, kd, kh, kw)));
 }
+
+ods_def<DepthwiseConvNHWCOp>:
+def depthwise_conv_2d_nhwc
+      (I: f32(N, IH, IW, C), K: f32(KH, KW, C))
+   -> (O: f32(N, OH, OW, C))
+  attr(strides: 2xi64)
+"""A depth-wise 2-D convolution operation.
+
+This operation performs depth-wise 2-D convolution over an input `I` and filter
+`F` and generates output `O` using the following computation:
+
+```
+O(n, oh, ow, c) = std_addf<kh, kw>(std_mulf(
+  I(n, oh * strides[0] + kh, ow * strides[1] + kw, c), K(kh, kw, c)))
+```
+
+where
+
+* `I` is a 4-D tensor with shape `(N, IH, IW, C)`.
+* `F` is a 3-D tensor with shape `(KH, KW, C)`.
+* `O` is a 4-D tensor with shape `(N, OH, OW, C)`.
+* `strides` is a 2-element vector attribute for window strides along the
+  height/width dimension.
+
+The indexing maps for these three tensors contain 6 dimensions, following the
+order of (`N`, `OH`, `OW`, `C`, `KH`, `KW`).
+"""
+{
+  O(n, oh, ow, c) = std_addf<kh, kw>(std_mulf(
+    I(n, oh * strides[0] + kh, ow * strides[1] + kw, c), K(kh, kw, c)));
+}

mlir/test/Dialect/Linalg/named-ops.mlir

+// RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
+
+// CHECK-LABEL: func @depthwise_conv_2d_nhwc_tensor
+func @depthwise_conv_2d_nhwc_tensor(%input: tensor<1x113x113x96xf32>, %filter: tensor<3x3x96xf32>) -> tensor<1x56x56x96xf32> {
+  %init = linalg.init_tensor [1, 56, 56, 96] : tensor<1x56x56x96xf32>
+  // CHECK:      %{{.+}} = linalg.depthwise_conv_2d_nhwc
+  // CHECK-SAME:   {strides = dense<2> : vector<2xi64>}
+  // CHECK-SAME:   ins(%{{.+}}, %{{.+}} : tensor<1x113x113x96xf32>, tensor<3x3x96xf32>)
+  // CHECK-SAME:   outs(%{{.+}} : tensor<1x56x56x96xf32>) -> tensor<1x56x56x96xf32>
+  %0 = linalg.depthwise_conv_2d_nhwc {strides = dense<2> : vector<2xi64>}
+         ins(%input, %filter: tensor<1x113x113x96xf32>, tensor<3x3x96xf32>)
+         outs(%init: tensor<1x56x56x96xf32>) -> tensor<1x56x56x96xf32>
+  return %0: tensor<1x56x56x96xf32>
+}
+
+// CHECK-LABEL: func @depthwise_conv_2d_nhwc_memref
+func @depthwise_conv_2d_nhwc_memref(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) {
+  // CHECK:      linalg.depthwise_conv_2d_nhwc
+  // CHECK-SAME:   {strides = dense<2> : vector<2xi64>}
+  // CHECK-SAME:   ins(%{{.+}}, %{{.+}} : memref<1x113x113x96xf32>, memref<3x3x96xf32>)
+  // CHECK-SAME:   outs(%{{.+}} : memref<1x56x56x96xf32>)
+  linalg.depthwise_conv_2d_nhwc {strides = dense<2> : vector<2xi64>}
+    ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>)
+    outs(%output: memref<1x56x56x96xf32>)
+  return
+}