[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

27

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

96

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

27

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_input_nhwc_filter_hwc
+      (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`).
+
+Note: this op only support channel multiplier == 1.
+"""
+{
+  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/generalize-named-ops.mlir

 // CHECK-SAME:  ins(%{{.+}}, %{{.+}} : tensor<16x8xf32>, tensor<8x32xf32>)
 // CHECK-SAME: outs(%{{.+}} : tensor<16x32xf32>)
 
 // CHECK:      ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32)
 // CHECK-NEXT:   %[[MUL:.+]] = mulf %[[A_ARG]], %[[B_ARG]] : f32
 // CHECK-NEXT:   %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32
 // CHECK-NEXT:   linalg.yield %[[ADD]] : f32
 // CHECK-NEXT: -> tensor<16x32xf32>
+
+// -----
+
+func @depthwise_conv_2d_input_nhwc_filter_hwc(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) {
+  linalg.depthwise_conv_2d_input_nhwc_filter_hwc {strides = dense<2> : vector<2xi64>}
+    ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>)
+    outs(%output: memref<1x56x56x96xf32>)
+  return
+}
+
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d4, d2 * 2 + d5, d3)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5, d3)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
+
+// CHECK: func @depthwise_conv_2d_input_nhwc_filter_hwc
+
+// CHECK: linalg.generic
+// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x113x113x96xf32>, memref<3x3x96xf32>)
+// CHECK-SAME: outs(%{{.+}} : memref<1x56x56x96xf32>)
+
+// CHECK:         ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
+// CHECK-NEXT:      %[[MUL:.+]] = mulf %[[BBARG0]], %[[BBARG1]] : f32
+// CHECK-NEXT:      %[[ADD:.+]] = addf %[[BBARG2]], %[[MUL]] : f32
+// CHECK-NEXT:      linalg.yield %[[ADD]] : f32

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

+// RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
+
+// CHECK-LABEL: func @depthwise_conv_2d_input_nhwc_filter_hwc_tensor
+func @depthwise_conv_2d_input_nhwc_filter_hwc_tensor(%input: tensor<1x113x113x96xf32>, %filter: tensor<3x3x96xf32>) -> tensor<1x56x56x96xf32> {
+  %init = linalg.init_tensor [1, 56, 56, 96] : tensor<1x56x56x96xf32>
+  // CHECK:      %{{.+}} = linalg.depthwise_conv_2d_input_nhwc_filter_hwc
+  // 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_input_nhwc_filter_hwc {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_input_nhwc_filter_hwc_memref
+func @depthwise_conv_2d_input_nhwc_filter_hwc_memref(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) {
+  // CHECK:      linalg.depthwise_conv_2d_input_nhwc_filter_hwc
+  // CHECK-SAME:   {strides = dense<2> : vector<2xi64>}
+  // CHECK-SAME:   ins(%{{.+}}, %{{.+}} : memref<1x113x113x96xf32>, memref<3x3x96xf32>)
+  // CHECK-SAME:   outs(%{{.+}} : memref<1x56x56x96xf32>)
+  linalg.depthwise_conv_2d_input_nhwc_filter_hwc {strides = dense<2> : vector<2xi64>}
+    ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>)
+    outs(%output: memref<1x56x56x96xf32>)
+  return
+}
+
+// -----
+
+func @depthwise_conv_2d_input_nhwc_filter_missing_stride(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) {
+  // expected-error @+1 {{missing indexing map required attribute 'strides'}}
+  linalg.depthwise_conv_2d_input_nhwc_filter_hwc
+    ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>)
+    outs(%output: memref<1x56x56x96xf32>)
+  return
+}
+
+// -----
+
+func @depthwise_conv_2d_input_nhwc_filter_wrong_stride_element_type(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) {
+  // expected-error @+1 {{incorrect element type for indexing map required attribute 'strides'}}
+  linalg.depthwise_conv_2d_input_nhwc_filter_hwc {strides = dense<2.0> : vector<2xf32>}
+    ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>)
+    outs(%output: memref<1x56x56x96xf32>)
+  return
+}
+
+// -----
+
+func @depthwise_conv_2d_input_nhwc_filter_wrong_stride_size(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) {
+  // expected-error @+1 {{incorrect shape for indexing map required attribute 'strides'}}
+  linalg.depthwise_conv_2d_input_nhwc_filter_hwc {strides = dense<2> : vector<3xi64> }
+    ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>)
+    outs(%output: memref<1x56x56x96xf32>)
+  return
+}