[tvm] branch main updated: [ETHOSN] Improved handling of 5d reshapes (#10860)

[ma...@apache.org]

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

manupa 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 fc736ed6dc [ETHOSN] Improved handling of 5d reshapes (#10860)
fc736ed6dc is described below

commit fc736ed6dcc1108b411af16e6dfae68b2ef7f96a
Author: Leo-arm <Le...@arm.com>
AuthorDate: Wed Apr 6 14:27:10 2022 +0100

    [ETHOSN] Improved handling of 5d reshapes (#10860)
    
    Resolves an issue with 5d reshapes in the Yolo network and added a
    test case. Refactored the reshape tests to use parametrization.
---
 src/relay/backend/contrib/ethosn/ethosn_api.cc   |  5 ++
 tests/python/contrib/test_ethosn/test_reshape.py | 72 +++++++++++++-----------
 2 files changed, 45 insertions(+), 32 deletions(-)

diff --git a/src/relay/backend/contrib/ethosn/ethosn_api.cc b/src/relay/backend/contrib/ethosn/ethosn_api.cc
index 2ed94fbc13..14fda4651f 100644
--- a/src/relay/backend/contrib/ethosn/ethosn_api.cc
+++ b/src/relay/backend/contrib/ethosn/ethosn_api.cc
@@ -275,6 +275,11 @@ EthosnError EthosnAPI::Reshape(const Expr& expr, ReshapeParams* params) {
   const auto* input_dtype = reshape->args[0]->checked_type().as<TensorTypeNode>();
   const auto& reshape_attrs = reshape->attrs.as<ReshapeAttrs>();
 
+  if (reshape_attrs->newshape.size() > params->new_shape.size()) {
+    return EthosnError(ErrStrm() << "reshape dimension=" << reshape_attrs->newshape.size()
+                                 << ", reshape dimension must be <= " << params->new_shape.size());
+  }
+
   sl::TensorShape input_tensor_shape = {1, 1, 1, 1};
   sl::DataType input_data_type;
   EthosnError err = Tvm2Npu(input_dtype->shape, &input_tensor_shape);
diff --git a/tests/python/contrib/test_ethosn/test_reshape.py b/tests/python/contrib/test_ethosn/test_reshape.py
index 2b40b9af9f..6266367e90 100644
--- a/tests/python/contrib/test_ethosn/test_reshape.py
+++ b/tests/python/contrib/test_ethosn/test_reshape.py
@@ -36,53 +36,61 @@ def _get_model(input_shape, output_shape, dtype):
 
 @requires_ethosn
 @pytest.mark.parametrize("dtype", ["uint8", "int8"])
-def test_reshape(dtype):
-    trials = [
+@pytest.mark.parametrize(
+    "input_shape, output_shape",
+    [
         ((1, 15, 4, 1), (1, 60)),
         ((1, 15, 4, 1), (1, 30, 2)),
         ((1, 15, 4, 1), (1, 4, 15, 1)),
         ((1, 15, 4, 1), (1, 12, 5, 1)),
         ((1, 15, 4, 1), (1, -1, 2, 1)),
-    ]
-
+    ],
+)
+def test_reshape(dtype, input_shape, output_shape):
     np.random.seed(0)
-    for input_shape, output_shape in trials:
-        inputs = {
-            "a": tvm.nd.array(
-                np.random.randint(
-                    low=np.iinfo(dtype).min,
-                    high=np.iinfo(dtype).max + 1,
-                    size=input_shape,
-                    dtype=dtype,
-                )
+    inputs = {
+        "a": tvm.nd.array(
+            np.random.randint(
+                low=np.iinfo(dtype).min,
+                high=np.iinfo(dtype).max + 1,
+                size=input_shape,
+                dtype=dtype,
             )
-        }
-        outputs = []
-        for npu in [False, True]:
-            model, params = _get_model(input_shape, output_shape, dtype)
-            mod = tei.make_module(model, params)
-            outputs.append(tei.build_and_run(mod, inputs, 1, params, npu=npu))
+        )
+    }
+    outputs = []
+    for npu in [False, True]:
+        model, params = _get_model(input_shape, output_shape, dtype)
+        mod = tei.make_module(model, params)
+        outputs.append(tei.build_and_run(mod, inputs, 1, params, npu=npu))
 
-        tei.verify(outputs, dtype, 1)
+    tei.verify(outputs, dtype, 1)
 
 
 @requires_ethosn
-def test_reshape_failure():
-    trials = [
+@pytest.mark.parametrize(
+    "input_shape, output_shape, dtype, err_msg",
+    [
         (
             (1, 15, 4, 1),
             (1, 15, -2),
             "uint8",
             "reshape dimension=-2, reshape dimension must be >= -1",
         ),
-    ]
-
+        (
+            (1, 1, 4, 1),
+            (1, 1, 2, 2, 1),
+            "uint8",
+            "reshape dimension=5, reshape dimension must be <= 4",
+        ),
+    ],
+)
+def test_reshape_failure(input_shape, output_shape, dtype, err_msg):
     np.random.seed(0)
-    for input_shape, output_shape, dtype, err_msg in trials:
-        model, params = _get_model(input_shape, output_shape, dtype)
-        mod = tei.make_module(model, params)
-        pattern = get_pattern_table("ethos-n")
-        mod = tei.make_module(model, params)
-        mod = relay.transform.MergeComposite(pattern)(mod)
-        mod = tei.make_ethosn_partition(mod["main"].body)
-        tei.test_error(mod, {}, err_msg)
+    model, params = _get_model(input_shape, output_shape, dtype)
+    mod = tei.make_module(model, params)
+    pattern = get_pattern_table("ethos-n")
+    mod = tei.make_module(model, params)
+    mod = relay.transform.MergeComposite(pattern)(mod)
+    mod = tei.make_ethosn_partition(mod["main"].body)
+    tei.test_error(mod, {}, err_msg)