MLECO-2745: Fix for arena-cache-size

Arena-cache-size is different from default when using
Dedicated_Sram memory mode.

Change-Id: Ie112146218e1ec456e17babd4ed3e7c7bc2009a8

3 files changed, 95 insertions, 41 deletions

diff --git a/docs/quick_start.md b/docs/quick_start.md
index d892ab1..6522bdd 100644
--- a/docs/quick_start.md
+++ b/docs/quick_start.md
@@ -136,7 +136,8 @@ curl -L https://github.com/ARM-software/ML-zoo/raw/7dd3b16bb84007daf88be8648983c
     --optimise Performance --config scripts/vela/default_vela.ini \
     --memory-mode=Shared_Sram \
     --system-config=Ethos_U55_High_End_Embedded \
-    --output-dir=resources_downloaded/kws
+    --output-dir=resources_downloaded/kws \
+    --arena-cache-size=2097152
 mv resources_downloaded/kws/ds_cnn_clustered_int8_vela.tflite resources_downloaded/kws/ds_cnn_clustered_int8_vela_H128.tflite
 
 . resources_downloaded/env/bin/activate && vela resources_downloaded/kws/ds_cnn_clustered_int8.tflite \
@@ -152,7 +153,8 @@ mv resources_downloaded/kws/ds_cnn_clustered_int8_vela.tflite resources_download
     --optimise Performance --config scripts/vela/default_vela.ini \
     --memory-mode=Shared_Sram \
     --system-config=Ethos_U55_High_End_Embedded \
-    --output-dir=resources_downloaded/kws_asr
+    --output-dir=resources_downloaded/kws_asr \
+    --arena-cache-size=2097152
 mv resources_downloaded/kws_asr/wav2letter_int8_vela.tflite resources_downloaded/kws_asr/wav2letter_int8_vela_H128.tflite
 
 . resources_downloaded/env/bin/activate && vela resources_downloaded/kws_asr/wav2letter_int8.tflite \
@@ -168,7 +170,8 @@ mv resources_downloaded/kws_asr/wav2letter_int8_vela.tflite resources_downloaded
     --optimise Performance --config scripts/vela/default_vela.ini \
     --memory-mode=Shared_Sram \
     --system-config=Ethos_U55_High_End_Embedded \
-    --output-dir=resources_downloaded/kws_asr
+    --output-dir=resources_downloaded/kws_asr \
+    --arena-cache-size=2097152
 mv resources_downloaded/kws_asr/ds_cnn_clustered_int8_vela.tflite resources_downloaded/kws_asr/ds_cnn_clustered_int8_vela_H128.tflite
 
 . resources_downloaded/env/bin/activate && vela resources_downloaded/kws_asr/ds_cnn_clustered_int8.tflite \
@@ -184,7 +187,8 @@ mv resources_downloaded/kws_asr/ds_cnn_clustered_int8_vela.tflite resources_down
     --optimise Performance --config scripts/vela/default_vela.ini \
     --memory-mode=Shared_Sram \
     --system-config=Ethos_U55_High_End_Embedded \
-    --output-dir=resources_downloaded/inference_runner
+    --output-dir=resources_downloaded/inference_runner \
+    --arena-cache-size=2097152
 mv resources_downloaded/inference_runner/dnn_s_quantized_vela.tflite resources_downloaded/inference_runner/dnn_s_quantized_vela_H128.tflite
 
 . resources_downloaded/env/bin/activate && vela resources_downloaded/inference_runner/dnn_s_quantized.tflite \
@@ -200,7 +204,8 @@ mv resources_downloaded/inference_runner/dnn_s_quantized_vela.tflite resources_d
     --optimise Performance --config scripts/vela/default_vela.ini \
     --memory-mode=Shared_Sram \
     --system-config=Ethos_U55_High_End_Embedded \
-    --output-dir=resources_downloaded/img_class
+    --output-dir=resources_downloaded/img_class \
+    --arena-cache-size=2097152
 mv resources_downloaded/img_class/mobilenet_v2_1.0_224_INT8_vela.tflite resources_downloaded/img_class/mobilenet_v2_1.0_224_INT8_vela_H128.tflite
 
 . resources_downloaded/env/bin/activate && vela resources_downloaded/img_class/mobilenet_v2_1.0_224_INT8.tflite \
@@ -216,7 +221,8 @@ mv resources_downloaded/img_class/mobilenet_v2_1.0_224_INT8_vela.tflite resource
     --optimise Performance --config scripts/vela/default_vela.ini \
     --memory-mode=Shared_Sram \
     --system-config=Ethos_U55_High_End_Embedded \
