Skip to content

Commit

Permalink
[Relay] Fix index order in conv2d computation for Arm CPU.
Browse files Browse the repository at this point in the history
When dilation is larger than value 1 in conv2d with NHWC
layout, the ordering of indexes when accessing data array
in computation of convolution appears to be incorrect.

'data_vec' is defined as

lambda n, oho, owo, kh, kw, ic, ohi, owi:

But accessed as

data_vec[n, oho, owo, kh, kw, ohi, owi, ic]

This patch fixes the order of indexes.
  • Loading branch information
Anastasia Stulova committed Jun 28, 2021
1 parent c586834 commit b5e3063
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion python/tvm/topi/arm_cpu/conv2d_spatial_pack.py
Original file line number Diff line number Diff line change
Expand Up @@ -344,7 +344,7 @@ def conv2d_spatial_pack_nhwc(cfg, data, kernel, strides, padding, dilation, out_
conv = te.compute(
ovshape,
lambda n, oho, owo, oco, ohi, owi, oci: te.sum(
data_vec[n, oho, owo, kh, kw, ohi, owi, ic].astype(out_dtype)
data_vec[n, oho, owo, kh, kw, ic, ohi, owi].astype(out_dtype)
* kernel_vec[oco, kh, kw, ic, oci].astype(out_dtype),
axis=[ic, kh, kw],
),
Expand Down

0 comments on commit b5e3063

Please sign in to comment.