forked from OSchip/llvm-project
Fix an issue with grouped conv2d op
Reviewed By: silvas Differential Revision: https://reviews.llvm.org/D128880
This commit is contained in:
parent
f0cd538985
commit
4dc8cf3a86
|
@ -1648,7 +1648,7 @@ metadata: !LinalgOpMetadata
|
|||
name: conv_2d_ngchw_fgchw
|
||||
cpp_class_name: Conv2DNgchwFgchwOp
|
||||
doc: |-
|
||||
Performs 2-D convolution.
|
||||
Performs 2-D grouped convolution.
|
||||
|
||||
Layout:
|
||||
* Input: NGCHW.
|
||||
|
@ -1664,44 +1664,44 @@ structured_op: !LinalgStructuredOpConfig
|
|||
name: I
|
||||
kind: input_tensor
|
||||
type_var: T1
|
||||
shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] -> (s0,
|
||||
s1, s2 * s3 + s4 * s5, s6 * s7 + s8 * s9)>
|
||||
shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] ->
|
||||
(s0, s1, s2, s3 * s4 + s5 * s6, s7 * s8 + s9 * s10)>
|
||||
- !LinalgOperandDefConfig
|
||||
name: K
|
||||
kind: input_tensor
|
||||
type_var: T2
|
||||
shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] -> (s10,
|
||||
s1, s11, s4, s8)>
|
||||
shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] ->
|
||||
(s11, s1, s2, s5, s9)>
|
||||
- !LinalgOperandDefConfig
|
||||
name: O
|
||||
kind: output_tensor
|
||||
type_var: U
|
||||
shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] -> (s0,
|
||||
s1, s10, s2, s6)>
|
||||
shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] ->
|
||||
(s0, s11, s1, s3, s7)>
|
||||
- !LinalgOperandDefConfig
|
||||
name: strides
|
||||
kind: index_attr
|
||||
index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] ->
|
||||
(s3, s7)>
|
||||
index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11]
|
||||
-> (s4, s8)>
|
||||
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, s10, s11] ->
|
||||
(s5, s9)>
|
||||
index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11]
|
||||
-> (s6, s10)>
|
||||
default_indices:
|
||||
- 1
|
||||
- 1
|
||||
indexing_maps: !LinalgIndexingMapsConfig
|
||||
static_indexing_maps:
|
||||
- affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
|
||||
s9, s10, s11] -> (d0, d1, d5, d3 * s3 + d6 * s5, d4 * s7 + d7 * s9)>
|
||||
- affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
|
||||
s9, s10, s11] -> (d2, d1, d5, d6, d7)>
|
||||
- affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
|
||||
s9, s10, s11] -> (d0, d1, d2, d3, d4)>
|
||||
- affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7,
|
||||
s8, s9, s10, s11] -> (d0, d1, d5, d3 * s4 + d6 * s6, d4 * s8 + d7 * s10)>
|
||||
- affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7,
|
||||
s8, s9, s10, s11] -> (d1, d2, d5, d6, d7)>
|
||||
- affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7,
|
||||
s8, s9, s10, s11] -> (d0, d1, d2, d3, d4)>
|
||||
iterator_types:
|
||||
- parallel
|
||||
- parallel
|
||||
|
|
|
@ -370,7 +370,7 @@ def conv_2d_nchw_fchw(I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH,
|
|||
def conv_2d_ngchw_fgchw(I=TensorDef(T1, S.N, S.G, S.C, S.OH * S.SH + S.KH * S.DH,
|
||||
S.OW * S.SW + S.KW * S.DW),
|
||||
K=TensorDef(T2, S.FG, S.G, S.C, S.KH, S.KW),
|
||||
O=TensorDef(U, S.N, S.G, S.FG, S.OH, S.OW, output=True),
|
||||
O=TensorDef(U, S.N, S.FG, S.G, 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 2-D grouped convolution.
|
||||
|
@ -386,7 +386,7 @@ def conv_2d_ngchw_fgchw(I=TensorDef(T1, S.N, S.G, S.C, S.OH * S.SH + S.KH * S.DH
|
|||
domain(D.n, D.g, D.fg, D.oh, D.ow, D.c, D.kh, D.kw)
|
||||
O[D.n, D.g, D.fg, D.oh, D.ow] += TypeFn.cast_signed(
|
||||
U, I[D.n, D.g, D.c, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW +
|
||||
D.kw * S.DW]) * TypeFn.cast_signed(U, K[D.fg, D.g, D.c, D.kh, D.kw])
|
||||
D.kw * S.DW]) * TypeFn.cast_signed(U, K[D.g, D.fg, D.c, D.kh, D.kw])
|
||||
|
||||
@linalg_structured_op
|
||||
def conv_3d_ndhwc_dhwcf(I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD,
|
||||
|
|
Loading…
Reference in New Issue