[tvm] branch main updated: [microNPU] enable striping for network tests. (#11883)

[lu...@apache.org]

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

lukhut 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 b733aa3ec8 [microNPU] enable striping for network tests. (#11883)
b733aa3ec8 is described below

commit b733aa3ec87a1c5a6e49350bd805e187db1aca70
Author: Manupa Karunaratne <ma...@arm.com>
AuthorDate: Tue Jun 28 16:47:40 2022 +0100

    [microNPU] enable striping for network tests. (#11883)
    
    This commit enables the striping for network tests.
    Currently it requires, storage_rewrite to be run if
    striping is enabled to produce correct results.
    
    Change-Id: I12b976bb77d339771f8b5a554817d192e7c99723
---
 .../relay/backend/contrib/ethosu/tir/compiler.py   |  8 +++-
 .../test_ethosu/cascader/test_memory_reduction.py  | 10 ++---
 tests/python/contrib/test_ethosu/test_networks.py  | 49 ++++++++++++++++++++++
 3 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/python/tvm/relay/backend/contrib/ethosu/tir/compiler.py b/python/tvm/relay/backend/contrib/ethosu/tir/compiler.py
index db216e43e2..0fd82378c3 100644
--- a/python/tvm/relay/backend/contrib/ethosu/tir/compiler.py
+++ b/python/tvm/relay/backend/contrib/ethosu/tir/compiler.py
@@ -91,9 +91,15 @@ def lower_ethosu(sch, args, const_dict, name="main"):
         mod, const_dict = ethosu_passes.EncodeConstants(const_dict)(mod)
         mod = ethosu_passes.HoistAllocates()(mod)
         mod = ethosu_passes.CopyComputeReordering()(mod)
+
+        # When striping is enabled and if storage_rewrite is not run
+        # the striping results in incorrect code generation. This needs
+        # further investigation. Until such a time that is fixed, disable_storage_rewrite
+        # user directive will be overridden if striping is enabled.
         disable_storage_rewrite = curr_cfg.get("tir.disable_storage_rewrite", False)
-        if not disable_storage_rewrite:
+        if not disable_storage_rewrite or util.is_striping_enabled():
             mod = tvm.tir.transform.StorageRewrite()(mod)
+
         mod = tvm.tir.transform.RemoveNoOp()(mod)
         mod = ethosu_passes.AnnotateAllocates()(mod)
         mod, const_dict = ethosu_passes.CreatePrimFuncWithoutConstants(const_dict)(mod)
diff --git a/tests/python/contrib/test_ethosu/cascader/test_memory_reduction.py b/tests/python/contrib/test_ethosu/cascader/test_memory_reduction.py
index 0bfc64fe04..5c3b745cb4 100644
--- a/tests/python/contrib/test_ethosu/cascader/test_memory_reduction.py
+++ b/tests/python/contrib/test_ethosu/cascader/test_memory_reduction.py
@@ -158,15 +158,13 @@ def test_double_conv2d(
     assert workspace_size_cascader_enabled_striping_enabled == expected_ws_size_with_striping
 
 
-# TODO(ekalda): Fix a bug in the block config selection that selects block config that is too large
-# for the smaller accelerators
 @pytest.mark.parametrize(
     "accel_type, expected_ws_size_without_striping, expected_ws_size_with_striping",
     [
-        ("ethos-u55-256", 180288, 15200),
-        ("ethos-u55-128", 180288, 15200),
-        ("ethos-u55-64", 180288, 14432),
-        ("ethos-u55-32", 180272, 14416),
+        ("ethos-u55-256", 180288, 15312),
+        ("ethos-u55-128", 180288, 15312),
+        ("ethos-u55-64", 180288, 14544),
+        ("ethos-u55-32", 180272, 14544),
     ],
 )
 def test_depthwise2d_conv2d_pooling(
diff --git a/tests/python/contrib/test_ethosu/test_networks.py b/tests/python/contrib/test_ethosu/test_networks.py
index c4081f911a..9b09132a9e 100644
--- a/tests/python/contrib/test_ethosu/test_networks.py
+++ b/tests/python/contrib/test_ethosu/test_networks.py
@@ -139,5 +139,54 @@ def test_networks_with_usmp_and_cascader_wo_striping(accel_type, model_url, work
     assert allocated_pool_info.allocated_size == workspace_size
 
 
+@pytest.mark.parametrize(
+    "accel_type, model_url, workspace_size",
+    [
+        ("ethos-u55-256", MOBILENET_V1_URL, 1010000),
+        ("ethos-u55-256", MOBILENET_V2_URL, 1180000),
+    ],
+)
+def test_networks_with_usmp_and_cascader_with_striping(accel_type, model_url, workspace_size):
+    np.random.seed(23)
+
+    pool_name = "my_memory_pool"
+    host_target = tvm.target.Target("c")
+    ethosu_target = tvm.target.Target("ethos-u")
+    workspace_pools = WorkspaceMemoryPools(
+        [
+            WorkspacePoolInfo(
+                pool_name,
+                [host_target, ethosu_target],
+                PoolInfoProperties(
+                    size_hint_bytes=workspace_size,
+                    read_bandwidth_bytes_per_cycle=16,
+                    write_bandwidth_bytes_per_cycle=16,
+                    target_burst_bytes={ethosu_target: 1},
+                ),
+            )
+        ]
+    )
+    tflite_model_buf = infra.get_tflite_model(model_url)
+    input_data, output_data = infra.generate_ref_data_tflite(tflite_model_buf)
+    mod, params = convert_to_relay(tflite_model_buf)
+    mod = partition_for_ethosu(mod, params)
+    test_runner = infra.create_test_runner(
+        accel_type,
+        enable_usmp=True,
+        enable_cascader=True,
+        enable_striping=True,
+        workspace_pools=workspace_pools,
+    )
+    compiled_models = infra.build_source(
+        mod, input_data, output_data, test_runner, workspace_pools=workspace_pools
+    )
+    infra.verify_source(compiled_models, test_runner)
+
+    allocated_pool_info = list(
+        dict(compiled_models[0].executor_factory.executor_codegen_metadata.pool_inputs).values()
+    )[0]
+    assert allocated_pool_info.allocated_size <= workspace_size
+
+
 if __name__ == "__main__":
     pytest.main([__file__])