Arm backend: Make the logic for temp_allocation_pool realistic

backends/arm/scripts/corstone_utils.cmake

@@ -81,20 +81,6 @@ function(add_corstone_subdirectory SYSTEM_CONFIG ETHOS_SDK_PATH)
   else()
     message(FATAL_ERROR "Unsupported SYSTEM_CONFIG ${SYSTEM_CONFIG}.")
   endif()
-  if(MEMORY_MODE MATCHES "Dedicated_Sram")
-    target_compile_definitions(
-      ethosu_target_common INTERFACE ETHOSU_MODEL=1 ETHOSU_ARENA=1
-    )
-  elseif(MEMORY_MODE MATCHES "Shared_Sram" OR MEMORY_MODE MATCHES "Sram_Only")
-    target_compile_definitions(
-      ethosu_target_common INTERFACE ETHOSU_MODEL=1 ETHOSU_ARENA=0
-    )
-  else()
-    message(
-      FATAL_ERROR
-        "Unsupported MEMORY_MODE ${MEMORY_MODE}. Memory_mode can be Shared_Sram, Sram_Only or Dedicated_Sram(applicable for the Ethos-U85)"
-    )
-  endif()
 endfunction()
 
 function(configure_timing_adapters SYSTEM_CONFIG MEMORY_MODE)

backends/arm/scripts/get_ethosu_scratch_from_pte.py

@@ -38,6 +38,10 @@ def get_scratch_size_from_delegate(blob: bytes) -> int | None:
 
 def _extract_pte_from_bundle(pte_data: bytes) -> bytes:
     try:
+        # If no mention of BP08(specifier for bundled program), return pte
+        if len(pte_data) < 8 or pte_data[4:8] != b"BP08":
+            return pte_data
+        # bundled program
         bundled = deserialize_from_flatbuffer_to_bundled_program(pte_data)
     except Exception:
         return pte_data
@@ -62,10 +66,8 @@ def get_scratch_from_pte(pte_path: str) -> int | None:
     if not sizes:
         return None
 
-    for did, s in sizes:
-        print(f"{did}: scratch_size={s} bytes")
     max_size = max(s for _, s in sizes)
-    print(f"max_scratch_size={max_size} bytes")
+    print(f"{max_size} bytes needed for the scratch buffer")
     return max_size
 
 

backends/arm/test/common.py

