[Mlir-commits] [mlir] 068bf9e - [mlir][linalg] Define a depthwise 2-D convolution op
Lei Zhang
llvmlistbot at llvm.org
Tue Feb 9 05:56:13 PST 2021
Author: Lei Zhang
Date: 2021-02-09T08:55:20-05:00
New Revision: 068bf9e80220329f7baeaf8df34fb4818ad8aa58
URL: https://github.com/llvm/llvm-project/commit/068bf9e80220329f7baeaf8df34fb4818ad8aa58
DIFF: https://github.com/llvm/llvm-project/commit/068bf9e80220329f7baeaf8df34fb4818ad8aa58.diff
LOG: [mlir][linalg] Define a depthwise 2-D convolution op
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.
Reviewed By: nicolasvasilache
Differential Revision: https://reviews.llvm.org/D94966
Added:
mlir/test/Dialect/Linalg/named-ops.mlir
Modified:
mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc
mlir/test/Dialect/Linalg/generalize-named-ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc
index 5406b3099293..6764e55e3610 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc
@@ -91,3 +91,36 @@ def conv_3d_ncdhw(I: f32(N, C, D, H, W), K: f32(F, C, KD, KH, KW)) -> (O: f32(N,
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 supports 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)));
+}
diff --git a/mlir/test/Dialect/Linalg/generalize-named-ops.mlir b/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
index c9f24844662f..44207308ab00 100644
--- a/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
+++ b/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
@@ -73,3 +73,29 @@ func @generalize_matmul_tensor(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>, %C:
// 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
diff --git a/mlir/test/Dialect/Linalg/named-ops.mlir b/mlir/test/Dialect/Linalg/named-ops.mlir
new file mode 100644
index 000000000000..f4a8fa94a7e0
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/named-ops.mlir
@@ -0,0 +1,56 @@
+// 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
+}
More information about the Mlir-commits
mailing list