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 @@ -72,3 +72,30 @@ 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_nhwc(I: f32(N, IH, IW, IC), K: f32(KH, KW, IC, CM)) -> (O: f32(N, OH, OW, IC, CM)) + attr(strides: 2xi32) +"""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, ic, cm) = std_addf(std_mulf( + I(n, oh * strides[0] + kh, ow * strides[1] + kw, ic), K(kh, kw, ic, cm))) +``` + +where + +* `I` is a 4-D tensor with shape `(N, IH, IW, IC)` +* `F` is a 4-D tensor with shape `(KH, KW, IC, CM)` +* `O` is a 5-D tensor with shape `(N, OH, OW, IC, CM)` + +The indexing maps for these three tensors contain 7 dimensions, following the +order of (`N`, `OH`, `OW`, `IC`, `CM`, `KH`, `KW`). +""" +{ + O(n, oh, ow, ic, cm) = std_addf(std_mulf( + I(n, oh * strides[0] + kh, ow * strides[1] + kw, ic), K(kh, kw, ic, cm))); +}