[tvm] branch main updated: [Relay][TOPI] Fix compute and schedule bugs for conv2d_winograd_nhwc on mali device. (#8091)

[jc...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

jcf94 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 0cf7ac3  [Relay][TOPI] Fix compute and schedule bugs for conv2d_winograd_nhwc on mali device. (#8091)
0cf7ac3 is described below

commit 0cf7ac3e436e5350c3db5f45857a633ef826f288
Author: hzshao <ha...@gmail.com>
AuthorDate: Tue May 25 16:45:47 2021 +0800

    [Relay][TOPI] Fix compute and schedule bugs for conv2d_winograd_nhwc on mali device. (#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 <ha...@gmail.com>
---
 python/tvm/relay/op/strategy/mali.py |  9 +++++++--
 python/tvm/topi/mali/conv2d.py       | 19 +++++++++++++++++--
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/python/tvm/relay/op/strategy/mali.py b/python/tvm/relay/op/strategy/mali.py
index fc47bd6..6c6440e 100644
--- a/python/tvm/relay/op/strategy/mali.py
+++ b/python/tvm/relay/op/strategy/mali.py
@@ -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,
@@ -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,
diff --git a/python/tvm/topi/mali/conv2d.py b/python/tvm/topi/mali/conv2d.py
index 6da3c5f..62be92c 100644
--- a/python/tvm/topi/mali/conv2d.py
+++ b/python/tvm/topi/mali/conv2d.py
@@ -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,
     )