-    --output-dir=resources_downloaded/asr
+    --output-dir=resources_downloaded/asr \
+    --arena-cache-size=2097152
 mv resources_downloaded/asr/wav2letter_int8_vela.tflite resources_downloaded/asr/wav2letter_int8_vela_H128.tflite
 
 . resources_downloaded/env/bin/activate && vela resources_downloaded/asr/wav2letter_int8.tflite \
@@ -232,7 +238,8 @@ mv resources_downloaded/asr/wav2letter_int8_vela.tflite resources_downloaded/asr
     --optimise Performance --config scripts/vela/default_vela.ini \
     --memory-mode=Shared_Sram \
     --system-config=Ethos_U55_High_End_Embedded \
-    --output-dir=resources_downloaded/ad
+    --output-dir=resources_downloaded/ad \
+    --arena-cache-size=2097152
 mv resources_downloaded/ad/ad_medium_int8_vela.tflite resources_downloaded/ad/ad_medium_int8_vela_H128.tflite
 
 . resources_downloaded/env/bin/activate && vela resources_downloaded/ad/ad_medium_int8.tflite \
@@ -248,7 +255,8 @@ mv resources_downloaded/ad/ad_medium_int8_vela.tflite resources_downloaded/ad/ad
     --optimise Performance --config scripts/vela/default_vela.ini \
     --memory-mode=Shared_Sram \
     --system-config=Ethos_U55_High_End_Embedded \
-    --output-dir=resources_downloaded/ad
+    --output-dir=resources_downloaded/ad \
+    --arena-cache-size=2097152
 mv resources_downloaded/vww/vww4_128_128_INT8_vela.tflite resources_downloaded/vww/vww4_128_128_INT8_vela_H128.tflite
 
 . resources_downloaded/env/bin/activate && vela resources_downloaded/vww/vww4_128_128_INT8.tflite \
@@ -264,7 +272,8 @@ mv resources_downloaded/vww/vww4_128_128_INT8_vela.tflite resources_downloaded/v
     --optimise Performance --config scripts/vela/default_vela.ini \
     --memory-mode=Shared_Sram \
     --system-config=Ethos_U55_High_End_Embedded \
-    --output-dir=resources_downloaded/ad
+    --output-dir=resources_downloaded/ad \
+    --arena-cache-size=2097152
 mv resources_downloaded/noise_reduction/rnnoise_INT8_vela.tflite resources_downloaded/noise_reduction/rnnoise_INT8_vela_H128.tflite
 
 . resources_downloaded/env/bin/activate && vela resources_downloaded/noise_reduction/rnnoise_INT8.tflite \
