forked from OSchip/llvm-project
[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
This commit is contained in:
Vivek Khandelwal
Prashant Kumar
parent
f1cfa461f2
commit
b20719dc7d
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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])
|
||||
|
||||
|
||||
@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),
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
Loading…
Reference in New Issue