Fix an issue with grouped conv2d op

Reviewed By: silvas

Differential Revision: https://reviews.llvm.org/D128880
This commit is contained in:
George Petterson 2022-07-11 15:37:03 -04:00 committed by Nirvedh
parent f0cd538985
commit 4dc8cf3a86
2 changed files with 19 additions and 19 deletions

View File

@ -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

View File

@ -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,