Skip to content

Commit

Permalink
[Relay][TOPI] Fix compute and schedule bugs for conv2d_winograd_nhwc …
Browse files Browse the repository at this point in the history
…on mali device. (apache#8091)

1. add argument `auto_scheduler_rewritten_layout=""` in conv2d_winograd_nhwc_mali;
2. add `need_auto_scheduler_layout=True` for conv2d_strategy_mali and
conv2d_winograd_without_weight_transfrom_strategy_mali.

Signed-off-by: haizhu.shao <haizhu.shao@gmail.com>
  • Loading branch information
i4oolish authored and trevor-m committed Jun 17, 2021
1 parent 8ecf290 commit 128335a
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 4 deletions.
9 changes: 7 additions & 2 deletions python/tvm/relay/op/strategy/mali.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,9 @@ def conv2d_strategy_mali(attrs, inputs, out_type, target):
)
if is_winograd_applicable:
strategy.add_implementation(
wrap_compute_conv2d(topi.nn.conv2d_winograd_nhwc),
wrap_compute_conv2d(
topi.nn.conv2d_winograd_nhwc, need_auto_scheduler_layout=True
),
naive_schedule, # this implementation should never be picked by autotvm
name="conv2d_nhwc.winograd",
plevel=15,
Expand Down Expand Up @@ -155,7 +157,10 @@ def conv2d_winograd_without_weight_transfrom_strategy_mali(attrs, inputs, out_ty
"Winograd conv2d NHWC is not enabled for mali without auto_scheduler."
)
strategy.add_implementation(
wrap_compute_conv2d(topi.nn.conv2d_winograd_nhwc_without_weight_transform),
wrap_compute_conv2d(
topi.nn.conv2d_winograd_nhwc_without_weight_transform,
need_auto_scheduler_layout=True,
),
naive_schedule, # this implementation should never be picked by autotvm
name="conv2d_nhwc_winograd_without_weight_transform",
plevel=15,
Expand Down
19 changes: 17 additions & 2 deletions python/tvm/topi/mali/conv2d.py
Original file line number Diff line number Diff line change
Expand Up @@ -579,14 +579,29 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):

@conv2d_winograd_nhwc.register(["mali"])
def conv2d_winograd_nhwc_mali(
data, weight, strides, padding, dilation, out_dtype, pre_computed=False
data,
weight,
strides,
padding,
dilation,
out_dtype,
pre_computed=False,
auto_scheduler_rewritten_layout="",
):
"""Conv2D Winograd in NHWC layout.
This is a clean version to be used by the auto-scheduler for mali.
"""
tile_size = _pick_tile_size(data, weight, layout="NHWC")
return _conv2d_winograd_nhwc_impl(
data, weight, strides, padding, dilation, out_dtype, tile_size, pre_computed
data,
weight,
strides,
padding,
dilation,
out_dtype,
tile_size,
pre_computed,
auto_scheduler_rewritten_layout,
)


Expand Down

0 comments on commit 128335a

Please sign in to comment.