@@ -121,8 +121,8 @@ def get_u55_compile_spec(
 def get_u85_compile_spec(
     macs: int = 128,
     system_config="Ethos_U85_SYS_DRAM_Mid",
-    memory_mode="Shared_Sram",
-    extra_flags="--output-format=raw --arena-cache-size=2097152",
+    memory_mode="Dedicated_Sram_384KB",
+    extra_flags="--output-format=raw",
     custom_path: Optional[str] = None,
     config: Optional[str] = None,
     tosa_debug_mode: EthosUCompileSpec.DebugMode | None = None,

backends/arm/test/test_model.py

@@ -70,13 +70,6 @@ def get_args():
         default="",
         help="Extra cmake flags to pass the when building the executor_runner",
     )
-    parser.add_argument(
-        "--specify_ethosu_scratch",
-        action="store_true",
-        required=False,
-        default=False,
-        help="Use PTE scratch size to set Ethos-U temp allocator pool size",
-    )
     parser.add_argument(
         "--extra_runtime_flags",
         required=False,
@@ -185,21 +178,7 @@ def build_ethosu_runtime(
     memory_mode: str,
     extra_flags: str,
     elf_build_path: str,
-    specify_ethosu_scratch: bool,
 ):
-    if specify_ethosu_scratch:
-        from executorch.backends.arm.scripts.get_ethosu_scratch_from_pte import (
-            get_scratch_from_pte,
-        )
-
-        scratch_size = get_scratch_from_pte(pte_file)
-        if scratch_size is not None:
-            extra_flags = (
-                f"{extra_flags} "
-                f"-DET_ARM_BAREMETAL_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE={scratch_size}"
-            )
-        else:
-            raise RuntimeError(f"Failed to derive Ethos-U scratch size from {pte_file}")
 
     elf_build_path = os.path.join(elf_build_path, "cmake-out")
     run_external_cmd(
@@ -349,7 +328,6 @@ def run_vkml(script_path: str, pte_file: str, runner_build_path: str, extra_flag
                 args.memory_mode,
                 args.extra_flags,
                 elf_build_path,
-                args.specify_ethosu_scratch,
             )
             end_time = time.perf_counter()
             print(

examples/arm/executor_runner/CMakeLists.txt

@@ -177,21 +177,32 @@ message(STATUS "SYSTEM_CONFIG is ${SYSTEM_CONFIG}")
 message(STATUS "MEMORY_MODE is ${MEMORY_MODE}")
 message(STATUS "ET_NUM_INFERENCES is ${ET_NUM_INFERENCES}")
 
-# By default, use 2MB of temporary scratch buffer. For Dedicated_Sram, use 64MB
-# for the temporary scratch buffer and 384KB for the fast scratch buffer (the
-# cache, applicable only for Ethos-U65 and Ethos-U85). Allow -D overrides.
-if(MEMORY_MODE MATCHES "Dedicated_Sram")
-  if(NOT DEFINED ET_ARM_BAREMETAL_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE)
-    set(ET_ARM_BAREMETAL_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE 0x4000000)
-  endif()
-  if(NOT DEFINED ET_ARM_BAREMETAL_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE)
-    set(ET_ARM_BAREMETAL_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE 0x60000)
-  endif()
-else()
-  if(NOT DEFINED ET_ARM_BAREMETAL_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE)
+if(NOT DEFINED ET_ARM_BAREMETAL_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE)
+  if(SYSTEM_CONFIG MATCHES "U55")
+    # The Corstone-300 has 2MB of SRAM, provide the Ethos-U with the full 2MB
+    # for the temp_allocation_pool array.
     set(ET_ARM_BAREMETAL_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE 0x200000)
+  elseif(SYSTEM_CONFIG MATCHES "U85")
+    if(MEMORY_MODE MATCHES "Dedicated_Sram")
+      # 32MB of scratch buffer for Dedicated_Sram memory mode.
+      set(ET_ARM_BAREMETAL_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE 0x2000000)
+      # For Dedicated_Sram, set the
+      # ET_ARM_BAREMETAL_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE to 384KB unless
+      # specified otherwise.
+      if(NOT DEFINED ET_ARM_BAREMETAL_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE)
+        set(ET_ARM_BAREMETAL_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE 0x60000)
+      endif()
+
+    elseif(MEMORY_MODE MATCHES "Shared_Sram" OR MEMORY_MODE MATCHES "Sram_Only")
+      # For Shared_Sram and Sram only, use scratch buffer of 4MB - 64KB The
+      # Corstone-320 provides 4MB of SRAM and we subtract 64KB because we have a
+      # few objects placed at the start of the SRAM, before the
+      # temp_allocation_pool array.
+      set(ET_ARM_BAREMETAL_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE 0x3F0000)
+    endif()
   endif()
 endif()
+
 message(
   STATUS
     "ET_ARM_BAREMETAL_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE = ${ET_ARM_BAREMETAL_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE}"
@@ -282,9 +293,24 @@ set(LINK_FILE_OUT
     "${CMAKE_CURRENT_BINARY_DIR}/${LINK_FILE_OUT_BASE}.${LINK_FILE_EXT}"
 )
 
+if(MEMORY_MODE MATCHES "Dedicated_Sram")
+  set(ETHOSU_MODEL 1)
+  set(ETHOSU_ARENA 1)
+elseif(MEMORY_MODE MATCHES "Shared_Sram" OR MEMORY_MODE MATCHES "Sram_Only")
+  set(ETHOSU_ARENA 0)
+  # On silicon, for Sram_Only, the model is assumed to be in the SRAM but that
+  # limits the number of models we test for Sram_Only. For testing coverage &
+  # consistency, we places the model in the external memory for Sram_Only
+  set(ETHOSU_MODEL 1)
+else()
+  message(FATAL_ERROR "Unsupported MEMORY_MODE ${MEMORY_MODE}")
+endif()
+
 execute_process(
-  COMMAND ${CMAKE_C_COMPILER} ${COMPILER_PREPROCESSOR_OPTIONS} -o
-          ${LINK_FILE_OUT} ${LINK_FILE_IN}
+  COMMAND
+    ${CMAKE_C_COMPILER} ${COMPILER_PREPROCESSOR_OPTIONS}
+    -DETHOSU_ARENA=${ETHOSU_ARENA} -DETHOSU_MODEL=${ETHOSU_MODEL} -o
+    ${LINK_FILE_OUT} ${LINK_FILE_IN}
 )
 target_link_options(arm_executor_runner PRIVATE "-T" "${LINK_FILE_OUT}")
 

examples/arm/run.sh

@@ -48,7 +48,6 @@ perf_overlay=false
 visualize_tosa=false
 visualize_pte=false
 model_converter=false
-specify_ethosu_scratch=false
 extra_build_flags=""
 preset_file="${et_root_dir}/tools/cmake/preset/arm_baremetal.cmake"
 cmake_cache_file=""
@@ -82,7 +81,6 @@ function help() {
     echo "  --config=<FILEPATH>                    Ethos-U: System configuration file that specifies system configurations (vela.ini)"
     echo "  --memory_mode=<MODE>                   Ethos-U: Memory mode to select from the Vela configuration file (see vela.ini), e.g. Shared_Sram/Sram_Only. Default: 'Shared_Sram' for Ethos-U55 targets, 'Sram_Only' for Ethos-U85 targets"
     echo "  --pte_placement=<elf|ADDR>             Ethos-U: Control if runtime has PTE baked into the elf or if its placed in memory outside of the elf, defaults to ${pte_placement}"
-    echo "  --specify_ethosu_scratch               Use actual Ethos-U scratch size for given model to size temp allocator"
     echo "  --et_build_root=<FOLDER>               Executorch build output root folder to use, defaults to ${et_build_root}"
     echo "  --scratch-dir=<FOLDER>                 Path to your Ethos-U scratch dir if you not using default ${arm_scratch_dir}"
     echo "  --qdq_fusion_op                        Enable QDQ fusion op"
@@ -119,7 +117,6 @@ for arg in "$@"; do
       --config=*) config="${arg#*=}";;
       --memory_mode=*) memory_mode="${arg#*=}";;
       --pte_placement=*) pte_placement="${arg#*=}";;
-      --specify_ethosu_scratch) specify_ethosu_scratch=true ;;
       --et_build_root=*) et_build_root="${arg#*=}";;
       --scratch-dir=*) arm_scratch_dir="${arg#*=}" ; scratch_dir_set=true ;;
       --qdq_fusion_op) qdq_fusion_op=true;;
