[Mlir-commits] [mlir] b20719d - [mlir][Linalg] Add pooling_nchw_sum op.
Prashant Kumar
llvmlistbot at llvm.org
Fri Apr 8 05:29:13 PDT 2022
Author: Vivek Khandelwal
Date: 2022-04-08T17:57:47+05:30
New Revision: b20719dc7d39492b8934a6b6943f836fa60ce2f4
URL: https://github.com/llvm/llvm-project/commit/b20719dc7d39492b8934a6b6943f836fa60ce2f4
DIFF: https://github.com/llvm/llvm-project/commit/b20719dc7d39492b8934a6b6943f836fa60ce2f4.diff
LOG: [mlir][Linalg] Add pooling_nchw_sum op.
This commit adds pooling_nchw_sum as a yaml op.
Reviewed By: cathyzhyi, gysit
Differential Revision: https://reviews.llvm.org/D123013
Added:
Modified:
mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
mlir/test/Dialect/Linalg/named-ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
index 5e2c6eab356cf..3cc8f8e32cc94 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
@@ -2184,6 +2184,10 @@ metadata: !LinalgOpMetadata
doc: |-
Performs sum pooling.
+ Layout:
+ * Input: NHWC.
+ * Kernel: HW.
+
Numeric casting is performed on the input operand, promoting it to the same
data type as the accumulator/output.
implements:
@@ -2257,6 +2261,89 @@ structured_op: !LinalgStructuredOpConfig
- !ScalarExpression
scalar_arg: I
--- !LinalgOpConfig
+metadata: !LinalgOpMetadata
+ name: pooling_nchw_sum
+ cpp_class_name: PoolingNchwSumOp
+ doc: |-
+ Performs sum pooling.
+
+ Layout:
+ * Input: NCHW.
+ * Kernel: HW.
+
+ Numeric casting is performed on the input operand, promoting it to the same
+ data type as the accumulator/output.
+ implements:
+ - LinalgConvolutionOpInterface
+structured_op: !LinalgStructuredOpConfig
+ args:
+ - !LinalgOperandDefConfig
+ name: I
+ kind: input_tensor
+ type_var: T1
+ shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1, s2
+ * s3 + s4 * s5, s6 * s7 + s8 * s9)>
+ - !LinalgOperandDefConfig
+ name: K
+ kind: input_tensor
+ type_var: T2
+ shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4, s8)>
+ - !LinalgOperandDefConfig
+ name: O
+ kind: output_tensor
+ type_var: U
+ shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1, s2,
+ s6)>
+ - !LinalgOperandDefConfig
+ name: strides
+ kind: index_attr
+ index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s3,
+ s7)>
+ default_indices:
+ - 1
+ - 1
+ - !LinalgOperandDefConfig
+ name: dilations
+ kind: index_attr
+ index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s5,
+ s9)>
+ default_indices:
+ - 1
+ - 1
+ indexing_maps: !LinalgIndexingMapsConfig
+ static_indexing_maps:
+ - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
+ -> (d0, d1, d2 * s3 + d4 * s5, d3 * s7 + d5 * s9)>
+ - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
+ -> (d4, d5)>
+ - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
+ -> (d0, d1, d2, d3)>
+ iterator_types:
+ - parallel
+ - parallel
+ - parallel
+ - parallel
+ - reduction
+ - reduction
+ assignments:
+ - !ScalarAssign
+ arg: O
+ value: !ScalarExpression
+ scalar_fn:
+ kind: binary
+ fn_name: add
+ operands:
+ - !ScalarExpression
+ scalar_arg: O
+ - !ScalarExpression
+ scalar_fn:
+ kind: type
+ fn_name: cast_signed
+ type_var: U
+ operands:
+ - !ScalarExpression
+ scalar_arg: I
+--- !LinalgOpConfig
metadata: !LinalgOpMetadata
name: pooling_nhwc_max
cpp_class_name: PoolingNhwcMaxOp
diff --git a/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py b/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
index 2c6291bad85f4..023a95dc71267 100644
--- a/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
+++ b/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
@@ -522,6 +522,10 @@ def pooling_nhwc_sum(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH,
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
"""Performs sum pooling.
+ Layout:
+ * Input: NHWC.
+ * Kernel: HW.
+
Numeric casting is performed on the input operand, promoting it to the same
data type as the accumulator/output.
"""
@@ -531,6 +535,28 @@ def pooling_nhwc_sum(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH,
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])
+ at linalg_structured_op
+def pooling_nchw_sum(I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH,
+ S.OW * S.SW + S.KW * S.DW),
+ K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
+ O=TensorDef(U, S.N, S.C, S.OH, S.OW, output=True),
+ strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
+ dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
+ """Performs sum pooling.
+
+ Layout:
+ * Input: NCHW.
+ * Kernel: HW.
+
+ Numeric casting is performed on the input operand, promoting it to the same
+ data type as the accumulator/output.
+ """
+ implements(ConvolutionOpInterface)
+ domain(D.n, D.c, D.oh, D.ow, D.kh, D.kw)
+ O[D.n, D.c, D.oh, D.ow] += TypeFn.cast_signed(
+ U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW])
+
+
@linalg_structured_op
def pooling_nhwc_max(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH,
S.OW * S.SW + S.KW * S.DW, S.C),
diff --git a/mlir/test/Dialect/Linalg/named-ops.mlir b/mlir/test/Dialect/Linalg/named-ops.mlir
index 39a13683bfa92..50ad331a3c9af 100644
--- a/mlir/test/Dialect/Linalg/named-ops.mlir
+++ b/mlir/test/Dialect/Linalg/named-ops.mlir
@@ -260,6 +260,40 @@ func @pooling_nhwc_sum(%input: memref<1x4x4x1xf32>, %fake: memref<3x3xf32>, %out
// -----
+// CHECK-LABEL: func @pooling_nchw_sum_tensor
+// CHECK: %{{.+}} = linalg.pooling_nchw_sum
+// CHECK-SAME: dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME: strides = dense<1> : tensor<2xi64>
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x1x4x4xf32>, tensor<3x3xf32>)
+// CHECK-SAME: outs(%{{.+}} : tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32>
+func @pooling_nchw_sum_tensor(%input: tensor<1x1x4x4xf32>) -> tensor<1x1x2x2xf32> {
+ %fake = linalg.init_tensor [3, 3] : tensor<3x3xf32>
+ %init = linalg.init_tensor [1, 1, 2, 2] : tensor<1x1x2x2xf32>
+ %cst = arith.constant 0.000000e+00 : f32
+ %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32>
+ %res = linalg.pooling_nchw_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+ ins(%input, %fake: tensor<1x1x4x4xf32>, tensor<3x3xf32>)
+ outs(%fill: tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32>
+ return %res : tensor<1x1x2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @pooling_nchw_sum
+// CHECK: linalg.pooling_nchw_sum
+// CHECK-SAME: dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME: strides = dense<1> : tensor<2xi64>
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x1x4x4xf32>, memref<3x3xf32>)
+// CHECK-SAME: outs(%{{.+}} : memref<1x1x2x2xf32>)
+func @pooling_nchw_sum(%input: memref<1x1x4x4xf32>, %fake: memref<3x3xf32>, %output: memref<1x1x2x2xf32>) {
+ linalg.pooling_nchw_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+ ins(%input, %fake: memref<1x1x4x4xf32>, memref<3x3xf32>)
+ outs(%output: memref<1x1x2x2xf32>)
+ return
+}
+
+// -----
+
// CHECK-LABEL: func @pooling_nhwc_max_tensor
// CHECK: %{{.+}} = linalg.pooling_nhwc_max
// CHECK-SAME: dilations = dense<1> : tensor<2xi64>
More information about the Mlir-commits
mailing list