diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc @@ -91,3 +91,36 @@ O(n, f, d, h, w) = std_addf(std_mulf( I(n, c, d + kd, h + kh, w + kw), K(f, c, kd, kh, kw))); } + +ods_def: +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(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(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 --- a/mlir/test/Dialect/Linalg/generalize-named-ops.mlir +++ b/mlir/test/Dialect/Linalg/generalize-named-ops.mlir @@ -73,3 +73,29 @@ // 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 --- /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 +}