@@ -254,21 +251,6 @@ function check_setup () {
     return 0
 }
 
-function get_ethosu_scratch_size() {
-    local pte_path="$1"
-    python3 - "$pte_path" <<'PY'
-from executorch.backends.arm.scripts.get_ethosu_scratch_from_pte import (
-    get_scratch_from_pte,
-)
-import sys
-
-size = get_scratch_from_pte(sys.argv[1])
-if size is None:
-    sys.exit(2)
-print(size)
-PY
-}
-
 sanitize_for_path() {
     local value="$1"
     printf '%s' "${value}" | tr -c '[:alnum:]._-' '_'
@@ -351,6 +333,7 @@ configure_runner_build_dir() {
         cmake_cmd+=("${extra_args[@]}")
     fi
     echo "[run.sh] Configuring ExecuTorch build at ${build_dir}"
+    echo "[run.sh] Configuring the build system with ${cmake_cmd[@]}"
     "${cmake_cmd[@]}"
     build_dir_initialized=false
 }
@@ -653,30 +636,6 @@ EOF
     echo "${cache_path}"
 }
 
-configure_ethosu_scratch_if_requested() {
-    local pte_path="$1"
-    if [ "$specify_ethosu_scratch" != true ] || [[ ! ${target} =~ "ethos-u" ]]; then
-        return
-    fi
-    local scratch_size
-    scratch_size=$(get_ethosu_scratch_size "$pte_path" | tail -n 1)
-    if [[ -z "${scratch_size}" ]]; then
-        echo "WARNING: Failed to derive Ethos-U scratch size from ${pte_path}" >&2
-        return
-    fi
-    local cmake_cmd=(
-        cmake -S "${runner_source_dir}" -B "${build_dir}"
-    )
-    if [[ -n "${extra_build_flags}" ]]; then
-        # shellcheck disable=SC2206
-        local extra_args=(${extra_build_flags})
-        cmake_cmd+=("${extra_args[@]}")
-    fi
-    cmake_cmd+=("-DET_ARM_BAREMETAL_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE=${scratch_size}")
-    echo "[run.sh] Updating scratch allocator size to ${scratch_size}"
-    "${cmake_cmd[@]}"
-}
-
 if [[ -z "$model_name" ]]; then
     echo "[run.sh] WARNING: Built-in test models executed when --model_name is omitted are deprecated and will be removed after the ExecuTorch 1.2 release." >&2
     # the test models run, and whether to delegate
@@ -809,7 +768,6 @@ for i in "${!test_model[@]}"; do
             model_data="--data=${pte_file}@${pte_placement}"
         fi
 
-        configure_ethosu_scratch_if_requested "${pte_file}"
 
         build_runner_target arm_executor_runner
         elf_file=$(locate_runner_binary arm_executor_runner) \