[microNPU] enable striping for network tests.

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

Author: manupak

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)

tests/python/contrib/test_ethosu/cascader/test_memory_reduction.py

@@ -158,15 +158,13 @@ def tf_graph(x):
     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(

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, 1005440),
+        ("ethos-u55-256", MOBILENET_V2_URL, 1162368),
+    ],
+)
+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=1200000,
+                    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__])