diff --git a/docs/sections/building.md b/docs/sections/building.md
index b289e52..6cbabd9 100644
--- a/docs/sections/building.md
+++ b/docs/sections/building.md
@@ -318,21 +318,32 @@ for the default 128 MACs configuration of the Arm® *Ethos™-U55* NPU and for t
 > **Note:** This script requires Python version 3.6 or higher. Please make sure all [build prerequisites](./building.md#build-prerequisites)
 > are satisfied.
 
-If you need to optimize the models for a different Ethos-U configuration, you can pass a
-list of additional configurations for Vela compiler. For example, to optimize the downloaded models
-for *Ethos™-U55* 32 MAC configuration and *Ethos™-U65* 512 MAC configuration, you can choose to run:
+Additional command line arguments supported by this script are:
 
-```sh
-python3 ./set_up_default_resources.py \
-  --additional-ethos-u-config-name ethos-u55-32 \
-  --additional-ethos-u-config-name ethos-u65-512
-```
-> **Note:** As the argument name suggests, the configuration names are **in addition to** the default ones: `ethos-u55-128` and `ethosu-u65-256`.
+- `--additional-ethos-u-config-name`: if you need to optimize the models for a different Ethos-U configuration,
+  you can pass a list of additional configurations for Vela compiler. For example, to optimize the downloaded models
+  for *Ethos™-U55* 32 MAC configuration and *Ethos™-U65* 512 MAC configuration, you can choose to run:
+
+  ```sh
+  python3 ./set_up_default_resources.py \
+    --additional-ethos-u-config-name ethos-u55-32 \
+    --additional-ethos-u-config-name ethos-u65-512
+  ```
+
+  > **Note:** As the argument name suggests, the configuration names are **in addition to** the default ones: `ethos-u55-128` and `ethosu-u65-256`.
+
+- `--arena-cache-size`: the size of the arena cache memory area, in bytes.
+  The default value is:
+  - the internal SRAM size for Corstone-300 implementation on MPS3 specified by AN552,
+  when optimizing for the default 128 MACs configuration of the Arm® *Ethos™-U55* NPU.
+  - the default value specified in the Vela configuration file  [default_vela.ini](../../scripts/vela/default_vela.ini),
+  when optimizing for the default 256 MACs configuration of the Arm® *Ethos™-U65* NPU.
 
 ### Building for default configuration
 
 A helper script `build_default.py` is provided to configure and build all the applications. It configures the project
-with default settings i.e., for `mps3` target, `sse-300` subsystem and *Ethos-U55* timing-adapter settings. Under the hood, it invokes all the necessary
+with default settings i.e., for `mps3` target, `sse-300` subsystem and *Ethos-U55* timing-adapter settings.
+Under the hood, it invokes all the necessary
 CMake commands that are described in the next sections.
 
 If using the `Arm GNU embedded toolchain`, execute:
@@ -353,16 +364,16 @@ Additional command line arguments supported by this script are:
 - `--skip-vela`: Do not run Vela optimizer on downloaded models.
 - `--npu-config-name`: Arm Ethos-U configuration to build for. The default value is
     `ethos-u55-128`. Valid values are:
-    - `ethos-u55-32`
-    - `ethos-u55-64`
-    - `ethos-u55-128`
-    - `ethos-u55-256`
-    - `ethos-u65-256`
-    - `ethos-u65-512`
+  - `ethos-u55-32`
+  - `ethos-u55-64`
+  - `ethos-u55-128`
+  - `ethos-u55-256`
+  - `ethos-u65-256`
+  - `ethos-u65-512`
 - `--make-jobs`: Specifies the number of concurrent jobs to use for compilation.
 The default value is equal to the number of cores in the system.
 Lowering this value can be useful in case of limited resources.
-- `--make-verbose`: Make the compile process verbose. This is equal to run ```make VERBOSE=1```. 
+- `--make-verbose`: Make the compile process verbose. This is equal to run ```make VERBOSE=1```.
 
 To build for *Ethos™-U55* 32 MAC configuration, using `Arm Compiler`, run:
 
diff --git a/set_up_default_resources.py b/set_up_default_resources.py
index 973c423..91007e4 100755
--- a/set_up_default_resources.py
+++ b/set_up_default_resources.py
@@ -22,7 +22,7 @@ import fnmatch
 import logging
 import sys
 
-from argparse import ArgumentParser
+from argparse import ArgumentParser, ArgumentTypeError
 from urllib.error import URLError
 from collections import namedtuple
 
@@ -132,7 +132,8 @@ NPUConfig = namedtuple('NPUConfig',['config_name',
                                     'memory_mode',
                                     'system_config',
                                     'ethos_u_npu_id',
-                                    'ethos_u_config_id'])
+                                    'ethos_u_config_id',
+                                    'arena_cache_size'])
 
 # The internal SRAM size for Corstone-300 implementation on MPS3 specified by AN552
 mps3_max_sram_sz = 2 * 1024 * 1024 # 2 MiB (2 banks of 1 MiB each)
@@ -157,14 +158,19 @@ def call_command(command: str) -> str:
     return log
 
 
-def get_default_npu_config_from_name(config_name: str) -> NPUConfig:
+def get_default_npu_config_from_name(config_name: str, arena_cache_size: int = 0) -> NPUConfig:
     """
     Gets the file suffix for the tflite file from the
     `accelerator_config` string.
 
     Parameters:
     ----------
-    config_name (str):   Ethos-U NPU configuration from valid_npu_config_names
+    config_name (str):      Ethos-U NPU configuration from valid_npu_config_names
+
+    arena_cache_size (int): Specifies arena cache size in bytes. If a value
+                            greater than 0 is provided, this will be taken
+                            as the cache size. If 0, the default values, as per
+                            the NPU config requirements, are used.
 
     Returns:
     -------
@@ -182,6 +188,13 @@ def get_default_npu_config_from_name(config_name: str) -> NPUConfig:
     prefix_ids = ["H", "Y"]
     memory_modes = ["Shared_Sram", "Dedicated_Sram"]
     system_configs = ["Ethos_U55_High_End_Embedded", "Ethos_U65_High_End"]
+    memory_modes_arena = {
+        # For shared SRAM memory mode, we use the MPS3 SRAM size by default
+        "Shared_Sram"    : mps3_max_sram_sz if arena_cache_size <= 0 else arena_cache_size,
+        # For dedicated SRAM memory mode, we do no override the arena size. This is expected to
+        # be defined in the vela configuration file instead.
+        "Dedicated_Sram" : None if arena_cache_size <= 0 else arena_cache_size
+    }
 
     for i in range(len(strings_ids)):
         if config_name.startswith(strings_ids[i]):
@@ -190,22 +203,26 @@ def get_default_npu_config_from_name(config_name: str) -> NPUConfig:
                               memory_mode=memory_modes[i],
                               system_config=system_configs[i],
                               ethos_u_npu_id=processor_ids[i],
-                              ethos_u_config_id=npu_config_id)
+                              ethos_u_config_id=npu_config_id,
+                              arena_cache_size=memory_modes_arena[memory_modes[i]])
 
     return None
 
 
-def set_up_resources(run_vela_on_models=False,
-                     additional_npu_config_names=[],
-                     arena_cache_size=mps3_max_sram_sz):
+def set_up_resources(run_vela_on_models: bool = False,
+                     additional_npu_config_names: list = [],
+                     arena_cache_size: int = 0):
     """
     Helpers function that retrieve the output from a command.
 
     Parameters:
     ----------
