Arm backend: Add options when using run.sh (pytorch#5627)

zingo · facebook-github-bot · commit b71926f93af3 · 2024-09-30T08:41:05.000-07:00
Summary: This will enable you to build for different target, specify models from the command line and also control output folder used. Pull Request resolved: pytorch#5627 Reviewed By: digantdesai Differential Revision: D63637343 Pulled By: mergennachin fbshipit-source-id: cec5830b0a6a406cd72b7a2fecf3ff229fa2d359
diff --git a/examples/arm/README.md b/examples/arm/README.md
@@ -25,7 +25,7 @@ $ ./setup.sh --i-agree-to-the-contained-eula [optional-scratch-dir]
 
 # Step [2] - build + run ExecuTorch and executor_runner baremetal application
 # suited for Corstone300 to run a simple PyTorch model.
-$ ./run.sh [same-optional-scratch-dir-as-before]
+$ ./run.sh [--scratch-dir=same-optional-scratch-dir-as-before]
 ```
 ### Online Tutorial
 
diff --git a/examples/arm/aot_arm_compiler.py b/examples/arm/aot_arm_compiler.py
@@ -148,8 +148,15 @@ def forward(self, x):
 }
 
 targets = [
-    "ethos-u85-128",
+    "ethos-u55-32",
+    "ethos-u55-64",
     "ethos-u55-128",
+    "ethos-u55-256",
+    "ethos-u85-128",
+    "ethos-u85-256",
+    "ethos-u85-512",
+    "ethos-u85-1024",
+    "ethos-u85-2048",
     "TOSA",
 ]
 
@@ -160,23 +167,23 @@ def get_compile_spec(target: str, intermediates: bool) -> ArmCompileSpecBuilder:
         spec_builder = (
             ArmCompileSpecBuilder().tosa_compile_spec().set_permute_memory_format(True)
         )
-    elif target == "ethos-u55-128":
+    elif "ethos-u55" in target:
         spec_builder = (
             ArmCompileSpecBuilder()
             .ethosu_compile_spec(
-                "ethos-u55-128",
+                target,
                 system_config="Ethos_U55_High_End_Embedded",
                 memory_mode="Shared_Sram",
                 extra_flags="--debug-force-regor --output-format=raw",
             )
             .set_permute_memory_format(args.model_name in MODEL_NAME_TO_MODEL.keys())
             .set_quantize_io(True)
         )
-    elif target == "ethos-u85-128":
+    elif "ethos-u85" in target:
         spec_builder = (
             ArmCompileSpecBuilder()
             .ethosu_compile_spec(
-                "ethos-u85-128",
+                target,
                 system_config="Ethos_U85_SYS_DRAM_Mid",
                 memory_mode="Shared_Sram",
                 extra_flags="--output-format=raw",
@@ -323,7 +330,9 @@ def get_args():
 
     model_name = os.path.basename(os.path.splitext(args.model_name)[0])
     output_name = f"{model_name}" + (
-        f"_arm_delegate_{args.target}" if args.delegate is True else ""
+        f"_arm_delegate_{args.target}"
+        if args.delegate is True
+        else f"_arm_{args.target}"
     )
 
     if args.output is not None:
diff --git a/examples/arm/executor_runner/arm_perf_monitor.cpp b/examples/arm/executor_runner/arm_perf_monitor.cpp
@@ -39,10 +39,20 @@ void ethosu_inference_begin(struct ethosu_driver* drv, void*) {
   ETHOSU_PMU_PMCCNTR_CFG_Set_Start_Event(drv, ETHOSU_PMU_NPU_ACTIVE);
 
   // Setup 4 counters
+#if defined(ETHOSU55) || defined(ETHOSU65)
   ETHOSU_PMU_Set_EVTYPER(drv, 0, ETHOSU_PMU_AXI0_RD_DATA_BEAT_RECEIVED);
   ETHOSU_PMU_Set_EVTYPER(drv, 1, ETHOSU_PMU_AXI1_RD_DATA_BEAT_RECEIVED);
   ETHOSU_PMU_Set_EVTYPER(drv, 2, ETHOSU_PMU_AXI0_WR_DATA_BEAT_WRITTEN);
   ETHOSU_PMU_Set_EVTYPER(drv, 3, ETHOSU_PMU_NPU_IDLE);
+#elif defined(ETHOSU85)
+  ETHOSU_PMU_Set_EVTYPER(drv, 0, ETHOSU_PMU_EXT0_RD_DATA_BEAT_RECEIVED);
+  ETHOSU_PMU_Set_EVTYPER(drv, 1, ETHOSU_PMU_EXT1_RD_DATA_BEAT_RECEIVED);
+  ETHOSU_PMU_Set_EVTYPER(drv, 2, ETHOSU_PMU_EXT0_WR_DATA_BEAT_WRITTEN);
+  ETHOSU_PMU_Set_EVTYPER(drv, 3, ETHOSU_PMU_NPU_IDLE);
+#else
+#error No NPU target defined
+#endif
+
   // Enable 4 counters
   ETHOSU_PMU_CNTR_Enable(drv, 0xf);
 
@@ -160,9 +170,17 @@ void StopMeasurements() {
   for (size_t i = 0; i < ethosu_pmuCountersUsed; i++) {
     ET_LOG(Info, "ethosu_pmu_cntr%zd : %" PRIu64, i, ethosu_pmuEventCounts[i]);
   }
+#if defined(ETHOSU55) || defined(ETHOSU65)
   ET_LOG(
       Info,
       "Ethos-U PMU Events:[ETHOSU_PMU_AXI0_RD_DATA_BEAT_RECEIVED, ETHOSU_PMU_AXI1_RD_DATA_BEAT_RECEIVED, ETHOSU_PMU_AXI0_WR_DATA_BEAT_WRITTEN, ETHOSU_PMU_NPU_IDLE]");
+#elif defined(ETHOSU85)
+  ET_LOG(
+      Info,
+      "Ethos-U PMU Events:[ETHOSU_PMU_EXT0_RD_DATA_BEAT_RECEIVED, ETHOSU_PMU_EXT1_RD_DATA_BEAT_RECEIVED, ETHOSU_PMU_EXT0_WR_DATA_BEAT_WRITTEN, ETHOSU_PMU_NPU_IDLE]");
+#else
+#error No NPU target defined
+#endif
 }
 
 #else
diff --git a/examples/arm/run.sh b/examples/arm/run.sh
@@ -9,20 +9,58 @@
 
 set -eu
 
-if [[ "${1:-'.'}" == "-h" || "${#}" -gt 2 ]]; then
-    echo "Usage: $(basename $0) [path-to-a-scratch-dir]"
-    echo "Supplied args: $*"
-    exit 1
-fi
+
 
 ########
 ### Hardcoded constants
 ########
 script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
 
-# Ethos-u
-root_dir=${1:-"${script_dir}/ethos-u-scratch"}
+# Default Ethos-u tool folder override with --scratch-dir=<FOLDER>
+root_dir=${script_dir}/ethos-u-scratch
+
+model_name=""
+aot_arm_compiler_flags="--delegate --quantize"
+target="ethos-u55-128"
+output_folder_set=false
+output_folder="."
+build_only=false
+
+help() {
+    echo "Usage: $(basename $0) [options]"
+    echo "Options:"
+    echo "  --model_name=<MODEL>                   Model to run, can be a builtin, examples/models or a filename Default to all builtin models"
+    echo "  --aot_arm_compiler_flags=<FLAGS>       Only used if --model_name is used Default: ${aot_arm_compiler_flags}"
+    echo "  --target=<TARGET>                      Target to build and run for Default: ${target}"
+    echo "  --output=<FOLDER>                      Output folder Default: ${output_folder}"
+    echo "  --build_only                           Only build, don't run FVP"
+    echo "  --scratch-dir=<FOLDER>                 Path to your Ethos-U scrach dir if you not using default"
+    exit 0
+}
+
+for arg in "$@"; do
+    case $arg in
+      -h|--help) help ;;
+      --model_name=*) model_name="${arg#*=}";;
+      --aot_arm_compiler_flags=*) aot_arm_compiler_flags="${arg#*=}";;
+      --target=*) target="${arg#*=}";;
+      --output=*) output_folder="${arg#*=}" ; output_folder_set=true ;;
+      --build_only) build_only=true ;;
+      --scratch-dir=*) root_dir="${arg#*=}";;
+      *)
+      ;;
+    esac
+done
+
 root_dir=$(realpath ${root_dir})
+output_folder=$(realpath ${output_folder})
+mkdir -p ${output_folder}
+if [ "$output_folder_set" = true ] ; then
+    executor_runner_path=${output_folder}
+else
+    executor_runner_path=${script_dir}/executor_runner
+fi
+executor_runner_path=$(realpath ${executor_runner_path})
 
 ethos_u_root_dir="$(cd ${root_dir}/ethos-u && pwd)"
 ethos_u_build_dir=${ethos_u_root_dir}/core_platform/build
@@ -38,26 +76,26 @@ _setup_msg="please refer to ${script_dir}/ethos-u-setup/setup.sh to properly ins
 
 # Generate a pte file
 function generate_pte_file() {
-    [[ $# -ne 2 ]] && { echo "[${FUNCNAME[0]}]" "Expecting model and delegate flag, got, $*"; exit 1; }
+    [[ $# -ne 2 ]] && { echo "[${FUNCNAME[0]}]" "Expecting model and model_compiler_flags flag, got, $*"; exit 1; }
     local model=${1}
-    local delegate=${2}
+    local model_compiler_flags=${2}
 
-    local model_filename=${model}.pte
-    if [[ "${delegate}" == *"--delegate"* ]]; then
+    local model_filename=${model}_arm_${target}.pte
+    if [[ "${model_compiler_flags}" == *"--delegate"* ]]; then
 	# Name aligned with default aot_arm_compiler output - run.sh only supports
 	# running on Corstone-300 with Ethos-U55 FVP at the moment.
-        model_filename=${model}_arm_delegate_ethos-u55-128.pte
+        model_filename=${model}_arm_delegate_${target}.pte
     fi
     cd $et_root_dir
 
     local pte_file
-    pte_file=$(realpath ${model_filename})
+    pte_file=$(realpath ${output_folder}/${model_filename})
     rm -f "${pte_file}"
 
     # We are using the aot_lib from build_quantization_aot_lib below
     SO_LIB=$(find cmake-out-aot-lib -name libquantized_ops_aot_lib.so)
 
-    python3 -m examples.arm.aot_arm_compiler --model_name="${model}" ${delegate} --so_library="$SO_LIB" 1>&2
+    python3 -m examples.arm.aot_arm_compiler --model_name="${model}" --target=${target} ${model_compiler_flags} --output ${output_folder} --so_library="$SO_LIB" 1>&2
     [[ -f ${pte_file} ]] || { >&2 echo "Failed to generate a pte file - ${pte_file}"; exit 1; }
     echo "${pte_file}"
 }
@@ -100,7 +138,7 @@ function build_executorch() {
         -DCMAKE_BUILD_TYPE=Release                        \
         -DEXECUTORCH_ENABLE_LOGGING=ON                    \
         -DEXECUTORCH_BUILD_ARM_BAREMETAL=ON               \
-        -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON                   \
+        -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON           \
         -DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=ON       \
         -DFLATC_EXECUTABLE="$(which flatc)"               \
         -DCMAKE_TOOLCHAIN_FILE="${toolchain_cmake}"       \
@@ -134,43 +172,55 @@ function build_executorch_runner() {
     echo "[${FUNCNAME[0]}] Generating ExecuTorch libraries"
     [[ $# -ne 1 ]] && { echo "[${FUNCNAME[0]}]" "Expecting a single pte file as argument got, $*"; exit 1; }
     local pte=${1}
+    if [[ ${target} == *"ethos-u55"*  ]]; then
+        local target_cpu=cortex-m55
+    else
+        local target_cpu=cortex-m85
+    fi
     cd ${script_dir}/executor_runner
     cmake -DCMAKE_TOOLCHAIN_FILE=${toolchain_cmake} \
-	  -DTARGET_CPU=cortex-m55 \
-	  -B cmake-out \
-	  -DETHOS_SDK_PATH:PATH=${ethos_u_root_dir} \
-	  -DET_DIR_PATH:PATH=${et_root_dir}         \
-	  -DET_BUILD_DIR_PATH:PATH=${et_build_dir}  \
-	  -DET_PTE_FILE_PATH:PATH="${pte}"          \
+	  -DTARGET_CPU=${target_cpu}                    \
+      -DETHOSU_TARGET_NPU_CONFIG=${target}          \
+	  -B ${executor_runner_path}/cmake-out          \
+	  -DETHOS_SDK_PATH:PATH=${ethos_u_root_dir}     \
+	  -DET_DIR_PATH:PATH=${et_root_dir}             \
+	  -DET_BUILD_DIR_PATH:PATH=${et_build_dir}      \
+	  -DET_PTE_FILE_PATH:PATH="${pte}"              \
 	  -DPYTHON_EXECUTABLE=$(which python3)
     echo "[${FUNCNAME[0]}] Configured CMAKE"
 
     n=$(nproc)
-    cmake --build cmake-out -- -j"$((n - 5))" arm_executor_runner
+    cmake --build ${executor_runner_path}/cmake-out -- -j"$((n - 5))" arm_executor_runner
     echo "[${FUNCNAME[0]}] Generated baremetal elf file:"
-    find cmake-out -name "arm_executor_runner"
-    echo "executable_text: $(find -name arm_executor_runner -exec size {} \; | grep -v filename | awk '{print $1}') bytes"
-    echo "executable_data: $(find -name arm_executor_runner -exec size {} \; | grep -v filename | awk '{print $2}') bytes"
-    echo "executable_bss:  $(find -name arm_executor_runner -exec size {} \; | grep -v filename | awk '{print $3}') bytes"
-    echo "pte_data_size:   $(stat -c%s ${pte}) bytes"
+    find ${executor_runner_path}/cmake-out -name "arm_executor_runner"
+    echo "executable_text: $(find ${executor_runner_path}/cmake-out -name arm_executor_runner -exec size {} \; | grep -v filename | awk '{print $1}') bytes"
+    echo "executable_data: $(find ${executor_runner_path}/cmake-out -name arm_executor_runner -exec size {} \; | grep -v filename | awk '{print $2}') bytes"
+    echo "executable_bss:  $(find ${executor_runner_path}/cmake-out -name arm_executor_runner -exec size {} \; | grep -v filename | awk '{print $3}') bytes"
 }
 
 # Execute the executor_runner on FVP Simulator
 function run_fvp() {
     [[ $# -ne 1 ]] && { echo "[${FUNCNAME[0]}]" "Expexted elf binary name, got $*"; exit 1; }
     local elf_name=${1}
-    elf=$(find ${script_dir}/executor_runner -name "${elf_name}")
+    elf=$(find ${executor_runner_path} -name "${elf_name}")
     [[ ! -f $elf ]] && { echo "[${FUNCNAME[0]}]: Unable to find executor_runner elf: ${elf}"; exit 1; }
-    FVP_Corstone_SSE-300_Ethos-U55                          \
-        -C cpu0.CFGITCMSZ=11                                \
-        -C ethosu.num_macs=128                              \
-        -C mps3_board.visualisation.disable-visualisation=1 \
-        -C mps3_board.telnetterminal0.start_telnet=0        \
-        -C mps3_board.uart0.out_file='-'                    \
-        -C mps3_board.uart0.shutdown_on_eot=1               \
-        -a "${elf}"                                         \
-        --timelimit 120 || true # seconds
-    echo "[${FUNCNAME[0]} Simulation complete, $?"
+    num_macs=$(echo ${target} | cut -d - -f 3)
+    if [[ ${target} == *"ethos-u55"*  ]]; then
+        echo "Running ${elf} for ${target} run with FVP_Corstone_SSE-300_Ethos-U55 num_macs:${num_macs}"
+        FVP_Corstone_SSE-300_Ethos-U55                          \
+            -C cpu0.CFGITCMSZ=11                                \
+            -C ethosu.num_macs=${num_macs}                      \
+            -C mps3_board.visualisation.disable-visualisation=1 \
+            -C mps3_board.telnetterminal0.start_telnet=0        \
+            -C mps3_board.uart0.out_file='-'                    \
+            -C mps3_board.uart0.shutdown_on_eot=1               \
+            -a "${elf}"                                         \
+            --timelimit 120 || true # seconds
+        echo "[${FUNCNAME[0]} Simulation complete, $?"
+    else
+        echo "Running ${elf} for ${target} is not supported"
+        exit 1
+    fi
 }
 
 #######
@@ -199,17 +249,31 @@ hash arm-none-eabi-gcc \
 build_executorch
 build_quantization_aot_lib
 
-# the test models run, and whether to delegate
-test_model=( "softmax" "add" "add3" "mv2" )
-test_delegate=( "" "--delegate" "--delegate" "--delegate --quantize" )
+if [[ -z "$model_name" ]]; then
+    # the test models run, and whether to delegate
+    test_model=( "softmax" "add" "add3" "mv2" )
+    model_compiler_flags=( "" "--delegate" "--delegate" "--delegate --quantize" )
+else
+    test_model=( "$model_name" )
+    model_compiler_flags=( "$aot_arm_compiler_flags" )
+fi
 
 # loop over running the AoT flow and executing the model on device
 for i in "${!test_model[@]}"; do
-    printf "Running e2e flow for model '%s' with flags '%s'\n" "${test_model[i]}" "${test_delegate[i]}"
-    pte=$(generate_pte_file "${test_model[i]}" "${test_delegate[i]}")
-    # Rebuild the application as the pte is imported as a header/c array
-    build_executorch_runner "${pte}"
-    run_fvp arm_executor_runner
+    echo "--------------------------------------------------------------------------------"
+    printf "Running e2e flow for model '%s' with flags '%s'\n" "${test_model[i]}" "${model_compiler_flags[i]}"
+    echo "--------------------------------------------------------------------------------"
+    pte=$(generate_pte_file "${test_model[i]}" "${model_compiler_flags[i]}")
+    stat --printf="Generated pte_data_size: %s bytes\npte_file:%n\n" ${pte}
+    if [[ ${target} == *"TOSA"*  ]]; then
+        echo "Build for ${target} skip generating .elf and running"
+    else
+        # Rebuild the application as the pte is imported as a header/c array
+        build_executorch_runner "${pte}"
+        if [ "$build_only" = false ] ; then
+            run_fvp arm_executor_runner
+        fi
+    fi
 done
 
 exit 0