-    run_vela_on_models (bool):    Specifies if run vela on downloaded models.
-    additional_npu_config_names(list): list of strings of Ethos-U NPU configs.
-    arena_cache_size(int):    Specifies arena cache size in bytes.
+    run_vela_on_models (bool):  Specifies if run vela on downloaded models.
+    additional_npu_config_names(list):  list of strings of Ethos-U NPU configs.
+    arena_cache_size (int): Specifies arena cache size in bytes. If a value
+                            greater than 0 is provided, this will be taken
+                            as the cache size. If 0, the default values, as per
+                            the NPU config requirements, are used.
     """
     current_file_dir = os.path.dirname(os.path.abspath(__file__))
     download_dir = os.path.abspath(os.path.join(current_file_dir, "resources_downloaded"))
@@ -301,18 +318,25 @@ def set_up_resources(run_vela_on_models=False,
         config_names = list(set(default_npu_config_names + additional_npu_config_names))
 
         # Get npu config tuple for each config name in a list:
-        npu_configs = [get_default_npu_config_from_name(name) for name in config_names]
+        npu_configs = [get_default_npu_config_from_name(name, arena_cache_size) for name in config_names]
 
         logging.info(f'All models will be optimised for these configs:')
         for config in npu_configs:
             logging.info(config)
 
+        optimisation_skipped = False
+
         for model in models:
             output_dir = os.path.dirname(model)
             # model name after compiling with vela is an initial model name + _vela suffix
             vela_optimised_model_path = str(model).replace(".tflite", "_vela.tflite")
 
             for config in npu_configs:
+                vela_command_arena_cache_size = ""
+
+                if config.arena_cache_size:
+                    vela_command_arena_cache_size = f"--arena-cache-size={config.arena_cache_size}"
+
                 vela_command = (f". {env_activate} && vela {model} " +
                        f"--accelerator-config={config.config_name} " +
                        "--optimise Performance " +
@@ -320,7 +344,7 @@ def set_up_resources(run_vela_on_models=False,
                        f"--memory-mode={config.memory_mode} " +
                        f"--system-config={config.system_config} " +
                        f"--output-dir={output_dir} " +
-                       f"--arena-cache-size={arena_cache_size} ")
+                       f"{vela_command_arena_cache_size}")
 
                 # we want the name to include the configuration suffix. For example: vela_H128,
                 # vela_Y512 etc.
@@ -330,6 +354,7 @@ def set_up_resources(run_vela_on_models=False,
 
                 if os.path.isfile(new_vela_optimised_model_path):
                     logging.info(f"File {new_vela_optimised_model_path} exists, skipping optimisation.")
+                    optimisation_skipped = True
                     continue
 
                 call_command(vela_command)
@@ -339,6 +364,12 @@ def set_up_resources(run_vela_on_models=False,
                 logging.info(f"Renaming {vela_optimised_model_path} to {new_vela_optimised_model_path}.")
 
 
+        # If any optimisation was skipped, show how to regenerate:
+        if optimisation_skipped:
+            logging.warning("One or more optimisations were skipped.")
+            logging.warning(f"To optimise all the models, please remove the directory {download_dir}.")
+
+
 if __name__ == '__main__':
     parser = ArgumentParser()
     parser.add_argument("--skip-vela",
@@ -349,11 +380,14 @@ if __name__ == '__main__':
                         {valid_npu_config_names}""",
                         default=[], action="append")
     parser.add_argument("--arena-cache-size",
-                        help="Arena cache size in bytes",
+                        help="Arena cache size in bytes (if overriding the defaults)",
                         type=int,
-                        default=mps3_max_sram_sz)
+                        default=0)
     args = parser.parse_args()
 
+    if args.arena_cache_size < 0:
+        raise ArgumentTypeError('Arena cache size cannot not be less than 0')
+
     logging.basicConfig(filename='log_build_default.log', level=logging.DEBUG)
     logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))