diff --git a/.ci/docker/common/install_android.sh b/.ci/docker/common/install_android.sh index d4dfe64b7d9..6e6f1e1568d 100755 --- a/.ci/docker/common/install_android.sh +++ b/.ci/docker/common/install_android.sh @@ -43,10 +43,10 @@ install_ndk() { ARCH=$(uname -m) if [ "${ARCH}" = "aarch64" ]; then # aarch64 NDK is not cached on S3, download from Google directly - curl -Os --retry 3 "https://dl.google.com/android/repository/android-ndk-${ANDROID_NDK_VERSION}-linux.zip" + curl -Os --retry 3 --retry-all-errors "https://dl.google.com/android/repository/android-ndk-${ANDROID_NDK_VERSION}-linux.zip" else # The NDK installation is cached on ossci-android S3 bucket - curl -Os --retry 3 "https://ossci-android.s3.amazonaws.com/android-ndk-${ANDROID_NDK_VERSION}-linux.zip" + curl -Os --retry 3 --retry-all-errors "https://ossci-android.s3.amazonaws.com/android-ndk-${ANDROID_NDK_VERSION}-linux.zip" fi unzip -qo "android-ndk-${ANDROID_NDK_VERSION}-linux.zip" @@ -62,7 +62,7 @@ install_cmdtools() { pushd /tmp # The file is cached on ossci-android S3 bucket - curl -Os --retry 3 "https://ossci-android.s3.us-west-1.amazonaws.com/${CMDTOOLS_FILENAME}" + curl -Os --retry 3 --retry-all-errors "https://ossci-android.s3.us-west-1.amazonaws.com/${CMDTOOLS_FILENAME}" unzip -qo "${CMDTOOLS_FILENAME}" -d /opt ls -lah /opt/cmdline-tools/bin diff --git a/.ci/docker/common/install_cache.sh b/.ci/docker/common/install_cache.sh index f17b4cc6068..7b7d39994ca 100755 --- a/.ci/docker/common/install_cache.sh +++ b/.ci/docker/common/install_cache.sh @@ -34,7 +34,7 @@ install_ubuntu() { install_binary() { echo "Downloading sccache binary from S3 repo" - curl --retry 3 https://s3.amazonaws.com/ossci-linux/sccache -o /opt/cache/bin/sccache + curl --retry 3 --retry-all-errors https://s3.amazonaws.com/ossci-linux/sccache -o /opt/cache/bin/sccache chmod +x /opt/cache/bin/sccache } diff --git a/.ci/docker/common/install_docs_reqs.sh b/.ci/docker/common/install_docs_reqs.sh index 13aba8e5bd5..3b6d10c5c2b 100755 --- a/.ci/docker/common/install_docs_reqs.sh +++ b/.ci/docker/common/install_docs_reqs.sh @@ -12,10 +12,10 @@ if [ -n "$BUILD_DOCS" ]; then # Ignore error if gpg-agent doesn't exist (for Ubuntu 16.04) apt-get install -y gpg-agent || : - curl --retry 3 -sL https://deb.nodesource.com/setup_16.x | sudo -E bash - + curl --retry 3 --retry-all-errors -sL https://deb.nodesource.com/setup_16.x | sudo -E bash - sudo apt-get install -y nodejs - curl --retry 3 -sS https://dl.yarnpkg.com/debian/pubkey.gpg | sudo apt-key add - + curl --retry 3 --retry-all-errors -sS https://dl.yarnpkg.com/debian/pubkey.gpg | sudo apt-key add - echo "deb https://dl.yarnpkg.com/debian/ stable main" | sudo tee /etc/apt/sources.list.d/yarn.list apt-get update diff --git a/.ci/docker/common/install_linter.sh b/.ci/docker/common/install_linter.sh index d262176e49b..52d2d262685 100755 --- a/.ci/docker/common/install_linter.sh +++ b/.ci/docker/common/install_linter.sh @@ -15,5 +15,5 @@ source "$(dirname "${BASH_SOURCE[0]}")/utils.sh" pip_install -r requirements-lintrunner.txt # Install google-java-format -curl -L --retry 3 https://github.com/google/google-java-format/releases/download/v1.23.0/google-java-format_linux-x86-64 > /opt/google-java-format +curl -L --retry 3 --retry-all-errors https://github.com/google/google-java-format/releases/download/v1.23.0/google-java-format_linux-x86-64 > /opt/google-java-format chmod +x /opt/google-java-format diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index 1f75d850e84..9adea394993 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -195,9 +195,17 @@ case "$HF_MODEL" in PREPROCESSOR_FEATURE_SIZE="" PREPROCESSOR_OUTPUT="" ;; + SocialLocalMobile/gemma-4-31B-it-HQQ-INT4) + MODEL_NAME="gemma4_31b" + TASK="" + MAX_SEQ_LEN="" + EXTRA_PIP="" + PREPROCESSOR_FEATURE_SIZE="" + PREPROCESSOR_OUTPUT="" + ;; *) echo "Error: Unsupported model '$HF_MODEL'" - echo "Supported models: mistralai/Voxtral-Mini-3B-2507, mistralai/Voxtral-Mini-4B-Realtime-2602, openai/whisper-{small, medium, large, large-v2, large-v3, large-v3-turbo}, google/gemma-3-4b-it, Qwen/Qwen3-0.6B, nvidia/diar_streaming_sortformer_4spk-v2, nvidia/parakeet-tdt, facebook/dinov2-small-imagenet1k-1-layer, SocialLocalMobile/Qwen3.5-35B-A3B-HQQ-INT4" + echo "Supported models: mistralai/Voxtral-Mini-3B-2507, mistralai/Voxtral-Mini-4B-Realtime-2602, openai/whisper-{small, medium, large, large-v2, large-v3, large-v3-turbo}, google/gemma-3-4b-it, Qwen/Qwen3-0.6B, nvidia/diar_streaming_sortformer_4spk-v2, nvidia/parakeet-tdt, facebook/dinov2-small-imagenet1k-1-layer, SocialLocalMobile/Qwen3.5-35B-A3B-HQQ-INT4, SocialLocalMobile/gemma-4-31B-it-HQQ-INT4" exit 1 ;; esac @@ -459,6 +467,50 @@ if [ "$MODEL_NAME" = "qwen3_5_moe" ]; then exit 0 fi +# Gemma 4 31B uses a prequantized checkpoint and custom export script +if [ "$MODEL_NAME" = "gemma4_31b" ]; then + pip install safetensors huggingface_hub gguf + + # Download prequantized model outside OUTPUT_DIR to avoid uploading on failure + LOCAL_MODEL_DIR=$(mktemp -d) + INDUCTOR_CACHE=$(mktemp -d) + trap 'rm -rf "$LOCAL_MODEL_DIR" "$INDUCTOR_CACHE"' EXIT + + python -c "from huggingface_hub import snapshot_download; snapshot_download('${HF_MODEL}', local_dir='${LOCAL_MODEL_DIR}')" + + # Sanity check: run inference on the prequantized model + echo "::group::Inference sanity check" + INFERENCE_OUTPUT=$(python -m executorch.examples.models.gemma4_31b.inference \ + --prequantized "$LOCAL_MODEL_DIR" \ + --prompt "What is the capital of France?" \ + --max-new-tokens 32 \ + --temperature 0 \ + --no-compile 2>&1) + echo "$INFERENCE_OUTPUT" + if ! echo "$INFERENCE_OUTPUT" | grep -q "Paris"; then + echo "ERROR: Inference sanity check failed — expected 'Paris' in output" + exit 1 + fi + echo "::endgroup::" + + # Copy tokenizer for the runner + cp "$LOCAL_MODEL_DIR/tokenizer.json" "${OUTPUT_DIR}/tokenizer.json" + + # Export to .pte/.ptd (short cache dir avoids objcopy symbol length issues) + echo "::group::Export" + TORCHINDUCTOR_CACHE_DIR="$INDUCTOR_CACHE" \ + python -m executorch.examples.models.gemma4_31b.export \ + --prequantized "$LOCAL_MODEL_DIR" \ + --output-dir "${OUTPUT_DIR}" + echo "::endgroup::" + + test -f "${OUTPUT_DIR}/model.pte" + test -f "${OUTPUT_DIR}/aoti_cuda_blob.ptd" + ls -al "${OUTPUT_DIR}" + + exit 0 +fi + MAX_SEQ_LEN_ARG="" if [ -n "$MAX_SEQ_LEN" ]; then MAX_SEQ_LEN_ARG="--max_seq_len $MAX_SEQ_LEN" diff --git a/.ci/scripts/setup-emscripten.sh b/.ci/scripts/setup-emscripten.sh index a4f4fd1a078..66477f24a7c 100644 --- a/.ci/scripts/setup-emscripten.sh +++ b/.ci/scripts/setup-emscripten.sh @@ -9,7 +9,7 @@ set -ex # need version >= 17 install_node() { - curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.3/install.sh | bash + curl --retry 3 --retry-all-errors -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.3/install.sh | bash source "$HOME/.nvm/nvm.sh" nvm install 22 } diff --git a/.ci/scripts/setup-macos.sh b/.ci/scripts/setup-macos.sh index 6bd26e0b171..38863597224 100755 --- a/.ci/scripts/setup-macos.sh +++ b/.ci/scripts/setup-macos.sh @@ -34,7 +34,7 @@ install_buck() { # team for help. BUCK2_VERSION=$(cat ci_commit_pins/buck2.txt) BUCK2=buck2-aarch64-apple-darwin-${BUCK2_VERSION}.zst - curl -s "https://ossci-macos.s3.amazonaws.com/${BUCK2}" -o "${BUCK2}" + curl -s --retry 3 --retry-all-errors "https://ossci-macos.s3.amazonaws.com/${BUCK2}" -o "${BUCK2}" zstd -d "${BUCK2}" -o buck2 @@ -68,7 +68,7 @@ install_sccache() { # NB: The function is adopted from PyTorch MacOS build workflow # https://github.com/pytorch/pytorch/blob/main/.github/workflows/_mac-build.yml if ! command -v sccache &> /dev/null; then - sudo curl --retry 3 "https://s3.amazonaws.com/ossci-macos/sccache/sccache-v0.4.1-${RUNNER_ARCH}" --output "${SCCACHE_PATH}/sccache" + sudo curl --retry 3 --retry-all-errors "https://s3.amazonaws.com/ossci-macos/sccache/sccache-v0.4.1-${RUNNER_ARCH}" --output "${SCCACHE_PATH}/sccache" sudo chmod +x "${SCCACHE_PATH}/sccache" fi diff --git a/.ci/scripts/setup-mediatek-deps.sh b/.ci/scripts/setup-mediatek-deps.sh index f93a319e11a..d24bcfaa9ef 100755 --- a/.ci/scripts/setup-mediatek-deps.sh +++ b/.ci/scripts/setup-mediatek-deps.sh @@ -14,7 +14,7 @@ install_neuropilot() { echo "Start installing neuropilot." mkdir -p "${MEDIATEK_INSTALLATION_DIR}" - curl -Lo /tmp/neuropilot-express.tar.gz "https://s3.ap-southeast-1.amazonaws.com/mediatek.neuropilot.com/06302508-4c94-4bf2-9789-b0ee44e83e27.gz" + curl -Lo /tmp/neuropilot-express.tar.gz --retry 3 --retry-all-errors "https://s3.ap-southeast-1.amazonaws.com/mediatek.neuropilot.com/06302508-4c94-4bf2-9789-b0ee44e83e27.gz" echo "Finishing downloading neuropilot sdk." tar zxvf /tmp/neuropilot-express.tar.gz --strip-components=1 --directory "${MEDIATEK_INSTALLATION_DIR}" echo "Finishing unzip neuropilot sdk." @@ -33,7 +33,7 @@ setup_neuropilot() { } setup_calibration_data() { - curl -Lo /tmp/imagenette2-160.tgz https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-160.tgz + curl -Lo /tmp/imagenette2-160.tgz --retry 3 --retry-all-errors https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-160.tgz tar zxvf /tmp/imagenette2-160.tgz --strip-components=1 --directory "${MEDIATEK_INSTALLATION_DIR}" } diff --git a/.ci/scripts/setup-openvino.sh b/.ci/scripts/setup-openvino.sh index 4f3de081536..6acc691649b 100755 --- a/.ci/scripts/setup-openvino.sh +++ b/.ci/scripts/setup-openvino.sh @@ -37,7 +37,7 @@ else echo "Using OpenVINO stable release: ${OPENVINO_BUILD}" fi -curl -Lo /tmp/openvino_toolkit.tgz --retry 3 --fail ${OPENVINO_URL} +curl -Lo /tmp/openvino_toolkit.tgz --retry 3 --retry-all-errors --fail ${OPENVINO_URL} tar -xzf /tmp/openvino_toolkit.tgz mv "${OPENVINO_EXTRACTED_DIR}" openvino diff --git a/.ci/scripts/setup-samsung-linux-deps.sh b/.ci/scripts/setup-samsung-linux-deps.sh index 9aa9c4380a5..5af62ec35ca 100644 --- a/.ci/scripts/setup-samsung-linux-deps.sh +++ b/.ci/scripts/setup-samsung-linux-deps.sh @@ -43,7 +43,7 @@ download_and_extract() { local out_file="$3" echo "Downloading from ${download_url}..." - curl -fsSL --retry 3 \ + curl -fsSL --retry 3 --retry-all-errors \ -H "apikey: ${API_KEY}" \ -o "${out_file}" \ "${download_url}" diff --git a/.ci/scripts/setup-vulkan-linux-deps.sh b/.ci/scripts/setup-vulkan-linux-deps.sh index cd99ff0d6ff..a0dcb75ad4a 100755 --- a/.ci/scripts/setup-vulkan-linux-deps.sh +++ b/.ci/scripts/setup-vulkan-linux-deps.sh @@ -16,7 +16,7 @@ install_swiftshader() { _tmp_archive="/tmp/${_swiftshader_archive}" - curl --silent --show-error --location --fail --retry 3 \ + curl --silent --show-error --location --fail --retry 3 --retry-all-errors \ --output "${_tmp_archive}" "$_https_amazon_aws/${_swiftshader_archive}" tar -C "${_swiftshader_dir}" -xzf "${_tmp_archive}" @@ -35,7 +35,7 @@ install_vulkan_sdk() { _tmp_archive="/tmp/vulkansdk.tar.gz" - curl --silent --show-error --location --fail --retry 3 \ + curl --silent --show-error --location --fail --retry 3 --retry-all-errors \ --output "${_tmp_archive}" "${_vulkan_sdk_url}" tar -C "${_vulkan_sdk_dir}" -xJf "${_tmp_archive}" diff --git a/.ci/scripts/test_cortex_m_e2e.sh b/.ci/scripts/test_cortex_m_e2e.sh index de47a45ea0d..ac6e6d46550 100755 --- a/.ci/scripts/test_cortex_m_e2e.sh +++ b/.ci/scripts/test_cortex_m_e2e.sh @@ -19,6 +19,7 @@ et_root_dir=$(realpath "${script_dir}/../..") # Quantization is the default for the cortex-m55 target; run.sh's # arg parser only recognizes --no_quantize, so we omit any explicit flag. +export ARM_FVP_INSTALL_I_AGREE_TO_THE_CONTAINED_EULA=True bash "${et_root_dir}/examples/arm/run.sh" \ --model_name="${MODEL}" \ --target=cortex-m55 \ diff --git a/.ci/scripts/test_ios_ci.sh b/.ci/scripts/test_ios_ci.sh index 46c3f71f021..a961fac0dd3 100755 --- a/.ci/scripts/test_ios_ci.sh +++ b/.ci/scripts/test_ios_ci.sh @@ -55,7 +55,7 @@ mv $MODEL_NAME*.pte "$APP_PATH/Resources/Models/MobileNet/" say "Downloading Labels" -curl https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt \ +curl --retry 3 --retry-all-errors https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt \ -o "$APP_PATH/Resources/Models/MobileNet/imagenet_classes.txt" say "Creating Simulator" diff --git a/.ci/scripts/test_model_e2e.sh b/.ci/scripts/test_model_e2e.sh index 1678b0a4fbb..e1ba976b0cc 100755 --- a/.ci/scripts/test_model_e2e.sh +++ b/.ci/scripts/test_model_e2e.sh @@ -228,9 +228,21 @@ case "$HF_MODEL" in AUDIO_FILE="" IMAGE_PATH="" ;; + SocialLocalMobile/gemma-4-31B-it-HQQ-INT4) + MODEL_NAME="gemma4_31b" + RUNNER_TARGET="gemma4_31b_runner" + RUNNER_PATH="gemma4_31b" + EXPECTED_OUTPUT="Paris" + PREPROCESSOR="" + TOKENIZER_URL="" + TOKENIZER_FILE="tokenizer.json" + AUDIO_URL="" + AUDIO_FILE="" + IMAGE_PATH="" + ;; *) echo "Error: Unsupported model '$HF_MODEL'" - echo "Supported models: mistralai/Voxtral-Mini-3B-2507, mistralai/Voxtral-Mini-4B-Realtime-2602, nvidia/diar_streaming_sortformer_4spk-v2, openai/whisper series (whisper-{small, medium, large, large-v2, large-v3, large-v3-turbo}), google/gemma-3-4b-it, Qwen/Qwen3-0.6B, nvidia/parakeet-tdt, facebook/dinov2-small-imagenet1k-1-layer, SocialLocalMobile/Qwen3.5-35B-A3B-HQQ-INT4" + echo "Supported models: mistralai/Voxtral-Mini-3B-2507, mistralai/Voxtral-Mini-4B-Realtime-2602, nvidia/diar_streaming_sortformer_4spk-v2, openai/whisper series (whisper-{small, medium, large, large-v2, large-v3, large-v3-turbo}), google/gemma-3-4b-it, Qwen/Qwen3-0.6B, nvidia/parakeet-tdt, facebook/dinov2-small-imagenet1k-1-layer, SocialLocalMobile/Qwen3.5-35B-A3B-HQQ-INT4, SocialLocalMobile/gemma-4-31B-it-HQQ-INT4" exit 1 ;; esac @@ -244,19 +256,19 @@ echo "::group::Prepare $MODEL_NAME Artifacts" # Download tokenizer files (skip for models that bundle tokenizer in export or do not use one) -if [ "$MODEL_NAME" != "parakeet" ] && [ "$MODEL_NAME" != "voxtral_realtime" ] && [ "$MODEL_NAME" != "sortformer" ] && [ "$MODEL_NAME" != "dinov2" ] && [ "$MODEL_NAME" != "qwen3_5_moe" ]; then +if [ "$MODEL_NAME" != "parakeet" ] && [ "$MODEL_NAME" != "voxtral_realtime" ] && [ "$MODEL_NAME" != "sortformer" ] && [ "$MODEL_NAME" != "dinov2" ] && [ "$MODEL_NAME" != "qwen3_5_moe" ] && [ "$MODEL_NAME" != "gemma4_31b" ]; then if [ "$TOKENIZER_FILE" != "" ]; then - curl -L $TOKENIZER_URL/$TOKENIZER_FILE -o $MODEL_DIR/$TOKENIZER_FILE + curl -L --retry 3 --retry-all-errors $TOKENIZER_URL/$TOKENIZER_FILE -o $MODEL_DIR/$TOKENIZER_FILE else - curl -L $TOKENIZER_URL/tokenizer.json -o $MODEL_DIR/tokenizer.json - curl -L $TOKENIZER_URL/tokenizer_config.json -o $MODEL_DIR/tokenizer_config.json - curl -L $TOKENIZER_URL/special_tokens_map.json -o $MODEL_DIR/special_tokens_map.json + curl -L --retry 3 --retry-all-errors $TOKENIZER_URL/tokenizer.json -o $MODEL_DIR/tokenizer.json + curl -L --retry 3 --retry-all-errors $TOKENIZER_URL/tokenizer_config.json -o $MODEL_DIR/tokenizer_config.json + curl -L --retry 3 --retry-all-errors $TOKENIZER_URL/special_tokens_map.json -o $MODEL_DIR/special_tokens_map.json fi fi # Download test files if [ "$AUDIO_URL" != "" ]; then - curl -L $AUDIO_URL -o ${MODEL_DIR}/$AUDIO_FILE + curl -L --retry 3 --retry-all-errors $AUDIO_URL -o ${MODEL_DIR}/$AUDIO_FILE elif [[ "$MODEL_NAME" == *whisper* ]] || [ "$MODEL_NAME" = "voxtral_realtime" ]; then if ! command -v ffmpeg >/dev/null; then if [ "$(uname -s)" = "Linux" ] && command -v apt-get >/dev/null; then @@ -278,7 +290,7 @@ fi # Download test image for vision models if [ -n "${IMAGE_URL:-}" ]; then - curl -L "$IMAGE_URL" -o "${MODEL_DIR}/test_image.jpg" + curl -L --retry 3 --retry-all-errors "$IMAGE_URL" -o "${MODEL_DIR}/test_image.jpg" fi ls -al @@ -368,6 +380,9 @@ EOF qwen3_5_moe) RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt 'What is the capital of France?' --max_new_tokens 128 --temperature 0 --cuda_graph" ;; + gemma4_31b) + RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt 'What is the capital of France?' --max_new_tokens 128 --temperature 0 --cuda_graph" + ;; voxtral_realtime) RUNNER_ARGS="--model_path ${MODEL_DIR}/model.pte --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --preprocessor_path ${MODEL_DIR}/$PREPROCESSOR --audio_path ${MODEL_DIR}/$AUDIO_FILE --temperature 0" # Add CUDA data path if present diff --git a/.ci/scripts/test_riscv_qemu.sh b/.ci/scripts/test_riscv_qemu.sh index 27ab57f3b09..2842542aa3a 100755 --- a/.ci/scripts/test_riscv_qemu.sh +++ b/.ci/scripts/test_riscv_qemu.sh @@ -5,7 +5,7 @@ # LICENSE file in the root directory of this source tree. # CI wrapper: install RISC-V cross-compile + qemu-user tooling, then run the -# RISC-V Phase 1 smoke test (export, cross-compile, qemu-user execution) via +# RISC-V smoke test (export, cross-compile, qemu-user execution) via # examples/riscv/run.sh. The bundled-IO comparison and Test_result: PASS # check are done by run.sh. @@ -14,5 +14,50 @@ set -eu script_dir=$(realpath "$(dirname "${BASH_SOURCE[0]}")") et_root_dir=$(realpath "${script_dir}/../..") +model="add" +xnnpack=false +quantize=false +verbose_xnnpack=false +debug_xnnpack=false + +usage() { + cat < Which model to export and run (default: add) + --xnnpack Enable the XNNPACK backend (AOT partitioner + runtime) + --quantize Produce an 8-bit quantized model + --verbose-xnnpack Build XNNPACK with XNN_LOG_LEVEL=4 to log microkernel dispatch + --debug-xnnpack Enable XNNPACK partitioner DEBUG logging and dump the lowered graph + -h, --help Show this help +EOF +} + +for arg in "$@"; do + case $arg in + --model=*) model="${arg#*=}" ;; + --xnnpack) xnnpack=true ;; + --quantize) quantize=true ;; + --debug-xnnpack) debug_xnnpack=true ;; + --verbose-xnnpack) verbose_xnnpack=true ;; + -h|--help) usage; exit 0 ;; + *) echo "Unknown option: $arg" >&2; usage; exit 1 ;; + esac +done + +run_extra_args=() +if ${xnnpack}; then + run_extra_args+=(--xnnpack) +fi +if ${quantize}; then + run_extra_args+=(--quantize) +fi +if ${debug_xnnpack}; then + run_extra_args+=(--debug-xnnpack) +fi +if ${verbose_xnnpack}; then + run_extra_args+=(--verbose-xnnpack) +fi + bash "${et_root_dir}/examples/riscv/setup.sh" -bash "${et_root_dir}/examples/riscv/run.sh" +bash "${et_root_dir}/examples/riscv/run.sh" --model="${model}" "${run_extra_args[@]}" diff --git a/.ci/scripts/test_zephyr.sh b/.ci/scripts/test_zephyr.sh new file mode 100755 index 00000000000..6ce8e082b84 --- /dev/null +++ b/.ci/scripts/test_zephyr.sh @@ -0,0 +1,178 @@ +#!/usr/bin/env bash +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +set -euo pipefail + +SCRIPT_DIR="$(realpath "$(dirname "${BASH_SOURCE[0]}")")" +EXECUTORCH_PROJ_ROOT="$(realpath "${SCRIPT_DIR}/../..")" +ZEPHYR_README_PATH="zephyr/README.md" + +ZEPHYR_SAMPLES_README_PATH="zephyr/samples/hello-executorch/README.md" +TARGETS_ARG="${TARGET_LIST:-}" + +usage() { + cat < README containing test_* command blocks + --targets Comma-separated target list, e.g. ethos-u55,cortex-m55,ethos-u85 + -h, --help Show this help +EOF +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --targets) + TARGETS_ARG="$2" + shift 2 + ;; + --zephyr-samples-readme-path) + ZEPHYR_SAMPLES_README_PATH="$2" + shift 2 + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "ERROR: Unknown argument: $1" >&2 + usage >&2 + exit 1 + ;; + esac +done + +if [[ -z "${TARGETS_ARG}" ]]; then + echo "ERROR: --targets or TARGET_LIST must be set" >&2 + usage >&2 + exit 1 +fi + +IFS=',' read -r -a TARGETS <<< "${TARGETS_ARG}" + +export EXECUTORCH_PROJ_ROOT + +cd "${EXECUTORCH_PROJ_ROOT}" + +# Source utility scripts. +. .ci/scripts/utils.sh +. .ci/scripts/zephyr-utils.sh + +run_target_test_blocks_from_readme() { + local readme_path="$1" + local target="$2" + local resolved_readme_path marker markers + + resolved_readme_path="$(_utils_path_from_root "${readme_path}")" + markers="$(awk -v target="${target}" ' + { + line = $0 + while (match(line, //)) { + marker = substr(line, RSTART, RLENGTH) + if (index(marker, " in ${readme_path}" >&2 + return 1 + fi + + while IFS= read -r marker; do + echo "---- ${target} ${marker} ----" + run_command_block_from_readme "${readme_path}" "${marker}" + done <<< "${markers}" +} + +# Check that zephyr/README.md and zephyr/executorch.yaml are in sync. +verify_zephyr_readme + +# Based on instructions in zephyr/README.md and the selected sample README. +run_command_block_from_readme "${ZEPHYR_README_PATH}" "" + +# Make sure to backup the zephyr_scratch folder if it exists to allow for local +# testing that does not lose code/data. +if [[ -d "zephyr_scratch" ]]; then + mv "zephyr_scratch" "zephyr_scratch.backup.$(date +%Y%m%d%H%M%S)" +fi +mkdir -p zephyr_scratch/ + +cd zephyr_scratch +export ZEPHYR_PROJ_ROOT="$(realpath "$(pwd)")" + +echo "---- Zephyr SDK ----" +# Use ZephyrSDK if on the disk (e.g. setup in the docker) +# Check for a zephyr-sdk-0.17.4 directory and make a symlink if found in parent directories +if sdk_dir=$(find ../../.. -maxdepth 4 -type d -name 'zephyr-sdk-0.17.4' -print -quit) && [ -n "${sdk_dir}" ]; then + echo "---- Found pre downloaded Zephyr SDK in ${sdk_dir} ----" + ln -s "${sdk_dir}" . +fi + +# Download and setup Zephyr SDK 0.17.4 if not already present +if [ ! -d "zephyr-sdk-0.17.4" ]; then + echo "---- Downloading Zephyr SDK ----" + wget https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v0.17.4/zephyr-sdk-0.17.4_linux-x86_64.tar.xz + tar -xf zephyr-sdk-0.17.4_linux-x86_64.tar.xz + rm -f zephyr-sdk-0.17.4_linux-x86_64.tar.xz* +fi + +./zephyr-sdk-0.17.4/setup.sh -c -t arm-zephyr-eabi +export ZEPHYR_SDK_INSTALL_DIR=$(realpath ./zephyr-sdk-0.17.4) + +cd ${ZEPHYR_PROJ_ROOT} + +run_command_block_from_readme "${ZEPHYR_README_PATH}" "" + +cp "${EXECUTORCH_PROJ_ROOT}/zephyr/executorch.yaml" zephyr/submanifests/ + +run_command_block_from_readme "${ZEPHYR_README_PATH}" "" + +# Switch to executorch in this PR e.g. replace modules/lib/executorch with the +# root folder of this repo instead of doing a re-checkout and figuring out the +# correct commit hash. +rm -Rf modules/lib/executorch +ln -s "${EXECUTORCH_PROJ_ROOT}" modules/lib/executorch + +# Setup git local user for Executorch git to allow +# modules/lib/executorch/examples/arm/setup.sh to run inside CI later. +if ! git config --get user.name >/dev/null 2>&1; then + git config --global user.name "Github Executorch" +fi +if ! git config --get user.email >/dev/null 2>&1; then + git config --global user.email "github_executorch@arm.com" +fi + +run_command_block_from_readme "${ZEPHYR_README_PATH}" "" +run_command_block_from_readme "${ZEPHYR_README_PATH}" "" + +for TARGET in "${TARGETS[@]}"; do + TARGET="$(echo "${TARGET}" | xargs)" + + echo "---- ${TARGET} ----" + rm -Rf build + + if [[ ${TARGET} == "ethos-u55" || ${TARGET} == "cortex-m55" ]]; then + BOARD="corstone300" + elif [[ ${TARGET} == "ethos-u85" ]]; then + BOARD="corstone320" + else + echo "Fail unsupported target selection ${TARGET}" + exit 1 + fi + + echo "---- ${TARGET} Board ${BOARD} FVP setup ----" + run_command_block_from_readme "${ZEPHYR_SAMPLES_README_PATH}" "" + + # Run all blocks that match + run_target_test_blocks_from_readme "${ZEPHYR_SAMPLES_README_PATH}" "${TARGET}" +done diff --git a/.ci/scripts/utils.sh b/.ci/scripts/utils.sh index b291374d667..486745f4bf6 100644 --- a/.ci/scripts/utils.sh +++ b/.ci/scripts/utils.sh @@ -242,8 +242,8 @@ cmake_install_executorch_lib() { download_stories_model_artifacts() { # Download stories110M.pt and tokenizer from Github - curl -Ls "https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.pt" --output stories110M.pt - curl -Ls "https://raw.githubusercontent.com/karpathy/llama2.c/master/tokenizer.model" --output tokenizer.model + curl -Ls --retry 3 --retry-all-errors "https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.pt" --output stories110M.pt + curl -Ls --retry 3 --retry-all-errors "https://raw.githubusercontent.com/karpathy/llama2.c/master/tokenizer.model" --output tokenizer.model # Create params.json file touch params.json echo '{"dim": 768, "multiple_of": 32, "n_heads": 12, "n_layers": 12, "norm_eps": 1e-05, "vocab_size": 32000}' > params.json diff --git a/.claude/skills/qualcomm/new_op_development.md b/.claude/skills/qualcomm/new_op_development.md index dc639655257..6e1abcf77f6 100644 --- a/.claude/skills/qualcomm/new_op_development.md +++ b/.claude/skills/qualcomm/new_op_development.md @@ -1,361 +1,258 @@ -# New Op Development +# New Op Development — QNN/HTP Backend -Full reference: `backends/qualcomm/builders/README.md` (op builder) and `backends/qualcomm/quantizer/README.md` (quantizer annotation). +## Decision Tree -## Overview - -Adding a new op requires three steps: -1. Implement the op builder (`builders/op_*.py`) -2. Register quantizer annotation (`quantizer/annotators/`) -3. Add unit tests (`tests/`) - -**Important**: If the torch op requires **multiple QNN ops** to implement (e.g., no direct QNN equivalent), use a **decompose pass** instead of creating multiple ops in a single builder. Skip Steps 3–6 and follow the **Decompose Pass Approach** section at the bottom of this file. +1. **QNN has a native op?** → Native builder approach (Steps 1–8) +2. **No native op, needs multiple QNN ops?** → Decompose pass approach --- ## Step 1: Identify the Unsupported Op -Run the model through the QNN backend. A missing op surfaces as: - -``` -KeyError: 'aten.native_layer_norm.default' -``` - -To trace back to the source PyTorch layer: - -```python -from executorch.backends.qualcomm.utils.utils import capture_program - -prog = capture_program(MyModel(), example_inputs) -for node in prog.exported_program.graph.nodes: - if node.op == "call_function" and node.target.__name__ == 'aten.native_layer_norm.default': - print(node.meta["source_fn_stack"]) -``` - ---- +Missing ops surface as `KeyError: 'aten.my_op.default'` when running through QNN backend. ## Step 2: Check Operator Spec -- **QNN side**: [Operator Definitions](https://docs.qualcomm.com/bundle/publicresource/topics/80-63442-10/MasterOpDef.html) — check IO order, required vs optional tensors, parameter names and shapes -- **PyTorch side**: [ATen Operator Definitions](https://github.com/pytorch/pytorch/tree/main/aten/src/ATen/native) — map PyTorch args to QNN IO/params -- **Fallback search**: [Supported Ops table](https://docs.qualcomm.com/bundle/publicresource/topics/80-63442-10/SupportedOps.html) -- **Header reference**: `$QNN_SDK_ROOT/include/QNN/QnnOpDef.h` — authoritative string literals +- [Master Op Definitions](https://docs.qualcomm.com/bundle/publicresource/topics/80-63442-10/MasterOpDef.html) — IO order, params, shapes +- [HTP Op Def Supplement](https://docs.qualcomm.com/doc/80-63442-10/topic/HtpOpDefSupplement.html) — HTP-specific constraints & supported dtypes +- [Supported Ops table](https://docs.qualcomm.com/bundle/publicresource/topics/80-63442-10/SupportedOps.html) +- `$QNN_SDK_ROOT/include/QNN/QnnOpDef.h` — authoritative string literals +- [ATen native ops](https://github.com/pytorch/pytorch/tree/main/aten/src/ATen/native) — PyTorch arg mapping ---- +**⚠️ Caveats:** +- An op in the Master def may **not exist** in the HTP supplement → not available on HTP +- HTP docs may claim a dtype is supported but **fail at runtime** → always test on-device -## Step 3: Add Op Constant - -In `builders/qnn_constants.py`, add a dataclass (alphabetical order): +## Step 3: Add Op Constant (`builders/qnn_constants.py`) ```python @dataclass(init=False, frozen=True) -class OpLayerNorm: - op_name: str = "LayerNorm" - param_epsilon = "epsilon" - param_axes = "axes" +class OpMyOp: + op_name: str = "MyOp" # Must match QnnOpDef.h exactly + param_axis: str = "axis" + param_epsilon: str = "epsilon" ``` -String values must exactly match `QnnOpDef.h`. - ---- - -## Step 4: Implement the Builder - -Create `builders/op_layer_norm.py`: +## Step 4: Implement Builder (`builders/op_my_op.py`) ```python -import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManager -import numpy as np -import torch -from executorch.backends.qualcomm.utils.constants import QCOM_DATA -from .node_visitor import NodeVisitor -from .node_visitor_manager import register_node_visitor -from .qnn_constants import OpLayerNorm, QNN_OP_PACKAGE_NAME_QTI_AISW -from .utils import get_parameter - @register_node_visitor -class LayerNormVisitor(NodeVisitor): - target = ["aten.native_layer_norm.default"] - - def __init__(self, *args) -> None: - super().__init__(*args) +class MyOpVisitor(NodeVisitor): + target = ["aten.my_op.default"] # Must be a list def define_node(self, node, nodes_to_wrappers): - # 1. Input activation - input_node = node.args[0] + input_node = self.get_node(node.args[0]) input_tensor = self.get_tensor(input_node, node) - input_tensor_wrapper = self.define_tensor( - input_node, node, input_tensor, - PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, - nodes_to_wrappers, - ) - - # 2. Weight (gamma) and bias (beta) — STATIC tensors - weight_node = self.get_node(node.args[2]) - weight_tensor = get_parameter(weight_node, self.edge_program) - weight_tensor_wrapper = self.define_tensor( - weight_node, node, weight_tensor, - PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC, - nodes_to_wrappers, - ) - - bias_node = self.get_node(node.args[3]) - bias_tensor = get_parameter(bias_node, self.edge_program) - bias_tensor_wrapper = self.define_tensor( - bias_node, node, bias_tensor, - PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC, - nodes_to_wrappers, - ) - - # 3. Parameters - normalized_shapes = node.args[1] - if len(normalized_shapes) != 1: - print("QNN only supports normalized output with rank 1") - return - axes = [len(input_tensor.shape) - 1] - axes_shape = [len(axes)] - epsilon = node.args[4] - - # 4. Output - output_tensor = self.get_tensor(node, node, 0) - output_tensor_wrapper = self.define_tensor( - node, node, output_tensor, - PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, - nodes_to_wrappers, - ) - - # 5. Build op - layer_norm_op = PyQnnManager.PyQnnOpWrapper( - node.name, QNN_OP_PACKAGE_NAME_QTI_AISW, OpLayerNorm.op_name, - ) - layer_norm_op.AddInputTensors( - [input_tensor_wrapper, weight_tensor_wrapper, bias_tensor_wrapper] - ) - layer_norm_op.AddOutputTensors([output_tensor_wrapper]) - layer_norm_op.AddScalarParam( - OpLayerNorm.param_epsilon, - PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_FLOAT_32, - {QCOM_DATA: np.float32(epsilon)}, - ) - layer_norm_op.AddTensorParam( - OpLayerNorm.param_axes, - PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_UINT_32, - len(axes_shape), axes_shape, - np.array(axes, dtype=np.uint32), - True, - ) - return layer_norm_op + input_wrapper = self.define_tensor(input_node, node, input_tensor, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, nodes_to_wrappers) + + output_tensor = self.get_tensor(node, node) + output_wrapper = self.define_tensor(node, node, output_tensor, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, nodes_to_wrappers) + + op = PyQnnManager.PyQnnOpWrapper(node.name, QNN_OP_PACKAGE_NAME_QTI_AISW, OpMyOp.op_name) + op.AddInputTensors([input_wrapper]) + op.AddOutputTensors([output_wrapper]) + op.AddScalarParam(OpMyOp.param_axis, PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_UINT_32, + {QCOM_DATA: np.uint32(axis)}) + return op # Return None → op falls back to CPU ``` -Key notes: -- `target` must be a list (multiple targets can share one visitor) -- Use `QNN_TENSOR_TYPE_NATIVE` for activations, `QNN_TENSOR_TYPE_STATIC` for weights/biases -- `define_tensor` handles `APP_READ`/`APP_WRITE` detection internally — always pass `NATIVE` -- `wrapper_idx` needed when node output is a tuple (e.g. split ops) -- Return `None` to signal validation failure → op falls back to CPU +**Key patterns:** +- `QNN_TENSOR_TYPE_NATIVE` for activations, `QNN_TENSOR_TYPE_STATIC` for weights/params +- `wrapper_idx=i` for multi-output ops (tuples); companion `getitem` skip op handles indexing +- Negative dims: `dim = dim % len(shape)` (QNN requires positive axes) +- Axis remapping: `if QCOM_AXIS_ORDER in node.meta: dim = node.meta[QCOM_AXIS_ORDER].index(dim)` +- Static params: `weight = get_parameter(self.get_node(node.args[1]), self.edge_program)` +- Scalar params → `AddScalarParam`; Array params → `AddTensorParam` +- Data types: axis/dims=`UINT_32`, epsilon=`FLOAT_32`, booleans=`BOOL_8` +- Int64 index tensors: use `.to(torch.int32)` in builder + add op to `I64_IN_OPS` in `_passes/i64_to_i32.py` for CPU fallback safety (see `op_gather.py` pattern) ---- +## Step 5: Register Builder (`builders/__init__.py`) + +Add `op_my_op` to both `from . import (...)` and `__all__ = [...]` (alphabetical). -## Step 5: Register the Builder +## Step 6: Add Quantizer Annotation -In `builders/__init__.py` (alphabetical order): +Add to BOTH `quantizer/annotators/htp_rules.py` AND `quantizer/annotators/lpai_rules.py`: ```python -from . import ( - ... - op_layer_norm, - ... -) -__all__ = [..., op_layer_norm, ...] +@register_annotator([torch.ops.aten.my_op.default], QnnConstants.OpMyOp.op_name) +class MyOp(GeneralOpDef): + pass # Default: annotate_single_in_single_out ``` ---- - -## Step 6: Add Quantizer Annotation +**Annotation function selection:** -In `quantizer/annotators/{backend}_rules.py`: +| Op type | Function | When | +|---------|----------|------| +| Compute (new scale) | `annotate_single_in_single_out` | Default — most ops | +| Pass-through (`is_math_invariant`) | `annotate_in_out_obs_sharing_op` + fallback `annotate_single_in_share_out` | Reshape, Permute, Squeeze, Gather | +| Two data inputs (same quant) | Custom `annotate` with `SharedQuantizationSpec` | Scatter, where both data+src need same spec | +| Two inputs | `annotate_binary` | Add, Mul, Sub | +| Conv/Linear (weight+bias) | `annotate_conv` | Convolution, Linear | +| Skip (no QNN mapping) | `qnn_op=None` | getitem, index_copy | +**Custom multi-input annotator** (e.g., scatter where args[0] and args[3] are both data tensors): ```python -@register_annotator( - [torch.ops.aten.native_layer_norm.default], - QnnConstants.OpLayerNorm.op_name, -) -class LayerNormAnnotator(GeneralOpDef): +@register_annotator([torch.ops.aten.scatter.src], qnn_op=None) +class ScatterElements(GeneralOpDef): @staticmethod def annotate(node, quantization_config): - annotate_single_in_single_out(node, quantization_config) + if _is_annotated([node]): return + input_qspec_map = {} + input_act = node.args[0] + input_qspec_map[input_act] = quantization_config.input_activation + if isinstance(node.args[3], Node) and _is_float_tensor(node.args[3]): + input_qspec_map[node.args[3]] = SharedQuantizationSpec((input_act, node)) + output_qspec = SharedQuantizationSpec((input_act, node)) if _is_float_tensor(node) else None + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, output_qspec=output_qspec, _annotated=True) ``` -- Use `qnn_op=None` for skip ops (e.g. `operator.getitem`) -- `annotate_single_in_single_out` covers most cases; implement custom logic for multi-input ops - -Full annotation tutorial: `backends/qualcomm/quantizer/README.md` - -### Choosing the right annotate function - -The QNN backend validates quantization constraints via `backend_opinfo` (QNN SDK ≥ 2.41). If validation fails with: - -``` -ValueError: Validation failed for node with target aten..default -``` - -Check the warning log above it — it will say which constraint failed. The most common case is `is_math_invariant=True`, which means the op does not change values (only rearranges data), so input and output **must share the same quantization parameters**. - -| Op type | annotate function | Example ops | -|---------|-------------------|-------------| -| General (input → output with new scale) | `annotate_single_in_single_out` | LayerNorm, Conv2d | -| Pass-through (rearranges data only) | `annotate_in_out_obs_sharing_op` + fallback | Reshape, ChannelShuffle, PixelShuffle | -| Multi-input | `annotate_binary` | Add, Mul | - -For **pass-through ops** (reshape, shuffle, permute — ops where `is_math_invariant=True`), override `annotate` like this: +## Step 7: Add Layout Transform Registration (`_passes/layout_transform.py`) +Add op to `layout_agnostic_ops` (most ops) or `layout_sensitive_ops` (conv, pool, etc.): ```python -@register_annotator( - [torch.ops.aten.channel_shuffle.default], QnnConstants.OpChannelShuffle.op_name -) -class ChannelShuffle(GeneralOpDef): - @staticmethod - def annotate(node: Node, quantization_config: QuantizationConfig) -> None: - annotate_in_out_obs_sharing_op(node, quantization_config) - if not _is_annotated([node]): - annotate_single_in_share_out(node, quantization_config) +exir_ops.edge.aten.my_op.default, ``` -`annotate_in_out_obs_sharing_op` shares the input's observer with the output (satisfies `is_math_invariant`). The fallback `annotate_single_in_share_out` handles the case where the input node is not yet annotated. - ---- - -## Step 7: Add Unit Tests - -In `tests/models.py` (alphabetical order): +## Step 8: Add Unit Tests +**Model** in `tests/models.py` (alphabetical, parameterize variants): ```python -class LayerNorm(torch.nn.Module): - def __init__(self): +class MyOp(torch.nn.Module): + def __init__(self, param=0): super().__init__() - self.layer_norm = torch.nn.LayerNorm([768], eps=1e-6) - + self.param = param def forward(self, x): - return self.layer_norm(x) + return torch.my_op(x, self.param) ``` -In `tests/test_qnn_delegate.py`, add to both `TestQNNFloatingPointOperator` and `TestQNNQuantizedOperator` (alphabetical order): - +**Tests** in `tests/test_qnn_delegate.py` — add to BOTH `TestQNNFloatingPointOperator` and `TestQNNQuantizedOperator`: ```python -def test_qnn_backend_layer_norm(self): - module = LayerNorm() - sample_input = (torch.randn(196, 768),) - module = self.get_qdq_module(module, sample_input) # quantized only - self.lower_module_and_test_output(module, sample_input) +def test_qnn_backend_my_op(self): + test_comb = [{ + QCOM_MODULE: [MyOp(), MyOp(param=1)], + QCOM_SAMPLE_INPUTS: [(torch.randn(3, 4),), (torch.randn(3, 4, dtype=torch.float16),)], + }] + index = 0 + for comb in test_comb: + for module in comb[QCOM_MODULE]: + for sample_input in comb[QCOM_SAMPLE_INPUTS]: + with self.subTest(i=index): + index += 1 + self.lower_module_and_test_output(module, sample_input) ``` -Expected result: 1 delegated node, only placeholders/output nodes remain outside the delegate. - ---- - -## Step 8: Prevent Decomposition (if needed) - -Some torch ops are in ExecuTorch's default decomposition table and will be broken into primitives **before** the QNN partitioner sees them. If QNN has a native op for it, you must explicitly skip decomposition. - -**Check first** with a quick Python snippet (run from the executorch root with the `executorch` conda env active): - +**Quantized test** — use separate variable to avoid overwriting module: ```python -import torch -from executorch.exir.tracer import _default_decomposition_table - -decomp_table = _default_decomposition_table() -op = torch.ops.aten.channel_shuffle.default -print(op in decomp_table) # True → will be decomposed +qdq_module = self.get_qdq_module(module, sample_input) +self.lower_module_and_test_output(qdq_module, sample_input) ``` -Output: -``` -True # in ExecuTorch decomp table +**Test data rules:** +- No duplicate indices for scatter/gather with `reduction=NONE` +- Deterministic inputs for precision-sensitive decompositions (avoid boundary values) +- Bounded inputs for ops with singularities (tan, reciprocal): `torch.rand() * 2 - 1` + +**Run on-device:** +```bash +python backends/qualcomm/tests/test_qnn_delegate.py \ + -k TestQNNFloatingPointOperator.test_qnn_backend_my_op \ + --model SM8750 --host --device --build_folder build-android ``` -If `True`, add the op to `get_skip_decomp_table()` in `partition/utils.py` (alphabetical order): +Always ask user for `--model`, `--host`, `--device`, `--build_folder` values. -```python -def get_skip_decomp_table() -> List[torch._ops.OperatorBase]: - do_not_decompose = [ - torch.ops.aten.adaptive_avg_pool2d.default, - torch.ops.aten.channel_shuffle.default, # ← add here - torch.ops.aten.col2im.default, - ... - ] -``` +## Step 9: Prevent Decomposition (if needed) -**Verification**: After adding, re-run the tests. The partitioner log should show: +If the ATen op exists in ExecuTorch's decomp table and you have a builder for it: +- Add to `partition/utils.py` → `get_skip_decomp_table()` +- Remove from `partition/common_defs.py` → `to_be_implemented_operator` if listed there -``` -[QNN Partitioner Op Support]: aten.channel_shuffle.default | True -``` +## Step 10: Update Documentation -If the op was decomposed (not in skip table), the partitioner would never see `aten.channel_shuffle.default` and the test would still pass but via decomposed primitives — not the native QNN op. +- `builders/README.md` — Update QNN ops table (✗ → ✓) and add to "Additional Operators Supported via Passes" table if using decomposition --- -## Decompose Pass Approach (for ops without direct QNN equivalent) +## Decompose Pass Approach -When a torch op has **no direct QNN equivalent** and requires multiple QNN ops to implement, use a **decompose pass** to rewrite the graph into primitive ops that QNN already supports. This is preferred over creating multiple ops in a single builder. +Use when QNN has **no native op** — decompose into supported primitives. -**Reference**: `backends/qualcomm/_passes/decompose_linalg_vector_norm.py` +### Approach A: Module Export +**Ref:** `_passes/decompose_linalg_vector_norm.py`. Write a `torch.nn.Module`, export, merge graph via `merge_decomposed_graph`. Simple but may produce unexpected ops. -### Pattern +### Approach B: Direct Graph Manipulation (RECOMMENDED) +**Ref:** `_passes/decompose_remainder.py`, `_passes/decompose_log_variants.py`. ```python -# 1. Define a torch.nn.Module that implements the op using supported primitives -class MyOpDecomposed(torch.nn.Module): - def __init__(self, param): +class DecomposeMyOp(ExportPass): + def __init__(self): super().__init__() - self.param = param - - def forward(self, x): - # Use only ops that QNN supports - return torch.some_supported_op(x, self.param) + self.targets = {torch.ops.aten.my_op.default, exir_ops.edge.aten.my_op.default} - -# 2. Create the ExportPass -class DecomposeMyOp(ExportPass): - def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + def call(self, graph_module): graph = graph_module.graph + const_cache = {} for node in list(graph.nodes): - if node.target == torch.ops.aten.my_op.default: - param = node.args[1] # extract params from node - model = MyOpDecomposed(param) - ep = torch.export.export(model, (node.args[0].meta["val"],), strict=True) - decomposed_module = ep.run_decompositions().graph_module - + if node.op == "call_function" and node.target in self.targets: + is_edge = isinstance(node.target, EdgeOpOverload) + op = exir_ops.edge.aten.div.Tensor if is_edge else torch.ops.aten.div.Tensor with graph.inserting_before(node): - remap = {"x": node.args[0]} - merge_decomposed_graph( - remap=remap, - target_node=node, - target_graph=graph, - decomposed_graph_module=decomposed_module, - ) - graph.erase_node(node) - + new_node = graph.create_node("call_function", op, (node.args[0],)) + new_node.meta = copy_meta(node.meta) + for user in node.users.copy(): + user.replace_input_with(node, new_node) graph.eliminate_dead_code() graph_module.recompile() return PassResult(graph_module, True) ``` -### Registration +**Critical rules:** (1) handle both dialects via `EdgeOpOverload` check, (2) `copy_meta` on every new node, (3) lift scalars to tensors in edge dialect with `get_const_node`, (4) cache constants with `const_cache`, (5) for bool-output nodes use `callback=lambda m: {**m, "val": m["val"].to(torch.bool)}` in `create_node`. + +### Approach C: Built-in Decomposition Table +**Ref:** `_passes/decompose_triu.py`. Uses `make_fx` + `get_decompositions`. Only works if PyTorch has a registered decomp. -1. Add to `_passes/__init__.py` (alphabetical order): - ```python - from .decompose_my_op import DecomposeMyOp - ``` +### Registration (all decompose passes) +1. `_passes/__init__.py` — import + `__all__` +2. `_passes/qnn_pass_manager.py` — import + `transform_for_annotation_pipeline` + `transform_for_export_pipeline` + `get_capture_program_passes` +3. `_passes/utils.py` — add to `get_passes_dependency_for_capture_program()` with `[RemoveRedundancy]` dependency + +--- + +## Common Gotchas + +- **Op name mismatch**: `aten.clamp`→`ReluMinMax`, `aten.expand`→`Tile`, `aten.select_copy`→`StridedSlice`. Search by functionality. +- **Multi-output ops**: Use `wrapper_idx=i` + `getitem` skip op +- **Negative dims**: QNN needs positive → `dim = dim % len(shape)` +- **QCOM_AXIS_ORDER**: `LayoutTransform` permutes NCHW→NHWC; remap axis with `.index(dim)`. `get_tensor()` auto-permutes data. +- **Int64 indices**: Add to `I64_IN_OPS` in `i64_to_i32.py` + `.to(torch.int32)` in builder (see `op_gather.py`) +- **Recompose passes**: Detect primitive sequences and replace with single native op. Ref: `recompose_pixel_unshuffle.py` +- **`partition/common_defs.py`**: Remove op from `to_be_implemented_operator` when adding support +- **HTP doc bugs**: If runtime fails but docs say supported → test on-device always. + +--- + +## Error Debugging + +| Error | Cause | Fix | +|-------|-------|-----| +| `KeyError: 'aten.my_op.default'` | Builder not registered | Check `builders/__init__.py` + `@register_node_visitor` | +| `was not decomposed or delegated` | Op in skip decomp but partitioner rejected | Check builder `define_node` errors; check `I64_IN_OPS` | +| `QNN_GRAPH_ERROR` / `validateOpConfig failed` | HTP doesn't support config | Check params vs HTP Op Def Supplement | +| `Tensor mismatching datatypes` | Quantized: not all inputs annotated | Use custom annotator with `SharedQuantizationSpec` | +| `ValueError: Validation failed` | Wrong annotation | Check `is_math_invariant`; use `annotate_in_out_obs_sharing_op` | +| `Expected dtype int64 for index` | Op fell back to CPU with int32 index | Add to `I64_IN_OPS` + `.to(int32)` in builder | +| `Numerical mismatch` | Precision issue | Quantized: check quant params. Float: HTP FP16 precision limit | + +**Debug order:** Float test first → then quantized. If float fails → builder/config issue. If only quantized fails → annotation issue. + +--- -2. Add to `_passes/qnn_pass_manager.py` imports and both pipeline methods: - - `transform_for_annotation_pipeline` (before quantizer) - - `transform_for_export_pipeline` (before `to_edge`) +## Quick Reference Checklists -3. Remove the op from `to_be_implemented_operator` in `partition/common_defs.py` +**Native QNN Op:** `qnn_constants.py` → `op_my_op.py` → `builders/__init__.py` → `htp_rules.py` → `lpai_rules.py` → `layout_transform.py` → `tests/models.py` → `test_qnn_delegate.py` → `partition/utils.py` (skip decomp) → `common_defs.py` (remove to_be_implemented) → `builders/README.md` -### Notes -- The decomposed module must only use ops that QNN already supports -- `ep.run_decompositions()` ensures the graph is in edge IR form -- `remap` maps placeholder names in the decomposed graph to actual nodes in the target graph -- No separate quantizer annotation needed — the decomposed ops already have their own annotations +**Decompose Pass:** `_passes/decompose_my_op.py` → `_passes/__init__.py` → `qnn_pass_manager.py` (annotation + export + capture) → `_passes/utils.py` (dependency) → `tests/models.py` → `test_qnn_delegate.py` → `common_defs.py` → `builders/README.md` diff --git a/.github/workflows/_test_riscv.yml b/.github/workflows/_test_riscv.yml index 79eec6cbb4c..e3b049bd614 100644 --- a/.github/workflows/_test_riscv.yml +++ b/.github/workflows/_test_riscv.yml @@ -12,13 +12,36 @@ on: required: false type: number default: 30 + model: + description: 'Which model to run. Possible values are: add, mv2 (mobilenetv2)' + required: false + type: string + default: 'add' + xnnpack: + description: 'Whether to enable XNNPACK' + required: false + type: boolean + default: false + quantize: + description: 'Produce an 8-bit quantized model' + required: false + type: boolean + default: false + qemu-cpu: + description: 'Configuration(s) for the CPU to emulate with QEMU, expecting a JSON array' + required: true + type: string + docker-image: + description: 'The docker image to use for this job' + required: false + type: string jobs: run: uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main with: runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-gcc11 + docker-image: ci-image:executorch-ubuntu-24.04-gcc14 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: ${{ inputs.timeout }} @@ -29,4 +52,37 @@ jobs: source .ci/scripts/utils.sh install_executorch "--use-pt-pinned-commit" - bash .ci/scripts/test_riscv_qemu.sh + # Allows failure in `echo | jq | while read` pipeline to bubble up and fail the workflow + set -o pipefail + + echo '${{ inputs.qemu-cpu }}' | jq -r '.[]' | while IFS= read -r qemu_cpu; do + export QEMU_CPU="${qemu_cpu}" + export GCC_VERSION=14 + bash .ci/scripts/test_riscv_qemu.sh \ + --model="${{ inputs.model }}" \ + ${{ inputs.xnnpack && '--xnnpack --verbose-xnnpack' || '' }} \ + ${{ inputs.quantize && '--quantize' || '' }} + + # Generate markdown table from riscv_test/${{ inputs.model }}_riscv.etdump.json, sorted by sum_ms + ( + etdump_json="riscv_test/${{ inputs.model }}_riscv.etdump.json" + echo "### Model=${{ inputs.model }} XNNPACK=${{ inputs.xnnpack }} Quantize=${{ inputs.quantize }} QEMU_CPU='${QEMU_CPU}'" + jq -r ' + def r3: (. * 1000 | round) / 1000; + ["Section","Op","Count","Sum (ms)","Avg (ms)","Max (ms)","Microkernels"], + ["---","---","---","---","---","---","---"], + ( [ (.ops[] | . + {section: "ops"}), + (.framework[] | . + {section: "framework"}) ] + | sort_by(-.sum_ms) | .[] + | [.section, .op, .count, (.sum_ms|r3), (.avg_ms|r3), (.max_ms|r3), ((.kernels // []) | join(", "))] ) + | "| " + (map(tostring) | join(" | ")) + " |" + ' "${etdump_json}" + echo + echo "
Registered XNNPACK microkernels" + echo + jq -r '.registered_kernels[] | "- `" + . + "`"' "${etdump_json}" + echo + echo "
" + echo + ) >> $GITHUB_STEP_SUMMARY + done diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 087917c1116..e1eaba6b7c1 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -150,7 +150,7 @@ jobs: # Run Gemma 4 31B tests (quant unit tests + pipeline integration tests) pip install gguf - python -m pytest examples/models/gemma4_31b/quant/tests/ examples/models/gemma4_31b/tests/ -v -o "addopts=" + python -m pytest examples/models/gemma4_31b/quant/tests/ examples/models/gemma4_31b/tests/ --ignore=examples/models/gemma4_31b/tests/test_mlx_pipeline.py -v -o "addopts=" export-model-cuda-artifact: name: export-model-cuda-artifact @@ -185,6 +185,8 @@ jobs: name: "dinov2-small-imagenet1k-1-layer" - repo: "SocialLocalMobile" name: "Qwen3.5-35B-A3B-HQQ-INT4" + - repo: "SocialLocalMobile" + name: "gemma-4-31B-it-HQQ-INT4" quant: - "non-quantized" - "quantized-int4-tile-packed" @@ -204,6 +206,15 @@ jobs: repo: "SocialLocalMobile" name: "Qwen3.5-35B-A3B-HQQ-INT4" quant: "quantized-int4-weight-only" + # Gemma 4 31B uses a prequantized checkpoint, only tile-packed + - model: + repo: "SocialLocalMobile" + name: "gemma-4-31B-it-HQQ-INT4" + quant: "non-quantized" + - model: + repo: "SocialLocalMobile" + name: "gemma-4-31B-it-HQQ-INT4" + quant: "quantized-int4-weight-only" # Voxtral Realtime only supports int4-tile-packed on CUDA - model: repo: "mistralai" @@ -258,7 +269,7 @@ jobs: with: timeout: 90 secrets-env: EXECUTORCH_HF_TOKEN - runner: ${{ matrix.model.name == 'Qwen3.5-35B-A3B-HQQ-INT4' && 'linux.aws.a100' || 'linux.g5.4xlarge.nvidia.gpu' }} + runner: ${{ (matrix.model.name == 'Qwen3.5-35B-A3B-HQQ-INT4' || matrix.model.name == 'gemma-4-31B-it-HQQ-INT4') && 'linux.aws.a100' || 'linux.g5.4xlarge.nvidia.gpu' }} gpu-arch-type: cuda gpu-arch-version: 12.6 use-custom-docker-registry: false @@ -315,6 +326,8 @@ jobs: name: "dinov2-small-imagenet1k-1-layer" - repo: "SocialLocalMobile" name: "Qwen3.5-35B-A3B-HQQ-INT4" + - repo: "SocialLocalMobile" + name: "gemma-4-31B-it-HQQ-INT4" quant: - "non-quantized" - "quantized-int4-tile-packed" @@ -334,6 +347,15 @@ jobs: repo: "SocialLocalMobile" name: "Qwen3.5-35B-A3B-HQQ-INT4" quant: "quantized-int4-weight-only" + # Gemma 4 31B uses a prequantized checkpoint, only tile-packed + - model: + repo: "SocialLocalMobile" + name: "gemma-4-31B-it-HQQ-INT4" + quant: "non-quantized" + - model: + repo: "SocialLocalMobile" + name: "gemma-4-31B-it-HQQ-INT4" + quant: "quantized-int4-weight-only" # Voxtral Realtime only supports int4-tile-packed on CUDA - model: repo: "mistralai" @@ -382,7 +404,7 @@ jobs: quant: "non-quantized" with: timeout: 90 - runner: ${{ matrix.model.name == 'Qwen3.5-35B-A3B-HQQ-INT4' && 'linux.aws.a100' || 'linux.g5.4xlarge.nvidia.gpu' }} + runner: ${{ (matrix.model.name == 'Qwen3.5-35B-A3B-HQQ-INT4' || matrix.model.name == 'gemma-4-31B-it-HQQ-INT4') && 'linux.aws.a100' || 'linux.g5.4xlarge.nvidia.gpu' }} gpu-arch-type: cuda gpu-arch-version: 12.6 use-custom-docker-registry: false diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index 4778d08fcdc..d429db16053 100644 --- a/.github/workflows/mlx.yml +++ b/.github/workflows/mlx.yml @@ -60,6 +60,7 @@ jobs: backends/mlx/test/test_passes.py \ backends/mlx/test/test_pattern_utils.py \ backends/mlx/test/test_partitioner.py \ + examples/models/gemma4_31b/tests/test_mlx_pipeline.py \ -v echo "::endgroup::" diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index bfe83b853d1..c56f9d16ddc 100644 --- a/.github/workflows/pull.yml +++ b/.github/workflows/pull.yml @@ -547,7 +547,7 @@ jobs: setup_script_args="" if [[ ${{ matrix.os}} == "bare_metal" ]]; then toolchain_prefix=arm-none-eabi- - threshold="111000" # 111 KiB + threshold="118000" # GCC 15.2 baseline: 115868 bytes, 2026-05-13 toolchain_cmake=examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake elif [[ ${{ matrix.os}} == "zephyr-preset" ]]; then setup_script_args="--target-toolchain zephyr" @@ -717,10 +717,10 @@ jobs: strategy: matrix: include: - - test_arm_baremetal: test_pytest_ops_no_target - - test_arm_baremetal: test_pytest_ops_tosa - - test_arm_baremetal: test_pytest_models_tosa - - test_arm_baremetal: test_run_tosa + - test_arm_backend: test_pytest_ops_no_target + - test_arm_backend: test_pytest_ops_tosa + - test_arm_backend: test_pytest_models_tosa + - test_arm_backend: test_run_tosa fail-fast: false with: runner: linux.2xlarge @@ -738,10 +738,10 @@ jobs: .ci/scripts/setup-arm-baremetal-tools.sh --disable-ethos-u-deps - ARM_TEST=${{ matrix.test_arm_baremetal }} + ARM_TEST=${{ matrix.test_arm_backend }} - # Test test_arm_baremetal.sh with test - backends/arm/test/test_arm_baremetal.sh "${ARM_TEST}" + # Test test_arm_backend.sh with test + backends/arm/test/test_arm_backend.sh "${ARM_TEST}" test-llama-runner-qnn-linux: name: test-llama-runner-qnn-linux diff --git a/.github/workflows/riscv64.yml b/.github/workflows/riscv64.yml index aa970dc6270..14b9ad62047 100644 --- a/.github/workflows/riscv64.yml +++ b/.github/workflows/riscv64.yml @@ -1,4 +1,4 @@ -name: RISC-V +name: Test RISC-V Backend on: push: @@ -25,6 +25,42 @@ jobs: test-riscv: name: test-riscv uses: ./.github/workflows/_test_riscv.yml + strategy: + fail-fast: false + matrix: + include: + - { model: add, xnnpack: false, quantize: false } + - { model: add, xnnpack: true, quantize: false } + - { model: mv2, xnnpack: false, quantize: false } + - { model: mv2, xnnpack: true, quantize: false } + - { model: mv2, xnnpack: true, quantize: true } + - { model: mobilebert, xnnpack: false, quantize: false } + - { model: mobilebert, xnnpack: true, quantize: false } + - { model: mobilebert, xnnpack: true, quantize: true } + - { model: llama2, xnnpack: false, quantize: false } + - { model: llama2, xnnpack: true, quantize: false } + - { model: llama2, xnnpack: true, quantize: true } + - { model: resnet18, xnnpack: false, quantize: false } + - { model: resnet18, xnnpack: true, quantize: false } + - { model: resnet18, xnnpack: true, quantize: true } permissions: id-token: write contents: read + with: + model: ${{ matrix.model }} + xnnpack: ${{ matrix.xnnpack }} + quantize: ${{ matrix.quantize }} + # If XNNPACK, test with multiple RVV length, disabled otherwise + qemu-cpu: >- + ${{ + case( + matrix.xnnpack, '[ + "rv64,zba=true,zbb=true,zbs=true,v=true,vlen=128,elen=64,vext_spec=v1.0", + "rv64,zba=true,zbb=true,zbs=true,v=true,vlen=256,elen=64,vext_spec=v1.0", + "rv64,zba=true,zbb=true,zbs=true,v=true,vlen=512,elen=64,vext_spec=v1.0" + ]', + '[ + "rv64,zba=true,zbb=true,zbs=true,v=false" + ]' + ) + }} diff --git a/.github/workflows/trunk.yml b/.github/workflows/trunk.yml index 303a942c74a..3dad255589c 100644 --- a/.github/workflows/trunk.yml +++ b/.github/workflows/trunk.yml @@ -62,7 +62,12 @@ jobs: uses: pytorch/test-infra/.github/workflows/linux_job.yml@main strategy: matrix: - target: [ethos-u55, cortex-m55, ethos-u85] + include: + - { readme: zephyr/samples/hello-executorch/README.md, target: ethos-u55 } + - { readme: zephyr/samples/hello-executorch/README.md, target: cortex-m55 } + - { readme: zephyr/samples/hello-executorch/README.md, target: ethos-u85 } + - { readme: zephyr/samples/mv2-ethosu/README.md, target: ethos-u55 } + - { readme: zephyr/samples/mv2-ethosu/README.md, target: ethos-u85 } fail-fast: false with: runner: linux.2xlarge @@ -79,143 +84,9 @@ jobs: # Test zephyr backend set -e - # Support comma-separated TARGET_LIST or ${{ matrix.target }} list, e.g., TARGET_LIST="ethos-u55,cortex-m55,ethos-u85" - if [ -z "${TARGET_LIST:-}" ]; then - IFS=',' read -r -a TARGETS <<< "${{ matrix.target }}" - else - IFS=',' read -r -a TARGETS <<< "${TARGET_LIST}" - fi - - export EXECUTORCH_PROJ_ROOT=$(realpath $(pwd)) - ZEPHYR_README_PATH="zephyr/README.md" - ZEPHYR_SAMPLES_README_PATH="zephyr/samples/hello-executorch/README.md" - - # Source utility scripts - . .ci/scripts/utils.sh - . .ci/scripts/zephyr-utils.sh - - # check that zephyr/README.md and zephyr/executorch.yaml are in sync - verify_zephyr_readme - - # Based on instructions in zephyr/README.md and zephyr/samples/hello-executorch/README.md - - run_command_block_from_readme "${ZEPHYR_README_PATH}" "" - - # Make sure to backup the zephyr_scratch folder if it exists to allow for local - # testing that does not lose code/data - if [ -d "zephyr_scratch" ]; then - mv "zephyr_scratch" "zephyr_scratch.backup.$(date +%Y%m%d%H%M%S)" - fi - mkdir -p zephyr_scratch/ - - cd zephyr_scratch - export ZEPHYR_PROJ_ROOT=$(realpath $(pwd)) - - echo "---- Zephyr SDK ----" - # Use ZephyrSDK if on the disk (e.g. setup in the docker) - # Check for a zephyr-sdk-0.17.4 directory and make a symlink if found in parent directories - if sdk_dir=$(find ../../.. -maxdepth 4 -type d -name 'zephyr-sdk-0.17.4' -print -quit) && [ -n "${sdk_dir}" ]; then - echo "---- Found pre downloaded Zephyr SDK in ${sdk_dir} ----" - ln -s "${sdk_dir}" . - fi - - # Download and setup Zephyr SDK 0.17.4 if not already present - if [ ! -d "zephyr-sdk-0.17.4" ]; then - echo "---- Downloading Zephyr SDK ----" - wget https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v0.17.4/zephyr-sdk-0.17.4_linux-x86_64.tar.xz - tar -xf zephyr-sdk-0.17.4_linux-x86_64.tar.xz - rm -f zephyr-sdk-0.17.4_linux-x86_64.tar.xz* - fi - - ./zephyr-sdk-0.17.4/setup.sh -c -t arm-zephyr-eabi - export ZEPHYR_SDK_INSTALL_DIR=$(realpath ./zephyr-sdk-0.17.4) - - cd ${ZEPHYR_PROJ_ROOT} - - run_command_block_from_readme "${ZEPHYR_README_PATH}" "" - - cp ${EXECUTORCH_PROJ_ROOT}/zephyr/executorch.yaml zephyr/submanifests/ - - run_command_block_from_readme "${ZEPHYR_README_PATH}" "" - - # Switch to executorch in this PR e.g. replace modules/lib/executorch with the root folder of this repo - # instead of doing a re-checkout and figure out the correct commit hash etc - rm -Rf modules/lib/executorch - ln -s ${EXECUTORCH_PROJ_ROOT} modules/lib/executorch - - # Setup git local user for Executorch git to allows modules/lib/executorch/examples/arm/setup.sh be run inside CI later - # Configure git user only if not already set - if ! git config --get user.name >/dev/null 2>&1; then - git config --global user.name "Github Executorch" - fi - if ! git config --get user.email >/dev/null 2>&1; then - git config --global user.email "github_executorch@arm.com" - fi - - run_command_block_from_readme "${ZEPHYR_README_PATH}" "" - - run_command_block_from_readme "${ZEPHYR_README_PATH}" "" - - for TARGET in "${TARGETS[@]}"; do - TARGET="$(echo "$TARGET" | xargs)" # trim whitespace - - echo "---- ${TARGET} ----" - rm -Rf build - - if [[ ${TARGET} == "ethos-u55" || ${TARGET} == "cortex-m55" ]]; then - BOARD="corstone300" - elif [[ ${TARGET} == "ethos-u85" ]]; then - BOARD="corstone320" - else - echo "Fail unsupport target selection ${TARGET}" - exit 1 - fi - - echo "---- ${TARGET} Board ${BOARD} FVP setup ----" - run_command_block_from_readme "${ZEPHYR_SAMPLES_README_PATH}" "" - - echo "---- ${TARGET} Create PTE ----" - run_command_block_from_readme "${ZEPHYR_SAMPLES_README_PATH}" "" - - echo "---- ${TARGET} Build and run ----" - run_command_block_from_readme "${ZEPHYR_SAMPLES_README_PATH}" "" - done - - # MV2 Ethos-U sample (NPU targets only — skips cortex-m55) - MV2_README_PATH="zephyr/samples/mv2-ethosu/README.md" - - for TARGET in "${TARGETS[@]}"; do - TARGET="$(echo "$TARGET" | xargs)" - - if [[ ${TARGET} == "cortex-m55" ]]; then - echo "---- Skipping MV2 for ${TARGET} (NPU-only sample) ----" - continue - fi - - if [[ ${TARGET} == "ethos-u55" ]]; then - BOARD="corstone300" - elif [[ ${TARGET} == "ethos-u85" ]]; then - BOARD="corstone320" - else - echo "Fail unsupported target selection ${TARGET}" - exit 1 - fi - - echo "---- MV2 ${TARGET} ----" - rm -Rf build - - echo "---- MV2 ${TARGET} Board ${BOARD} FVP setup ----" - run_command_block_from_readme "${ZEPHYR_SAMPLES_README_PATH}" "" - - echo "---- MV2 ${TARGET} Create PTE ----" - run_command_block_from_readme "${MV2_README_PATH}" "" - - # Build only — FVP cycle-accurate simulation of MV2 is too slow - # for the CI timeout. Corstone-300 also lacks enough ISRAM for - # the runtime pools. The build step still catches link regressions. - echo "---- MV2 ${TARGET} Build only ----" - run_command_block_from_readme "${MV2_README_PATH}" "" - done + .ci/scripts/test_zephyr.sh \ + --targets "${TARGET_LIST:-${{ matrix.target }}}" \ + --zephyr-samples-readme-path "${{ matrix.readme }}" test-models-linux-aarch64: name: test-models-linux-aarch64 @@ -349,14 +220,14 @@ jobs: strategy: matrix: include: - - test_arm_baremetal: test_pytest_ops_ethos_u55 - - test_arm_baremetal: test_pytest_models_ethos_u55 - - test_arm_baremetal: test_run_ethos_u55 - - test_arm_baremetal: test_pytest_ops_ethos_u85 - - test_arm_baremetal: test_pytest_models_ethos_u85 - - test_arm_baremetal: test_run_ethos_u85 - - test_arm_baremetal: test_smaller_stories_llama - - test_arm_baremetal: test_memory_allocation + - test_arm_backend: test_pytest_ops_ethos_u55 + - test_arm_backend: test_pytest_models_ethos_u55 + - test_arm_backend: test_run_ethos_u55 + - test_arm_backend: test_pytest_ops_ethos_u85 + - test_arm_backend: test_pytest_models_ethos_u85 + - test_arm_backend: test_run_ethos_u85 + - test_arm_backend: test_smaller_stories_llama + - test_arm_backend: test_memory_allocation fail-fast: false with: runner: linux.2xlarge.memory @@ -378,15 +249,15 @@ jobs: # Hopefully this is high enough for this setup. sudo sysctl fs.inotify.max_user_watches=1048576 # 1024 * 1024 - ARM_TEST=${{ matrix.test_arm_baremetal }} + ARM_TEST=${{ matrix.test_arm_backend }} # Output test report on pytest runs so that github can surface failing tests. if [[ -n "${RUNNER_TEST_RESULTS_DIR:-}" ]]; then export PYTEST_ADDOPTS="--junit-xml=${RUNNER_TEST_RESULTS_DIR}/${ARM_TEST}.xml ${PYTEST_ADDOPTS:-}" fi - # Test test_arm_baremetal.sh with test - backends/arm/test/test_arm_baremetal.sh "${ARM_TEST}" + # Test test_arm_backend.sh with test + backends/arm/test/test_arm_backend.sh "${ARM_TEST}" test-arm-backend-vkml: name: test-arm-backend-vkml @@ -397,7 +268,7 @@ jobs: strategy: matrix: include: - - test_arm_baremetal: test_pytest_ops_vkml + - test_arm_backend: test_pytest_ops_vkml fail-fast: false with: runner: linux.2xlarge.memory @@ -418,14 +289,14 @@ jobs: # Hopefully this is high enough for this setup. sudo sysctl fs.inotify.max_user_watches=1048576 # 1024 * 1024 - ARM_TEST=${{ matrix.test_arm_baremetal }} + ARM_TEST=${{ matrix.test_arm_backend }} # Output test report on pytest runs so that github can surface failing tests. if [[ -n "${RUNNER_TEST_RESULTS_DIR:-}" ]]; then export PYTEST_ADDOPTS="--junit-xml=${RUNNER_TEST_RESULTS_DIR}/${ARM_TEST}.xml ${PYTEST_ADDOPTS:-}" fi - backends/arm/test/test_arm_baremetal.sh "${ARM_TEST}" + backends/arm/test/test_arm_backend.sh "${ARM_TEST}" test-arm-ootb-linux: name: test-arm-ootb-linux diff --git a/.lintrunner.toml b/.lintrunner.toml index 2955e54c70c..43a53342abd 100644 --- a/.lintrunner.toml +++ b/.lintrunner.toml @@ -361,6 +361,7 @@ include_patterns = [ 'devtools/visualization/**/*.py', 'docs/**/*.py', # 'examples/**/*.py', + 'examples/arm/**/*.py', 'examples/openvino/**/*.py', # 'exir/**/*.py', # 'extension/**/*.py', diff --git a/CMakeLists.txt b/CMakeLists.txt index 359a0e0f5e4..ac40d86d273 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -160,6 +160,23 @@ announce_configured_options(BUILD_TESTING) load_build_preset() include(${PROJECT_SOURCE_DIR}/tools/cmake/preset/default.cmake) +# Keep bare-metal installs enabled only when ExecuTorch owns the top-level +# build. Standalone consumers (e.g., the runner) set +# EXECUTORCH_BAREMETAL_SKIP_INSTALL=ON but still add ExecuTorch as a subproject, +# which cannot satisfy our install() export dependencies until their own targets +# are configured. +if(DEFINED EXECUTORCH_BAREMETAL_SKIP_INSTALL + AND EXECUTORCH_BAREMETAL_SKIP_INSTALL + AND NOT (CMAKE_SOURCE_DIR STREQUAL PROJECT_SOURCE_DIR) +) + set(CMAKE_SKIP_INSTALL_RULES + ON + CACHE BOOL + "Skip install() rules when ExecuTorch is consumed as a subproject" + FORCE + ) +endif() + # Enable ccache if available find_program(CCACHE_PROGRAM ccache) if(CCACHE_PROGRAM) diff --git a/Makefile b/Makefile index ba61dddce44..9b7f24b2f83 100644 --- a/Makefile +++ b/Makefile @@ -91,7 +91,7 @@ # # ============================================================================== -.PHONY: voxtral-cuda voxtral-cpu voxtral-metal voxtral-mlx voxtral_realtime-cuda voxtral_realtime-cpu voxtral_realtime-metal voxtral_realtime-mlx voxtral_tts-cpu voxtral_tts-cuda whisper-cuda whisper-cuda-debug whisper-cpu whisper-metal parakeet-cuda parakeet-cuda-debug parakeet-cpu parakeet-metal parakeet-mlx parakeet-vulkan dinov2-cuda dinov2-cuda-debug sortformer-cuda sortformer-cpu silero-vad-cpu llama-cuda llama-cuda-debug llama-cpu llava-cpu gemma3-cuda gemma3-cpu gemma4_31b-cuda qwen3_5_moe-cuda qwen3_5_moe-metal clean help +.PHONY: voxtral-cuda voxtral-cpu voxtral-metal voxtral-mlx voxtral_realtime-cuda voxtral_realtime-cpu voxtral_realtime-metal voxtral_realtime-mlx voxtral_tts-cpu voxtral_tts-cuda whisper-cuda whisper-cuda-debug whisper-cpu whisper-metal parakeet-cuda parakeet-cuda-debug parakeet-cpu parakeet-metal parakeet-mlx parakeet-vulkan dinov2-cuda dinov2-cuda-debug sortformer-cuda sortformer-cpu silero-vad-cpu llama-cuda llama-cuda-debug llama-cpu llava-cpu gemma3-cuda gemma3-cpu gemma4_31b-cuda gemma4_31b-mlx qwen3_5_moe-cuda qwen3_5_moe-metal clean help help: @echo "This Makefile adds targets to build runners for various models on various backends. Run using \`make \`. Available targets:" @@ -127,6 +127,7 @@ help: @echo " gemma3-cuda - Build Gemma3 runner with CUDA backend" @echo " gemma3-cpu - Build Gemma3 runner with CPU backend" @echo " gemma4_31b-cuda - Build Gemma 4 31B runner with CUDA backend" + @echo " gemma4_31b-mlx - Build Gemma 4 31B runner with MLX backend" @echo " qwen3_5_moe-cuda - Build Qwen3.5 MoE runner with CUDA backend" @echo " qwen3_5_moe-metal - Build Qwen3.5 MoE runner with Metal backend" @echo " clean - Clean build artifacts" @@ -435,6 +436,15 @@ gemma4_31b-cuda: @echo "✓ Build complete!" @echo " Binary: cmake-out/examples/models/gemma4_31b/gemma4_31b_runner" +gemma4_31b-mlx: + @echo "==> Building and installing ExecuTorch with MLX..." + cmake --workflow --preset mlx-release + @echo "==> Building Gemma 4 31B runner with MLX..." + cd examples/models/gemma4_31b && cmake --workflow --preset gemma4-31b-mlx + @echo "" + @echo "✓ Build complete!" + @echo " Binary: cmake-out/examples/models/gemma4_31b/gemma4_31b_runner" + qwen3_5_moe-metal: @echo "==> Building and installing ExecuTorch with Metal..." cmake --workflow --preset llm-release-metal diff --git a/backends/apple/coreml/compiler/torch_ops.py b/backends/apple/coreml/compiler/torch_ops.py index 29c7120feb7..7f6abe980e1 100644 --- a/backends/apple/coreml/compiler/torch_ops.py +++ b/backends/apple/coreml/compiler/torch_ops.py @@ -28,6 +28,21 @@ from executorch.exir.dim_order_utils import get_memory_format +_IOS18_QUANT_HINT = ( + "ExecuTorch hint: pass `compile_specs=CoreMLBackend.generate_compile_specs(" + "minimum_deployment_target=ct.target.iOS18)` (or higher) to " + "`CoreMLPartitioner` when lowering models that use `quantize_(...)`." +) + + +def _raise_with_executorch_hint(err: Exception) -> "BaseException": + """Re-raise a coremltools quantization error with ExecuTorch-specific guidance.""" + msg = str(err) + if "iOS18" in msg or "iOS 18" in msg: + raise ValueError(f"{msg}\n{_IOS18_QUANT_HINT}") from err + raise err + + # https://github.com/apple/coremltools/pull/2563 @register_torch_op(override=False) def split_copy(context, node): @@ -159,12 +174,15 @@ def dequantize_affine(context, node): f"Unsupported quantization range: {quant_min} to {quant_max}. CoreML only supports 4-bit and 8-bit quantization." ) - output = _utils._construct_constexpr_dequant_op( - int_data.astype(quantized_np_dtype), - zero_point, - scale, - name=node.name, - ) + try: + output = _utils._construct_constexpr_dequant_op( + int_data.astype(quantized_np_dtype), + zero_point, + scale, + name=node.name, + ) + except ValueError as e: + _raise_with_executorch_hint(e) context.add(output, node.name) @@ -211,9 +229,12 @@ def dequantize_codebook(context, node): f"Core ML ignores output_dtype {out_np_dtype} on torchao.dequantize_affine and instead uses the native precision." ) - output = _utils._construct_constexpr_lut_op( - codes.astype(np.int8), - codebook, - name=node.name, - ) + try: + output = _utils._construct_constexpr_lut_op( + codes.astype(np.int8), + codebook, + name=node.name, + ) + except ValueError as e: + _raise_with_executorch_hint(e) context.add(output, node.name) diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index 57bd793a10a..b4e550d8479 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/backends/apple/coreml/partition/coreml_partitioner.py @@ -63,6 +63,26 @@ ) +_ARG_MIN_MAX_TARGETS = ( + torch.ops.aten.argmax.default, + torch.ops.aten.argmin.default, + exir_ops.edge.aten.argmax.default, + exir_ops.edge.aten.argmin.default, +) + + +def _is_arg_min_max_over_flattened_input(node: torch.fx.Node) -> bool: + """``argmin``/``argmax`` with ``dim=None`` reduces over the flattened input. + + CoreML doesn't support that reduction shape and intermittently crashes + the process at runtime — see pytorch/executorch#11715. + """ + if node.target not in _ARG_MIN_MAX_TARGETS: + return False + dim = node.args[1] if len(node.args) >= 2 else node.kwargs.get("dim", None) + return dim is None + + def _is_view_op(op: torch._ops.OpOverload) -> bool: schema = op._schema if len(schema.arguments) == 0: @@ -132,6 +152,13 @@ def should_override_support(self, node) -> bool: ) return True + if _is_arg_min_max_over_flattened_input(node): + self.log_once( + "torch.ops.aten.{argmax, argmin}.default with dim=None is " + "not supported by CoreML. Overriding op support." + ) + return True + # TODO: enable this after bugs in ExecuTorch's partitioner are fixed # # If lower_full_graph=False, do not partition nodes with symbolic args because it can result in symbolic args # # in the placeholders due to partitioning, which CoreML does not support diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py index a2321ee199f..0e75d6024e4 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -386,6 +386,53 @@ def forward(self, x): self.assertIn("executorch_call_delegate", op_names) self.assertNotIn("aten.randn.default", op_names) + def test_argmax_argmin_dim_none_is_skipped(self): + """ + Regression test for https://github.com/pytorch/executorch/issues/11715. + + argmax/argmin with dim=None reduces over the flattened tensor, which + CoreML does not support; the resulting model intermittently crashes + the process at runtime. The partitioner must reject these so they + fall back to the portable backend, while still delegating the + ordinary dim=int form. + """ + + class FlatModel(torch.nn.Module): + def forward(self, x): + return torch.argmax(x, dim=None, keepdim=False) + torch.argmin( + x, dim=None + ) + + ep = torch.export.export( + FlatModel().eval(), (torch.randn(10, 10),), strict=True + ) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner()] + ) + op_names = [ + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + ] + self.assertIn("aten.argmax.default", op_names) + self.assertIn("aten.argmin.default", op_names) + + class DimModel(torch.nn.Module): + def forward(self, x): + return torch.argmax(x, dim=1) + + ep = torch.export.export(DimModel().eval(), (torch.randn(10, 10),), strict=True) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner()] + ) + op_names = [ + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + ] + self.assertIn("executorch_call_delegate", op_names) + self.assertNotIn("aten.argmax.default", op_names) + def test_deprecation_warning_for_to_backend_workflow(self): """ Test that the deprecated to_edge + to_backend workflow shows a deprecation warning. diff --git a/backends/apple/coreml/test/test_torch_ops.py b/backends/apple/coreml/test/test_torch_ops.py index de54b684ee7..10c3f01a585 100644 --- a/backends/apple/coreml/test/test_torch_ops.py +++ b/backends/apple/coreml/test/test_torch_ops.py @@ -317,6 +317,33 @@ def forward(self, x): et_prog = delegated_program.to_executorch() self._compare_outputs(et_prog, model, example_inputs) + def test_dequantize_affine_below_ios18_raises_with_hint(self): + """ + Regression test for https://github.com/pytorch/executorch/issues/13122. + + `quantize_(...)` with blockwise / int4 configurations requires iOS18. + coremltools raises a ValueError that does not mention how to fix the + deployment target on the ExecuTorch side; we wrap it to add the + partitioner-level guidance. + """ + model = torch.nn.Linear(64, 64) + quantize_( + model, + IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32)), + ) + ep = torch.export.export(model.eval(), (torch.randn(1, 64),), strict=True) + with self.assertRaises(ValueError) as cm: + executorch.exir.to_edge_transform_and_lower( + ep, + partitioner=[ + self._coreml_partitioner(minimum_deployment_target=ct.target.iOS17) + ], + ) + msg = str(cm.exception) + self.assertIn("iOS18", msg) + self.assertIn("CoreMLPartitioner", msg) + self.assertIn("minimum_deployment_target", msg) + if __name__ == "__main__": test_runner = TestTorchOps() diff --git a/backends/arm/CMakeLists.txt b/backends/arm/CMakeLists.txt index 0c8b241522c..d8a6c1afce7 100644 --- a/backends/arm/CMakeLists.txt +++ b/backends/arm/CMakeLists.txt @@ -63,17 +63,20 @@ if(EXECUTORCH_BUILD_ARM_BAREMETAL OR EXECUTORCH_BUILD_ARM_ETHOSU_LINUX) add_library(executorch_delegate_ethos_u STATIC ${_arm_backend_sources}) target_link_libraries(executorch_delegate_ethos_u PUBLIC executorch_core) + target_include_directories( + executorch_delegate_ethos_u PRIVATE ${_common_include_directories} + ) if(EXECUTORCH_BUILD_ARM_BAREMETAL) target_sources( executorch_delegate_ethos_u PRIVATE ${EXECUTORCH_ROOT}/backends/arm/runtime/EthosUBackend_Cortex_M.cpp ) - set(DRIVER_ETHOSU_INCLUDE_DIR + set(_ethosu_core_driver_include "${THIRD_PARTY_ROOT}/ethos-u-core-driver/include" ) target_include_directories( - executorch_delegate_ethos_u PRIVATE ${DRIVER_ETHOSU_INCLUDE_DIR} + executorch_delegate_ethos_u PRIVATE ${_ethosu_core_driver_include} ) target_link_libraries(executorch_delegate_ethos_u PUBLIC ethosu_core_driver) elseif(EXECUTORCH_BUILD_ARM_ETHOSU_LINUX) @@ -110,7 +113,25 @@ if(EXECUTORCH_BUILD_ARM_BAREMETAL OR EXECUTORCH_BUILD_ARM_ETHOSU_LINUX) ) endif() - install(TARGETS executorch_delegate_ethos_u EXPORT ExecuTorchTargets) + if(NOT CMAKE_SKIP_INSTALL_RULES) + install(TARGETS executorch_delegate_ethos_u EXPORT ExecuTorchTargets) + + if(TARGET ethosu_core_driver) + get_property( + _et_ethosu_core_driver_exported GLOBAL + PROPERTY ET_ETHOSU_CORE_DRIVER_EXPORTED + ) + if(NOT _et_ethosu_core_driver_exported) + install( + TARGETS ethosu_core_driver + EXPORT ExecuTorchTargets + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + ) + set_property(GLOBAL PROPERTY ET_ETHOSU_CORE_DRIVER_EXPORTED TRUE) + endif() + endif() + endif() endif() diff --git a/backends/arm/MODELS.md b/backends/arm/MODELS.md index 6f84d694297..bcb410764bf 100644 --- a/backends/arm/MODELS.md +++ b/backends/arm/MODELS.md @@ -10,6 +10,7 @@ - Inception v3 (IC3) - Llama - Gemma3n +- Qwen3-VL - Long Short-Term Memory (LSTM) - MobileNet V1 0.25 - MobileNet v2 (MV2) diff --git a/backends/arm/README.md b/backends/arm/README.md index f844b1d6422..e9afa5a928d 100644 --- a/backends/arm/README.md +++ b/backends/arm/README.md @@ -116,6 +116,33 @@ Developers who need local source builds can use: The current flow lowers to TOSA and converts to VGF for use in external projects, so the `executor_runner` is not typically used here. +### Compiling models with the Python API + +Use the Python API as the primary way to compile your own models. It lets you +keep model construction, export inputs, quantization, custom passes, and artifact +generation in your application code. The `aot_arm_compiler.py` script is useful +for simple examples and smoke tests, but production code should call the +ExecuTorch and Arm backend APIs directly. + +The delegated Python API flow is: + +1. Prepare the model and representative example inputs. +2. Create a target-specific Arm compile spec. +3. Export the model with `torch.export.export`. +4. Optionally quantize with the target-specific Arm quantizer and re-export the + quantized graph. +5. Create the matching Arm partitioner from the compile spec. +6. Lower with `to_edge_transform_and_lower`. +7. Convert to an ExecuTorch program and save the PTE file. + +For complete examples of that flow, including quantization and target-specific +compile specs, see: + +- `docs/source/backends/arm-ethos-u/tutorials/ethos-u-getting-started.md` +- `docs/source/backends/arm-vgf/tutorials/vgf-getting-started.md` + +Additional examples are available in `examples/arm`. + ### Direct Drive (experimental, Ethos-U85 on Linux) workflow Direct Drive enables execution on Ethos-U85 via the Linux driver stack. @@ -159,7 +186,8 @@ scp -P 2222 arm_test/cmake-out/executor_runner root@127.0.0.1:/tmp/ #### Direct Drive model (PTE) workflow -Create a PTE file: +For a quick test with the example `add` model, +`aot_arm_compiler.py` can be used: ``` python3 -m backends.arm.scripts.aot_arm_compiler \ @@ -170,50 +198,64 @@ python3 -m backends.arm.scripts.aot_arm_compiler \ --direct_drive ``` +For production use, the Python API described in +[Compiling models with the Python API](#compiling-models-with-the-python-api) +should be used. Use an Ethos-U85 target and set the Direct Drive `extra_flags` when creating the `EthosUCompileSpec`: + +```python +compile_spec = EthosUCompileSpec( + target="ethos-u85-256", + extra_flags=["--separate-io-regions", "--cop-format=COP2"], +) +``` + +Then save the generated program as e.g. `model.pte` or +update the copy and run commands below to match your output file name. + Copy the `executor_runner` binary and the generated PTE file to the running FVP: ``` -scp -P 2222 arm_test/cmake-out/executor_runner add_arm_delegate_ethos-u85-256.pte root@127.0.0.1:/tmp/ +scp -P 2222 arm_test/cmake-out/executor_runner model.pte root@127.0.0.1:/tmp/ ``` Run the model on the FVP: ``` -ssh -p 2222 root@127.0.0.1 -t "/tmp/executor_runner -model_path /tmp/add_arm_delegate_ethos-u85-256.pte -num_executions 1" +ssh -p 2222 root@127.0.0.1 -t "/tmp/executor_runner -model_path /tmp/model.pte -num_executions 1" ``` ## Testing There are two approaches for running the tests for the Arm backend. This section will explain these two approaches: -### Using test_arm_baremetal.sh +### Using test_arm_backend.sh -The backend provides a script `backends/arm/test/test_arm_baremetal.sh`, which is used in the `trunk` CI workflow. +The backend provides a script `backends/arm/test/test_arm_backend.sh`, which is used in the `trunk` CI workflow. This approach is useful for checking your change against this workflow on your own machine. These scripts also install the necessary dependencies to run the tests. Below is an overview of some of the testing options this script provides: | Command | Description | | ---------------------------------------------------- | ------------------------------------------------------------ | -| `test_arm_baremetal.sh test_pytest_ops_no_target` | Runs operator unit tests for non-target specific use-cases. | -| `test_arm_baremetal.sh test_pytest_models_no_target` | Runs model unit tests for non-target specific use-cases. | -| `test_arm_baremetal.sh test_pytest_ops_tosa` | Runs operator unit tests for TOSA specific use-cases. | -| `test_arm_baremetal.sh test_pytest_models_tosa` | Runs model unit tests for TOSA specific use-cases. | -| `test_arm_baremetal.sh test_run_tosa` | Runs end-to-end unit tests for TOSA specific use-cases. | -| `test_arm_baremetal.sh test_pytest_ops_ethos_u55` | Runs operator unit tests for Ethos-U55 specific use-cases. | -| `test_arm_baremetal.sh test_pytest_models_ethos_u55` | Runs model unit tests for Ethos-U55 specific use-cases. | -| `test_arm_baremetal.sh test_run_ethos_u55` | Runs end-to-end unit tests for Ethos-U55 specific use-cases. | -| `test_arm_baremetal.sh test_pytest_ops_ethos_u85` | Runs operator unit tests for Ethos-U85 specific use-cases. | -| `test_arm_baremetal.sh test_pytest_models_ethos_u85` | Runs model unit tests for Ethos-U85 specific use-cases. | -| `test_arm_baremetal.sh test_run_ethos_u85` | Runs end-to-end unit tests for Ethos-U85 specific use-cases. | -| `test_arm_baremetal.sh test_pytest_ops_vkml` | Runs operator unit tests for VGF specific use-cases. | -| `test_arm_baremetal.sh test_pytest_models_vkml` | Runs model unit tests for VGF specific use-cases. | -| `test_arm_baremetal.sh test_run_vkml` | Runs end-to-end unit tests for VGF specific use-cases. | -| `test_arm_baremetal.sh test_model_smollm2-135M` | Runs some models with Corstone FVP. | -| `test_arm_baremetal.sh test_smaller_stories_llama` | Runs E2E model tests on Corstone FVP. | -| `test_arm_baremetal.sh test_memory_allocation` | Runs memory allocation tests for Ethos-U specific targets | - -For more information, please refer to the `backends/arm/test/test_arm_baremetal.sh` script. +| `test_arm_backend.sh test_pytest_ops_no_target` | Runs operator unit tests for non-target specific use-cases. | +| `test_arm_backend.sh test_pytest_models_no_target` | Runs model unit tests for non-target specific use-cases. | +| `test_arm_backend.sh test_pytest_ops_tosa` | Runs operator unit tests for TOSA specific use-cases. | +| `test_arm_backend.sh test_pytest_models_tosa` | Runs model unit tests for TOSA specific use-cases. | +| `test_arm_backend.sh test_run_tosa` | Runs end-to-end unit tests for TOSA specific use-cases. | +| `test_arm_backend.sh test_pytest_ops_ethos_u55` | Runs operator unit tests for Ethos-U55 specific use-cases. | +| `test_arm_backend.sh test_pytest_models_ethos_u55` | Runs model unit tests for Ethos-U55 specific use-cases. | +| `test_arm_backend.sh test_run_ethos_u55` | Runs end-to-end unit tests for Ethos-U55 specific use-cases. | +| `test_arm_backend.sh test_pytest_ops_ethos_u85` | Runs operator unit tests for Ethos-U85 specific use-cases. | +| `test_arm_backend.sh test_pytest_models_ethos_u85` | Runs model unit tests for Ethos-U85 specific use-cases. | +| `test_arm_backend.sh test_run_ethos_u85` | Runs end-to-end unit tests for Ethos-U85 specific use-cases. | +| `test_arm_backend.sh test_pytest_ops_vkml` | Runs operator unit tests for VGF specific use-cases. | +| `test_arm_backend.sh test_pytest_models_vkml` | Runs model unit tests for VGF specific use-cases. | +| `test_arm_backend.sh test_run_vkml` | Runs end-to-end unit tests for VGF specific use-cases. | +| `test_arm_backend.sh test_model_smollm2_135M` | Runs some models with Corstone FVP. | +| `test_arm_backend.sh test_smaller_stories_llama` | Runs E2E model tests on Corstone FVP. | +| `test_arm_backend.sh test_memory_allocation` | Runs memory allocation tests for Ethos-U specific targets | + +For more information, please refer to the `backends/arm/test/test_arm_backend.sh` script. ### Using pytest diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index 6b39d9a25c3..4a8857be82a 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -97,6 +97,7 @@ from .decompose_var_pass import DecomposeVarPass # noqa from .decompose_where_scalar_other_pass import DecomposeWhereScalarOtherPass # noqa from .decorate_fp32_to_int32_casting_pass import DecorateFp32toInt32CastingPass # noqa +from .deduplicate_get_attr_pass import DeduplicateGetAttrPass # noqa from .ensure_unique_output_nodes_pass import EnsureUniqueOutputNodesPass # noqa from .fold_qdq_with_annotated_qparams_pass import ( # noqa FoldAndAnnotateQParamsPass, diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 6bdb43d79f5..bf39cbe44ea 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -97,6 +97,7 @@ DecomposeVarPass, DecomposeWhereScalarOtherPass, DecorateFp32toInt32CastingPass, + DeduplicateGetAttrPass, EnsureUniqueOutputNodesPass, FoldAndAnnotateQParamsPass, FuseBatchNorm2dPass, @@ -651,6 +652,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule): [ ReplaceInfAndLimitValuesPass(tfa_pass=True), DecomposeMaskedFillPass(tfa_pass=True), + DeduplicateGetAttrPass(tfa_pass=True), ] ) diff --git a/backends/arm/_passes/arm_pass_utils.py b/backends/arm/_passes/arm_pass_utils.py index 9176f761220..000f92135eb 100644 --- a/backends/arm/_passes/arm_pass_utils.py +++ b/backends/arm/_passes/arm_pass_utils.py @@ -413,3 +413,16 @@ def to_2tuple(value): if len(value) == 1: return (value[0], value[0]) return tuple(value) + + +def permute_fake_tensor_metadata( + fake_tensor: FakeTensor, permute_dims: tuple[int, ...] +) -> FakeTensor: + permuted_shape = tuple(fake_tensor.shape[dim] for dim in permute_dims) + meta_tensor = torch.empty( + permuted_shape, + dtype=fake_tensor.dtype, + device="meta", + requires_grad=fake_tensor.requires_grad, + ) + return FakeTensor(fake_tensor.fake_mode, meta_tensor, fake_tensor.fake_device) diff --git a/backends/arm/_passes/decompose_int_pow_pass.py b/backends/arm/_passes/decompose_int_pow_pass.py index 2df8d3b2522..bb29d34d6bf 100644 --- a/backends/arm/_passes/decompose_int_pow_pass.py +++ b/backends/arm/_passes/decompose_int_pow_pass.py @@ -32,12 +32,16 @@ def call_operator(self, op, args, kwargs, meta): x = args[0] exp = args[1] - # Handle zero first and return early if exp == 0: - # return a tensor of ones with the same shape as x - return super().call_operator( + zeros = super().call_operator( + exir_ops.edge.aten.sub.Tensor, (x, x), {}, meta, True + ) + ones = super().call_operator( exir_ops.edge.aten.full_like.default, (x, 1), {}, meta, True ) + return super().call_operator( + exir_ops.edge.aten.add.Tensor, (zeros, ones), {}, meta, True + ) if not isinstance(exp, int): return super().call_operator(op, args, kwargs, meta) diff --git a/backends/arm/_passes/deduplicate_get_attr_pass.py b/backends/arm/_passes/deduplicate_get_attr_pass.py new file mode 100644 index 00000000000..201a9036e34 --- /dev/null +++ b/backends/arm/_passes/deduplicate_get_attr_pass.py @@ -0,0 +1,71 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Any, Set, Type + +import torch +from executorch.backends.arm._passes import ArmPass +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx import GraphModule, Node +from torchao.quantization.pt2e.utils import get_new_attr_name_with_prefix + + +class DeduplicateGetAttrPass(ArmPass): + """Give duplicate get_attr nodes distinct backing attributes. + + Torchao's constant folder can delete a shared backing attribute while + another get_attr node still refers to it. Keep separate graph nodes so PT2E + can attach per-use observers and backend lowering can process constants per + use. + + """ + + _passes_required_after: Set[Type[ExportPass]] = set() + + def _get_attr(self, graph_module: GraphModule, target: str) -> Any: + attr: Any = graph_module + for target_atom in target.split("."): + attr = getattr(attr, target_atom) + return attr + + def _copy_attr(self, graph_module: GraphModule, node: Node) -> str: + """Register a new attribute referring to the same data as the original + one. + """ + + assert isinstance(node.target, str) + attr = self._get_attr(graph_module, node.target) + get_new_attr_name = get_new_attr_name_with_prefix( + f"_deduplicated_get_attr_{node.name}_" + ) + attr_name = get_new_attr_name(graph_module) + + if isinstance(attr, torch.nn.Parameter): + graph_module.register_parameter(attr_name, attr) + elif isinstance(attr, torch.Tensor): + graph_module.register_buffer(attr_name, attr) + else: + setattr(graph_module, attr_name, attr) + + return attr_name + + def call(self, graph_module: GraphModule) -> PassResult: + seen_targets: set[str] = set() + modified = False + + for node in graph_module.graph.find_nodes(op="get_attr"): + + if node.target not in seen_targets: + seen_targets.add(node.target) + continue + + node.target = self._copy_attr(graph_module, node) + modified = True + + if modified: + graph_module.graph.lint() + graph_module.recompile() + + return PassResult(graph_module, modified) diff --git a/backends/arm/_passes/normalize_delegate_io_layout_pass.py b/backends/arm/_passes/normalize_delegate_io_layout_pass.py index d1b1d964b87..c55eec5c851 100644 --- a/backends/arm/_passes/normalize_delegate_io_layout_pass.py +++ b/backends/arm/_passes/normalize_delegate_io_layout_pass.py @@ -11,6 +11,7 @@ create_node, get_first_fake_tensor, is_param_node, + permute_fake_tensor_metadata, ) from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops @@ -63,8 +64,8 @@ def _normalize_input_layout(self, graph_module: torch.fx.GraphModule) -> bool: args=(node, list(transpose_perm)), from_node=node, ) - permute_node.meta["val"] = exir_ops.edge.aten.permute_copy.default( - node.meta["val"], list(transpose_perm) + permute_node.meta["val"] = permute_fake_tensor_metadata( + get_first_fake_tensor(node), transpose_perm ) users = [user for user in node.users if user != permute_node] @@ -91,8 +92,8 @@ def _rewrite_output_arg( args=(arg, list(dim_order)), from_node=arg, ) - permute_node.meta["val"] = exir_ops.edge.aten.permute_copy.default( - output_fake, list(dim_order) + permute_node.meta["val"] = permute_fake_tensor_metadata( + output_fake, dim_order ) return permute_node, True diff --git a/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py b/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py index e000b3d6fe8..fa6f6f7988c 100644 --- a/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py +++ b/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py @@ -11,13 +11,15 @@ class RemovePermutesAroundElementwiseTosaOps(RemovePermutesAroundElementwiseOps): - permutable_ops = { - *RemovePermutesAroundElementwiseOps.permutable_ops, - *TableOps.unary_table_ops.keys(), - *TableOps.special_table_ops, - exir_ops.backend.tosa.RESCALE.default, - exir_ops.backend.tosa.TABLE.default, - } + def __init__(self) -> None: + super().__init__( + extra_permutable_ops={ + *TableOps.unary_table_ops.keys(), + *TableOps.special_table_ops, + exir_ops.backend.tosa.RESCALE.default, + exir_ops.backend.tosa.TABLE.default, + } + ) def permute_subgraph(self, subgraph): # Original function will always permute constant nodes which is wrong for table ops diff --git a/backends/arm/_passes/rewrite_conv_pass.py b/backends/arm/_passes/rewrite_conv_pass.py index 27565e93452..a51f1ae0555 100644 --- a/backends/arm/_passes/rewrite_conv_pass.py +++ b/backends/arm/_passes/rewrite_conv_pass.py @@ -16,6 +16,7 @@ get_first_fake_tensor, get_param_tensor, is_persistent_buffer, + permute_fake_tensor_metadata, ) from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( get_input_qparams, @@ -153,7 +154,13 @@ def _add_bias( output_dtype = node.meta["val"].dtype bias_data = torch.zeros(size=(output_channels,), dtype=output_dtype) - with graph_module.graph.inserting_after(weight_node): + # Constant placeholders must appear before user-input placeholders in + # the graph. Insert the synthetic bias at the first placeholder slot + # instead of near the conv node. + first_placeholder = next( + n for n in graph_module.graph.nodes if n.op == "placeholder" + ) + with graph_module.graph.inserting_before(first_placeholder): bias_node = create_constant_placeholder( self.exported_program, graph=graph_module.graph, @@ -415,9 +422,10 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 args=(x, list(pre_permute_dims)), from_node=node, ) - x.meta["val"] = exir_ops.edge.aten.permute_copy.default( - input_fake_tensor, list(pre_permute_dims) + input_tensor_for_tosa_fake = permute_fake_tensor_metadata( + input_fake_tensor, pre_permute_dims ) + x.meta["val"] = input_tensor_for_tosa_fake weight = self._rewrite_weight( graph_module, weight, @@ -425,7 +433,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 permute_dims=OHWI_ORDER, name_suffix="ohwi", ) - input_tensor_for_tosa_fake = input_fake_tensor.permute(pre_permute_dims) weight_fake_tensor = get_first_fake_tensor(weight) conv_args = ( x, @@ -465,9 +472,10 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 args=(x, list(pre_permute_dims)), from_node=node, ) - x.meta["val"] = exir_ops.edge.aten.permute_copy.default( - input_fake_tensor, list(pre_permute_dims) + input_tensor_for_tosa_fake = permute_fake_tensor_metadata( + input_fake_tensor, pre_permute_dims ) + x.meta["val"] = input_tensor_for_tosa_fake weight = self._rewrite_weight( graph_module, weight, @@ -475,9 +483,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 permute_dims=ODHWI_ORDER, name_suffix="odhwi", ) - input_tensor_for_tosa_fake = input_fake_tensor.permute( - pre_permute_dims - ) weight_fake_tensor = get_first_fake_tensor(weight) elif self._is_depthwise_conv2d(node): target_op = exir_ops.backend.tosa.DEPTHWISE_CONV2D.default @@ -490,9 +495,10 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 args=(x, list(pre_permute_dims)), from_node=node, ) - x.meta["val"] = exir_ops.edge.aten.permute_copy.default( - input_fake_tensor, list(pre_permute_dims) + input_tensor_for_tosa_fake = permute_fake_tensor_metadata( + input_fake_tensor, pre_permute_dims ) + x.meta["val"] = input_tensor_for_tosa_fake kh, kw = weight_shape[2], weight_shape[3] in_channels = input_fake_tensor.shape[1] m_length = weight_shape[0] // in_channels @@ -504,9 +510,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 name_suffix="hwicm", reshape_dims=(kh, kw, in_channels, m_length), ) - input_tensor_for_tosa_fake = input_fake_tensor.permute( - pre_permute_dims - ) weight_fake_tensor = get_first_fake_tensor(weight) else: target_op = exir_ops.backend.tosa.CONV2D.default @@ -519,9 +522,10 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 args=(x, list(pre_permute_dims)), from_node=node, ) - x.meta["val"] = exir_ops.edge.aten.permute_copy.default( - input_fake_tensor, list(pre_permute_dims) + input_tensor_for_tosa_fake = permute_fake_tensor_metadata( + input_fake_tensor, pre_permute_dims ) + x.meta["val"] = input_tensor_for_tosa_fake weight = self._rewrite_weight( graph_module, weight, @@ -529,9 +533,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 permute_dims=NHWC_ORDER, name_suffix="ohwi", ) - input_tensor_for_tosa_fake = input_fake_tensor.permute( - pre_permute_dims - ) weight_fake_tensor = get_first_fake_tensor(weight) conv_args = ( @@ -606,8 +607,8 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 TosaSpecialDtype.meta_key() ): node_replacement.meta[TosaSpecialDtype.meta_key()] = special_dtype - node_replacement.meta["val"] = exir_ops.edge.aten.permute_copy.default( - node_replacement_fake_tensor, list(post_permute_dims) + node_replacement.meta["val"] = permute_fake_tensor_metadata( + node_replacement_fake_tensor, post_permute_dims ) node.replace_all_uses_with(node_replacement) diff --git a/backends/arm/_passes/rewrite_upsample.py b/backends/arm/_passes/rewrite_upsample.py index 9f81f5cbbe5..68a088286fa 100644 --- a/backends/arm/_passes/rewrite_upsample.py +++ b/backends/arm/_passes/rewrite_upsample.py @@ -13,6 +13,7 @@ create_node, create_shape_node, get_first_fake_tensor, + permute_fake_tensor_metadata, ) from executorch.backends.arm.tosa.mapping import TosaSpecialDtype from executorch.exir.dialects._ops import ops as exir_ops @@ -196,7 +197,7 @@ def call(self, graph_module): args=(x, list(self._NHWC_ORDER)), from_node=node, ) - pre_permute.meta["val"] = exir_ops.edge.aten.permute_copy.default( + pre_permute.meta["val"] = permute_fake_tensor_metadata( get_first_fake_tensor(x), list(self._NHWC_ORDER) ) @@ -255,9 +256,8 @@ def call(self, graph_module): args=(node_replacement, list(self._NHWC_INVERSE_ORDER)), from_node=node, ) - post_permute.meta["val"] = exir_ops.edge.aten.permute_copy.default( - node_replacement_fake, - list(self._NHWC_INVERSE_ORDER), + post_permute.meta["val"] = permute_fake_tensor_metadata( + node_replacement_fake, self._NHWC_INVERSE_ORDER ) node.replace_all_uses_with(post_permute) graph_module.graph.erase_node(node) diff --git a/backends/arm/cmake/ArmEthosUSDK.cmake b/backends/arm/cmake/ArmEthosUSDK.cmake new file mode 100644 index 00000000000..03affdf69bb --- /dev/null +++ b/backends/arm/cmake/ArmEthosUSDK.cmake @@ -0,0 +1,60 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +include_guard(GLOBAL) + +function(arm_ethos_u_content_ready SDK_PATH OUT_VAR) + if(EXISTS "${SDK_PATH}/core_platform" AND EXISTS "${SDK_PATH}/core_software") + set(${OUT_VAR} + TRUE + PARENT_SCOPE + ) + else() + set(${OUT_VAR} + FALSE + PARENT_SCOPE + ) + endif() +endfunction() + +function(arm_ethos_u_default_fetch SDK_PATH OUT_VAR) + arm_ethos_u_content_ready("${SDK_PATH}" _arm_ethos_ready) + if(_arm_ethos_ready) + set(${OUT_VAR} + OFF + PARENT_SCOPE + ) + else() + set(${OUT_VAR} + ON + PARENT_SCOPE + ) + endif() +endfunction() + +function(arm_ensure_ethos_u_content SDK_PATH EXECUTORCH_ROOT FETCH_REQUESTED) + arm_ethos_u_content_ready("${SDK_PATH}" _arm_ethos_ready_before) + + if(_arm_ethos_ready_before) + return() + endif() + + if(NOT FETCH_REQUESTED) + message( + FATAL_ERROR + "No Ethos-U content found at ${SDK_PATH}. Run examples/arm/setup.sh or enable FETCH_ETHOS_U_CONTENT=ON." + ) + endif() + + fetch_ethos_u_content(${SDK_PATH} ${EXECUTORCH_ROOT}) + + arm_ethos_u_content_ready("${SDK_PATH}" _arm_ethos_ready_after) + if(NOT _arm_ethos_ready_after) + message( + FATAL_ERROR + "Failed to fetch Ethos-U content into ${SDK_PATH}. Inspect the logs above." + ) + endif() +endfunction() diff --git a/backends/arm/cmake/ArmRunnerUtils.cmake b/backends/arm/cmake/ArmRunnerUtils.cmake new file mode 100644 index 00000000000..e67f38eec22 --- /dev/null +++ b/backends/arm/cmake/ArmRunnerUtils.cmake @@ -0,0 +1,69 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +include_guard(GLOBAL) + +# Helper routines shared by the standalone runner and any superbuild that reuses +# the runner targets. + +function(arm_runner_require_baremetal_targets) + if(NOT TARGET extension_runner_util) + message( + FATAL_ERROR + "extension_runner_util target missing. Configure ExecuTorch (or the standalone runner) with EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=ON." + ) + endif() + + if(NOT TARGET quantized_ops_lib OR NOT TARGET quantized_kernels) + message( + FATAL_ERROR + "quantized kernels not found. Ensure EXECUTORCH_BUILD_KERNELS_QUANTIZED=ON when configuring ExecuTorch." + ) + endif() + + if(NOT TARGET cortex_m_ops_lib OR NOT TARGET cortex_m_kernels) + message( + FATAL_ERROR + "cortex_m backend not found. Ensure EXECUTORCH_BUILD_CORTEX_M=ON when configuring ExecuTorch." + ) + endif() +endfunction() + +# Ensure a runner target emits its binary to a predictable location. Uses +# FALLBACK_DIR when TARGET_NAME has no runtime output directory set, and also +# fills per-configuration runtime output directories for multi-config generators +# when they are unset. +function(arm_runner_configure_runtime_output TARGET_NAME FALLBACK_DIR) + if(NOT TARGET ${TARGET_NAME}) + return() + endif() + + get_target_property(_base_runtime_dir ${TARGET_NAME} RUNTIME_OUTPUT_DIRECTORY) + if(NOT _base_runtime_dir + OR _base_runtime_dir STREQUAL "_base_runtime_dir-NOTFOUND" + OR "${_base_runtime_dir}" STREQUAL "" + ) + set_target_properties( + ${TARGET_NAME} PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${FALLBACK_DIR}" + ) + set(_base_runtime_dir "${FALLBACK_DIR}") + endif() + + if(CMAKE_CONFIGURATION_TYPES) + foreach(_cfg ${CMAKE_CONFIGURATION_TYPES}) + string(TOUPPER ${_cfg} _cfg_upper) + set(_cfg_prop "RUNTIME_OUTPUT_DIRECTORY_${_cfg_upper}") + get_target_property(_cfg_dir ${TARGET_NAME} ${_cfg_prop}) + if(NOT _cfg_dir + OR _cfg_dir STREQUAL "_cfg_dir-NOTFOUND" + OR "${_cfg_dir}" STREQUAL "" + ) + set_target_properties( + ${TARGET_NAME} PROPERTIES ${_cfg_prop} "${_base_runtime_dir}/${_cfg}" + ) + endif() + endforeach() + endif() +endfunction() diff --git a/backends/arm/runtime/VGFSetup.cpp b/backends/arm/runtime/VGFSetup.cpp index 0e1a66dd5a0..b62a6b2ec23 100644 --- a/backends/arm/runtime/VGFSetup.cpp +++ b/backends/arm/runtime/VGFSetup.cpp @@ -727,8 +727,7 @@ bool VgfRepr::process_vgf( VkDataGraphPipelineCreateInfoARM graph_pipeline_info{ .sType = VK_STRUCTURE_TYPE_DATA_GRAPH_PIPELINE_CREATE_INFO_ARM, .pNext = &shader_info, - .flags = VK_PIPELINE_CREATE_2_FAIL_ON_PIPELINE_COMPILE_REQUIRED_BIT | - VK_PIPELINE_CREATE_2_EARLY_RETURN_ON_FAILURE_BIT_KHR, + .flags = VK_PIPELINE_CREATE_2_EARLY_RETURN_ON_FAILURE_BIT_KHR, .layout = vk_pipeline_layout, .resourceInfoCount = static_cast(data_graph_resources.size()), .pResourceInfos = data_graph_resources.data(), diff --git a/backends/arm/scripts/build_executor_runner.sh b/backends/arm/scripts/build_executor_runner.sh index f2ffd2e27a7..55f1a272b9e 100755 --- a/backends/arm/scripts/build_executor_runner.sh +++ b/backends/arm/scripts/build_executor_runner.sh @@ -9,6 +9,9 @@ set -eu script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) et_root_dir=$(cd ${script_dir}/../../.. && pwd) et_root_dir=$(realpath ${et_root_dir}) +runner_source_dir=${et_root_dir}/examples/arm/executor_runner/standalone +runner_source_dir=$(realpath ${runner_source_dir}) +preset_file=${et_root_dir}/tools/cmake/preset/arm_baremetal.cmake toolchain=arm-none-eabi-gcc setup_path_script=${et_root_dir}/examples/arm/arm-scratch/setup_path.sh _setup_msg="please refer to ${et_root_dir}/examples/arm/setup.sh to properly install necessary tools." @@ -101,6 +104,9 @@ toolchain_cmake=$(realpath ${toolchain_cmake}) source ${setup_path_script} +[[ -f ${preset_file} ]] \ + || { echo "Missing ${preset_file}. ${_setup_msg}"; exit 1; } + if [[ ${pte_file} == "semihosting" ]]; then pte_data="-DSEMIHOSTING=ON" else @@ -122,13 +128,13 @@ else fi fi ethosu_tools_dir=$(realpath ${ethosu_tools_dir}) -ethos_u_root_dir="$ethosu_tools_dir/ethos-u" +ethos_u_root_dir="${ethosu_tools_dir}/ethos-u" mkdir -p "${ethos_u_root_dir}" -ethosu_tools_dir=$(realpath ${ethos_u_root_dir}) - -et_build_dir=${et_build_root}/cmake-out -mkdir -p ${et_build_dir} -et_build_dir=$(realpath ${et_build_dir}) +ethos_u_root_dir=$(realpath ${ethos_u_root_dir}) +cmsis_nn_local_path="" +if [[ -d "${ethos_u_root_dir}/core_software/cmsis-nn" ]]; then + cmsis_nn_local_path=$(realpath "${ethos_u_root_dir}/core_software/cmsis-nn") +fi if [[ ${system_config} == "" ]] then @@ -160,34 +166,47 @@ echo "-------------------------------------------------------------------------- echo "Build Arm ${toolchain/-gcc/} executor_runner for ${target} PTE: ${pte_file} using ${system_config} ${memory_mode} ${extra_build_flags} to '${output_folder}'" echo "--------------------------------------------------------------------------------" -cd ${et_root_dir}/examples/arm/executor_runner - if [ "$bundleio" = true ] ; then build_bundleio_flags=" -DET_BUNDLE_IO=ON " + candidate_build_dir="${et_build_root}/cmake-out" + if [[ -d "${candidate_build_dir}" ]]; then + candidate_build_dir=$(realpath "${candidate_build_dir}") + build_bundleio_flags+=" -DET_BUILD_DIR_PATH=${candidate_build_dir} " + fi + if [[ -n "${BUNDLED_PROGRAM_LIBRARY_DIR:-}" ]]; then + build_bundleio_flags+=" -DBUNDLED_PROGRAM_LIBRARY_DIR=${BUNDLED_PROGRAM_LIBRARY_DIR} " + fi fi if [ "$build_with_etdump" = true ] ; then build_with_etdump_flags=" -DEXECUTORCH_ENABLE_EVENT_TRACER=ON -DET_DUMP_INTERMEDIATE_OUTPUTS=ON " fi +devtools_flags="" +if [ "$bundleio" = true ] || [ "$build_with_etdump" = true ] ; then + devtools_flags=" -DEXECUTORCH_BUILD_DEVTOOLS=ON " +fi -echo "Building with BundleIO/etdump/extra flags: ${build_bundleio_flags} ${build_with_etdump_flags} ${extra_build_flags}" +echo "Building with BundleIO/etdump/extra flags: ${build_bundleio_flags} ${build_with_etdump_flags} ${devtools_flags} ${extra_build_flags}" cmake \ - -DCMAKE_BUILD_TYPE=${build_type} \ - -DCMAKE_TOOLCHAIN_FILE=${toolchain_cmake} \ - -DTARGET_CPU=${target_cpu} \ - -DET_DIR_PATH:PATH=${et_root_dir} \ - -DET_BUILD_DIR_PATH:PATH=${et_build_dir} \ - -DETHOS_SDK_PATH:PATH=${ethos_u_root_dir} \ - -DETHOSU_TARGET_NPU_CONFIG=${target} \ - ${pte_data} \ - ${build_bundleio_flags} \ - ${build_with_etdump_flags} \ - -DPYTHON_EXECUTABLE=$(which python3) \ - -DSYSTEM_CONFIG=${system_config} \ - -DMEMORY_MODE=${memory_mode} \ + -S ${runner_source_dir} \ + -B ${output_folder} \ + -DEXECUTORCH_ROOT=${et_root_dir} \ + -DCMAKE_BUILD_TYPE=${build_type} \ + -DCMAKE_TOOLCHAIN_FILE=${toolchain_cmake} \ + -DTARGET_CPU=${target_cpu} \ + -DETHOSU_TARGET_NPU_CONFIG=${target} \ + -DEXECUTORCH_BUILD_PRESET_FILE=${preset_file} \ + -DEXECUTORCH_BAREMETAL_SKIP_INSTALL=OFF \ + ${pte_data} \ + ${build_bundleio_flags} \ + ${build_with_etdump_flags} \ + ${devtools_flags} \ + -DSYSTEM_CONFIG=${system_config} \ + -DMEMORY_MODE=${memory_mode} \ -DEXECUTORCH_SELECT_OPS_LIST="${select_ops_list}" \ - ${extra_build_flags} \ - -B ${output_folder} + -DETHOS_SDK_PATH:PATH=${ethos_u_root_dir} \ + ${cmsis_nn_local_path:+-DCMSIS_NN_LOCAL_PATH:PATH=${cmsis_nn_local_path}} \ + ${extra_build_flags} echo "[${BASH_SOURCE[0]}] Configured CMAKE" diff --git a/backends/arm/scripts/build_executorch.sh b/backends/arm/scripts/build_executorch.sh index cf7e327b9ce..828cec77ae8 100755 --- a/backends/arm/scripts/build_executorch.sh +++ b/backends/arm/scripts/build_executorch.sh @@ -85,6 +85,7 @@ cmake_args=( -DCMAKE_BUILD_TYPE=${build_type} -DEXECUTORCH_BUILD_DEVTOOLS=${build_devtools} -DEXECUTORCH_BUILD_ARM_ETDUMP=${build_with_etdump} + -DEXECUTORCH_BAREMETAL_SKIP_INSTALL=OFF ) if [[ ${is_linux_musl} -eq 1 ]]; then @@ -108,7 +109,7 @@ parallel_jobs="$(get_parallel_jobs)" if [[ ${is_linux_musl} -eq 1 ]]; then cmake --build ${et_build_dir} -j"${parallel_jobs}" --target executorch_delegate_ethos_u executor_runner --config ${build_type} -- else - cmake --build ${et_build_dir} -j"${parallel_jobs}" --target install --config ${build_type} -- + cmake --build ${et_build_dir} -j"${parallel_jobs}" --config ${build_type} fi set +x diff --git a/backends/arm/scripts/docgen/ethos-u/backends-arm-ethos-u-overview.md.in b/backends/arm/scripts/docgen/ethos-u/backends-arm-ethos-u-overview.md.in index 1990bc6d946..555c61fd13b 100644 --- a/backends/arm/scripts/docgen/ethos-u/backends-arm-ethos-u-overview.md.in +++ b/backends/arm/scripts/docgen/ethos-u/backends-arm-ethos-u-overview.md.in @@ -4,7 +4,7 @@ The Arm® Ethos™-U backend targets Edge/IoT-type AI use-cases by enabli [Arm® Ethos™-U55 NPU](https://www.arm.com/products/silicon-ip-cpu/ethos/ethos-u55), [Arm® Ethos™-U65 NPU](https://www.arm.com/products/silicon-ip-cpu/ethos/ethos-u65), and [Arm® Ethos™-U85 NPU](https://www.arm.com/products/silicon-ip-cpu/ethos/ethos-u85), leveraging [TOSA](https://www.mlplatform.org/tosa/) and the [ethos-u-vela](https://pypi.org/project/ethos-u-vela/) graph compiler. This document is a technical reference for using the Ethos-U backend, for a top level view with code examples -please refer to the [Arm Ethos-U Backend Tutorial](https://docs.pytorch.org/executorch/stable/tutorial-arm-ethos-u.html). +please refer to the [Arm Ethos-U Backend Tutorial](tutorials/ethos-u-getting-started.md). ## Features @@ -27,7 +27,7 @@ For the AOT flow, compilation of a model to `.pte` format using the Ethos-U back - [TOSA Serialization Library](https://www.mlplatform.org/tosa/software.html) for serializing the Exir IR graph into TOSA IR. - [Ethos-U Vela graph compiler](https://pypi.org/project/ethos-u-vela/) for compiling TOSA flatbuffers into an Ethos-U command stream. -And for building and running the example application available in `examples/arm/executor_runner/`: +And for building and running the example application available in `examples/arm/executor_runner/` through the standalone CMake entry point: - [Arm GNU Toolchain](https://developer.arm.com/Tools%20and%20Software/GNU%20Toolchain) for cross compilation. - [Arm® Corstone™ SSE-300 FVP](https://developer.arm.com/documentation/100966/1128/Arm--Corstone-SSE-300-FVP) for testing on a Arm® Cortex®-M55+Ethos-U55 reference design. - [Arm® Corstone™ SSE-320 FVP](https://developer.arm.com/documentation/109760/0000/SSE-320-FVP) for testing on a Arm® Cortex®-M85+Ethos-U85 reference design. @@ -55,7 +55,7 @@ For more information on quantization, see [Quantization](arm-ethos-u-quantizatio ## Runtime Integration -An example runtime application is available in [examples/arm/executor_runner](https://github.com/pytorch/executorch/blob/main/examples/arm/executor_runner/), and the steps requried for building and deploying it on a FVP it is explained in the previously mentioned [Arm Ethos-U Backend Tutorial](https://docs.pytorch.org/executorch/stable/tutorial-arm-ethos-u.html). +An example runtime application is available in [examples/arm/executor_runner](https://github.com/pytorch/executorch/blob/main/examples/arm/executor_runner/), with a standalone CMake entry point in `examples/arm/executor_runner/standalone`. The steps required for building and deploying it on an FVP are explained in the previously mentioned [Arm Ethos-U Backend Tutorial](tutorials/ethos-u-getting-started.md). The example application is recommended to use for testing basic functionality of your lowered models, as well as a starting point for developing runtime integrations for your own targets. For an in-depth explanation of the architecture of the executor_runner and the steps required for doing such an integration, please refer to [Ethos-U porting guide](https://github.com/pytorch/executorch/blob/main/examples/arm/ethos-u-porting-guide.md). @@ -153,7 +153,7 @@ ExecuTorch for the Ethos-U backend, you automatically install the compiler conta **→{doc}`/backends/arm-ethos-u/arm-ethos-u-troubleshooting` — Troubleshooting and common issues.** -**→{doc}`/backends/arm-ethos-u/tutorials/arm-ethos-u-tutorials` — Tutorials.** +**→{doc}`/backends/arm-ethos-u/tutorials/ethos-u-getting-started` — Getting started tutorial.** **→{doc}`/backends/arm-ethos-u/U55_op_support` — Ethos-U55 supported operators.** @@ -168,7 +168,7 @@ ExecuTorch for the Ethos-U backend, you automatically install the compiler conta arm-ethos-u-partitioner arm-ethos-u-quantization arm-ethos-u-troubleshooting -tutorials/arm-ethos-u-tutorials +tutorials/ethos-u-getting-started U55_op_support U85_op_support ``` diff --git a/backends/arm/scripts/docgen/ethos-u/ethos-u-getting-started-tutorial.md.in b/backends/arm/scripts/docgen/ethos-u/ethos-u-getting-started-tutorial.md.in index 68b73755317..ecd63afd8ba 100644 --- a/backends/arm/scripts/docgen/ethos-u/ethos-u-getting-started-tutorial.md.in +++ b/backends/arm/scripts/docgen/ethos-u/ethos-u-getting-started-tutorial.md.in @@ -76,35 +76,28 @@ To produce a pte file equivalent to the one above, run ### Runtime: -After the AOT compilation flow is done, the runtime can be cross compiled and linked to the produced `.pte`-file using the Arm cross-compilation toolchain. This is done in two steps: +After the AOT compilation flow is done, the runtime can be cross compiled and linked to the produced `.pte`-file using the Arm cross-compilation toolchain. Configure the standalone Arm executor runner CMake project to pull in the ExecuTorch build graph, link the Ethos-U delegate, and generate kernel bindings for any non-delegated ops. This produces the `arm_executor_runner` program that will run on target. -First, build and install the ExecuTorch libraries and EthosUDelegate: ``` # In ExecuTorch top-level, with sourced setup_path.sh -cmake -DCMAKE_BUILD_TYPE=Release --preset arm-baremetal -B cmake-out-arm . -cmake --build cmake-out-arm --target install -j$(nproc) -``` -Second, build and link the `arm_executor_runner` and generate kernel bindings for any non delegated ops. This is the actual program that will run on target. - -``` -# In ExecuTorch top-level, with sourced setup_path.sh -cmake -DCMAKE_TOOLCHAIN_FILE=`pwd`/examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake \ +cmake -S examples/arm/executor_runner/standalone \ + -B ethos_u_minimal_example \ + -DEXECUTORCH_ROOT=$(pwd) \ + -DCMAKE_TOOLCHAIN_FILE=$(pwd)/examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake \ -DCMAKE_BUILD_TYPE=Release \ -DET_PTE_FILE_PATH=ethos_u_minimal_example.pte \ -DTARGET_CPU=cortex-m55 \ -DETHOSU_TARGET_NPU_CONFIG=ethos-u55-128 \ -DMEMORY_MODE=Shared_Sram \ - -DSYSTEM_CONFIG=Ethos_U55_High_End_Embedded \ - -Bethos_u_minimal_example \ - examples/arm/executor_runner + -DSYSTEM_CONFIG=Ethos_U55_High_End_Embedded cmake --build ethos_u_minimal_example -j$(nproc) -- arm_executor_runner ``` ```{tip} -For a quick start, you can use the script `backends/arm/scripts/build_executor_runner.sh` to build the runner. +For a quick start, you can use the script `backends/arm/scripts/build_executor_runner.sh` to configure and build the standalone runner. To build a runner equivalent to the one above, run `./backends/arm/scripts/build_executor_runner.sh --pte=ethos_u_minimal_example.pte` -```` +``` The block diagram below shows, at the high level, how the various build artifacts are generated and are linked together to generate the final bare-metal executable. @@ -123,10 +116,14 @@ The example application is by default built with an input of ones, so the expect ## Takeaways In this tutorial you have learned how to use ExecuTorch to export a PyTorch model to an executable that can run on an embedded target, and then run that executable on simulated hardware. -To learn more, check out these learning paths: +To learn more, check out the [ExecuTorch on Arm Practical Labs](https://github.com/arm-education/executorch_on_arm_labs) series. This series provides a structured entry-point to developing with ExecuTorch on Arm, across both CPU and Ethos-U NPU. + +For quick learning paths showcasing short tutorials: -https://learn.arm.com/learning-paths/embedded-and-microcontrollers/rpi-llama3/ -https://learn.arm.com/learning-paths/embedded-and-microcontrollers/visualizing-ethos-u-performance/ +- [Run Llama3 on Raspberry Pi 5 with ExecuTorch](https://learn.arm.com/learning-paths/embedded-and-microcontrollers/rpi-llama3/) +- [Visualize Ethos-U NPU Performance on FVP](https://learn.arm.com/learning-paths/embedded-and-microcontrollers/visualizing-ethos-u-performance/) +- [Image Classification with ExecuTorch on NXP i.MX 93 (Ethos-U65)](https://learn.arm.com/learning-paths/embedded-and-microcontrollers/observing-ethos-u-on-nxp/) +- [Image Classification with ExecuTorch on Alif E8 DevKit (Ethos-U85)](https://learn.arm.com/learning-paths/embedded-and-microcontrollers/alif-image-classification/) ## FAQs diff --git a/backends/arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in b/backends/arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in index e12aafc55a7..53af3e2499d 100644 --- a/backends/arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in +++ b/backends/arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in @@ -79,7 +79,7 @@ described in the rest of this guide but with a concrete end-to-end sample. **→{doc}`/backends/arm-vgf/arm-vgf-troubleshooting` — Debug common issues.** -**→{doc}`/backends/arm-vgf/tutorials/arm-vgf-tutorials` — Tutorials.** +**→{doc}`/backends/arm-vgf/tutorials/vgf-getting-started` — Getting started tutorial.** **→{doc}`/backends/arm-vgf/VGF_op_support` — VGF supported operators.** @@ -92,6 +92,6 @@ described in the rest of this guide but with a concrete end-to-end sample. arm-vgf-partitioner arm-vgf-quantization arm-vgf-troubleshooting -tutorials/arm-vgf-tutorials +tutorials/vgf-getting-started VGF_op_support ``` diff --git a/backends/arm/scripts/docgen/vgf/vgf-getting-started-tutorial.md.in b/backends/arm/scripts/docgen/vgf/vgf-getting-started-tutorial.md.in index cf0d5a0dbad..531dea14b37 100644 --- a/backends/arm/scripts/docgen/vgf/vgf-getting-started-tutorial.md.in +++ b/backends/arm/scripts/docgen/vgf/vgf-getting-started-tutorial.md.in @@ -83,9 +83,9 @@ To produce a pte file equivalent to the one above, run `python -m backends.arm.scripts.aot_arm_compiler --model_name=add --delegate --quantize --output=simple_example.pte --target=vgf` ``` -### Runtime: +## Runtime -## Build executor runtime +### Build executor runtime After the AOT compilation flow is done, we can build the executor runner target. For this tutorial, the default runner can be used. Build it with the following configuration: @@ -115,7 +115,7 @@ The block diagram below demonstrates, at the high level, how the various build a ![](arm-delegate-runtime-build.svg) -## Deploying and running on device +### Deploying and running on device Since we are using the Vulkan emulation layer, we can run the executor runner with the VGF delegate on the host machine: diff --git a/backends/arm/scripts/run_fvp.sh b/backends/arm/scripts/run_fvp.sh index 9f0010189af..9da309fbe41 100755 --- a/backends/arm/scripts/run_fvp.sh +++ b/backends/arm/scripts/run_fvp.sh @@ -151,7 +151,7 @@ elif [[ ${target} == *"ethos-u55"* ]]; then -C mps3_board.telnetterminal0.start_telnet=0 \ -C mps3_board.uart0.out_file='-' \ -C mps3_board.uart0.shutdown_on_eot=1 \ - "${extra_args_u55[@]}" \ + ${extra_args_u55[@]+"${extra_args_u55[@]}"} \ -a "${elf_file}" \ ${data_file} \ --timelimit ${timeout} 2>&1 | sed 's/\r$//' | tee ${log_file} || true # seconds @@ -164,7 +164,7 @@ elif [[ ${target} == *"ethos-u85"* ]]; then -C mps4_board.telnetterminal0.start_telnet=0 \ -C mps4_board.uart0.out_file='-' \ -C mps4_board.uart0.shutdown_on_eot=1 \ - "${extra_args_u85[@]}" \ + ${extra_args_u85[@]+"${extra_args_u85[@]}"} \ -a "${elf_file}" \ ${data_file} \ --timelimit ${timeout} 2>&1 | sed 's/\r$//' | tee ${log_file} || true # seconds diff --git a/backends/arm/scripts/toolchain_utils.sh b/backends/arm/scripts/toolchain_utils.sh index 5b37bcee7b4..0ed1058fa28 100644 --- a/backends/arm/scripts/toolchain_utils.sh +++ b/backends/arm/scripts/toolchain_utils.sh @@ -22,20 +22,20 @@ source "${script_dir}/utils.sh" function gcc_select_toolchain() { if [[ "${ARCH}" == "x86_64" ]] ; then - toolchain_url="https://armkeil.blob.core.windows.net/developer/Files/downloads/gnu/13.3.rel1/binrel/arm-gnu-toolchain-13.3.rel1-x86_64-arm-none-eabi.tar.xz" - toolchain_dir="arm-gnu-toolchain-13.3.rel1-x86_64-arm-none-eabi" - toolchain_md5_checksum="0601a9588bc5b9c99ad2b56133b7f118" + toolchain_url="https://developer.arm.com/-/media/Files/downloads/gnu/15.2.rel1/binrel/arm-gnu-toolchain-15.2.rel1-x86_64-arm-none-eabi.tar.xz" + toolchain_dir="arm-gnu-toolchain-15.2.rel1-x86_64-arm-none-eabi" + toolchain_md5_checksum="da62bef8821e7fc2a9b5d023871036e0" toolchain_archive="${toolchain_dir}.tar.xz" elif [[ "${ARCH}" == "aarch64" ]] || [[ "${ARCH}" == "arm64" ]] ; then if [[ "${OS}" == "Darwin" ]]; then - toolchain_url="https://armkeil.blob.core.windows.net/developer/Files/downloads/gnu/13.3.rel1/binrel/arm-gnu-toolchain-13.3.rel1-darwin-arm64-arm-none-eabi.tar.xz" - toolchain_dir="arm-gnu-toolchain-13.3.rel1-darwin-arm64-arm-none-eabi" - toolchain_md5_checksum="f1c18320bb3121fa89dca11399273f4e" + toolchain_url="https://developer.arm.com/-/media/Files/downloads/gnu/15.2.rel1/binrel/arm-gnu-toolchain-15.2.rel1-darwin-arm64-arm-none-eabi.tar.xz" + toolchain_dir="arm-gnu-toolchain-15.2.rel1-darwin-arm64-arm-none-eabi" + toolchain_md5_checksum="e91fd6348ba0f3e5ec35eeba1ad7e2b8" toolchain_archive="${toolchain_dir}.tar.xz" elif [[ "${OS}" == "Linux" ]]; then - toolchain_url="https://armkeil.blob.core.windows.net/developer/Files/downloads/gnu/13.3.rel1/binrel/arm-gnu-toolchain-13.3.rel1-aarch64-arm-none-eabi.tar.xz" - toolchain_dir="arm-gnu-toolchain-13.3.rel1-aarch64-arm-none-eabi" - toolchain_md5_checksum="303102d97b877ebbeb36b3158994b218" + toolchain_url="https://developer.arm.com/-/media/Files/downloads/gnu/15.2.rel1/binrel/arm-gnu-toolchain-15.2.rel1-aarch64-arm-none-eabi.tar.xz" + toolchain_dir="arm-gnu-toolchain-15.2.rel1-aarch64-arm-none-eabi" + toolchain_md5_checksum="458c5d9b362726c9ac20c96f1894ae13" toolchain_archive="${toolchain_dir}.tar.xz" fi else diff --git a/backends/arm/test/misc/test_runner_utils.py b/backends/arm/test/misc/test_runner_utils.py index 10a8b6df3a6..3c78b21e008 100644 --- a/backends/arm/test/misc/test_runner_utils.py +++ b/backends/arm/test/misc/test_runner_utils.py @@ -77,3 +77,39 @@ def _fake_run_cmd(cmd, check=True): assert "-i i1.bin" in semihosting_cmd_arg assert long_input_paths[0] not in semihosting_cmd_arg assert long_input_paths[1] not in semihosting_cmd_arg + + +def test_get_elf_path_uses_repo_root_candidates(monkeypatch, tmp_path: Path) -> None: + elf_path = ( + tmp_path + / "arm_test" + / "arm_semihosting_executor_runner_corstone-300" + / "arm_executor_runner" + ) + elf_path.parent.mkdir(parents=True) + elf_path.write_bytes(b"") + + monkeypatch.setattr(runner_utils, "_elf_search_roots", lambda: [tmp_path]) + other_cwd = tmp_path / "elsewhere" + other_cwd.mkdir() + monkeypatch.chdir(other_cwd) + + assert runner_utils.get_elf_path("corstone-300") == str(elf_path) + + +def test_get_elf_path_accepts_nested_runner_output(monkeypatch, tmp_path: Path) -> None: + elf_path = ( + tmp_path + / "arm_test" + / "arm_semihosting_executor_runner_corstone-300" + / "examples" + / "arm" + / "executor_runner" + / "arm_executor_runner" + ) + elf_path.parent.mkdir(parents=True) + elf_path.write_bytes(b"") + + monkeypatch.setattr(runner_utils, "_elf_search_roots", lambda: [tmp_path]) + + assert runner_utils.get_elf_path("corstone-300") == str(elf_path) diff --git a/backends/arm/test/misc/test_tosa_dialect_avg_pool2d_adaptive.py b/backends/arm/test/misc/tosa_dialect/test_tosa_avg_pool2d_adaptive.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_avg_pool2d_adaptive.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_avg_pool2d_adaptive.py diff --git a/backends/arm/test/misc/test_tosa_dialect_conv2d.py b/backends/arm/test/misc/tosa_dialect/test_tosa_conv2d.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_conv2d.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_conv2d.py diff --git a/backends/arm/test/misc/test_tosa_dialect_dw_conv2d.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dw_conv2d.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_dw_conv2d.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_dw_conv2d.py diff --git a/backends/arm/test/misc/test_tosa_dialect_identity.py b/backends/arm/test/misc/tosa_dialect/test_tosa_identity.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_identity.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_identity.py diff --git a/backends/arm/test/misc/test_tosa_dialect_resize.py b/backends/arm/test/misc/tosa_dialect/test_tosa_resize.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_resize.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_resize.py diff --git a/backends/arm/test/misc/test_tosa_dialect_shape_ops.py b/backends/arm/test/misc/tosa_dialect/test_tosa_shape_ops.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_shape_ops.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_shape_ops.py diff --git a/backends/arm/test/models/Qwen3_VL/qwen3_vl_test_config.py b/backends/arm/test/models/Qwen3_VL/qwen3_vl_test_config.py new file mode 100644 index 00000000000..4bf47826ea6 --- /dev/null +++ b/backends/arm/test/models/Qwen3_VL/qwen3_vl_test_config.py @@ -0,0 +1,64 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from transformers.models.qwen3_vl.configuration_qwen3_vl import ( + Qwen3VLConfig, + Qwen3VLTextConfig, + Qwen3VLVisionConfig, +) + + +def get_qwen3_vl_2b_instruct_checkpoint_config() -> Qwen3VLConfig: + text_config = Qwen3VLTextConfig( + attention_bias=False, + attention_dropout=0.0, + bos_token_id=151643, # type: ignore[call-arg] + dtype="bfloat16", + eos_token_id=151645, # type: ignore[call-arg] + head_dim=128, + hidden_act="silu", + hidden_size=2048, + initializer_range=0.02, + intermediate_size=6144, + max_position_embeddings=262144, + num_attention_heads=16, + num_hidden_layers=28, + num_key_value_heads=8, + rms_norm_eps=1e-6, + rope_parameters={ + "mrope_interleaved": True, # type: ignore[dict-item] + "mrope_section": [24, 20, 20], # type: ignore[dict-item] + "rope_type": "default", # type: ignore[dict-item] + "rope_theta": 5_000_000, # type: ignore[dict-item] + }, + tie_word_embeddings=True, # type: ignore[call-arg] + use_cache=True, + vocab_size=151936, + ) + vision_config = Qwen3VLVisionConfig( + deepstack_visual_indexes=[5, 11, 17], + depth=24, + hidden_act="gelu_pytorch_tanh", + hidden_size=1024, + in_channels=3, + initializer_range=0.02, + intermediate_size=4096, + num_heads=16, + num_position_embeddings=2304, + out_hidden_size=2048, + patch_size=16, + spatial_merge_size=2, + temporal_patch_size=2, + ) + return Qwen3VLConfig( + architectures=["Qwen3VLForConditionalGeneration"], + image_token_id=151655, + text_config=text_config.to_dict(), + tie_word_embeddings=True, + video_token_id=151656, + vision_config=vision_config.to_dict(), + vision_end_token_id=151653, + vision_start_token_id=151652, + ) diff --git a/backends/arm/test/models/Qwen3_VL/test_qwen3_vl_layers.py b/backends/arm/test/models/Qwen3_VL/test_qwen3_vl_layers.py new file mode 100644 index 00000000000..77b2739167a --- /dev/null +++ b/backends/arm/test/models/Qwen3_VL/test_qwen3_vl_layers.py @@ -0,0 +1,479 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Tuple + +import torch +from executorch.backends.arm.test import common +from executorch.backends.arm.test.models.Qwen3_VL.qwen3_vl_test_config import ( + get_qwen3_vl_2b_instruct_checkpoint_config, +) +from executorch.backends.arm.test.tester.test_pipeline import ( + TosaPipelineFP, + VgfPipeline, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + apply_rotary_pos_emb, + apply_rotary_pos_emb_vision, + Qwen3VLTextAttention, + Qwen3VLTextDecoderLayer, + Qwen3VLTextMLP, + Qwen3VLTextRMSNorm, + Qwen3VLTextRotaryEmbedding, + Qwen3VLVisionAttention, + Qwen3VLVisionBlock, + Qwen3VLVisionMLP, + Qwen3VLVisionPatchEmbed, + Qwen3VLVisionPatchMerger, + Qwen3VLVisionRotaryEmbedding, +) + +input_t = Tuple[torch.Tensor, ...] + + +def _make_qwen3_vl_2b_instruct_layer_config(): + config = get_qwen3_vl_2b_instruct_checkpoint_config() + config.text_config._attn_implementation = "sdpa" + config.vision_config._attn_implementation = "sdpa" + return config + + +def _make_text_position_ids( + batch_size: int, seq_length: int, device: torch.device +) -> torch.Tensor: + return torch.arange(seq_length, device=device).unsqueeze(0).repeat(batch_size, 1) + + +def _make_causal_mask( + batch_size: int, seq_length: int, device: torch.device +) -> torch.Tensor: + mask = torch.full( + (seq_length, seq_length), torch.finfo(torch.float32).min, device=device + ) + mask = torch.triu(mask, diagonal=1) + return mask.unsqueeze(0).unsqueeze(0).repeat(batch_size, 1, 1, 1) + + +def _make_image_grid_thw(device: torch.device) -> torch.Tensor: + return torch.tensor([[1, 4, 4]], dtype=torch.long, device=device) + + +def _make_pixel_values(config, device: torch.device) -> torch.Tensor: + grid_thw = _make_image_grid_thw(device) + patch_volume = ( + config.vision_config.in_channels + * config.vision_config.temporal_patch_size + * config.vision_config.patch_size + * config.vision_config.patch_size + ) + num_patches = int(torch.prod(grid_thw[0]).item()) + return torch.randn(num_patches, patch_volume, device=device) + + +def _make_vision_position_embeddings( + config, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + grid_thw = _make_image_grid_thw(device) + num_patches = int(torch.prod(grid_thw[0]).item()) + head_dim = config.vision_config.hidden_size // config.vision_config.num_heads + return ( + torch.randn(num_patches, head_dim, device=device), + torch.randn(num_patches, head_dim, device=device), + ) + + +def _make_vision_cu_seqlens(device: torch.device) -> torch.Tensor: + grid_thw = _make_image_grid_thw(device) + num_patches = int(torch.prod(grid_thw[0]).item()) + return torch.tensor([0, num_patches], dtype=torch.int32, device=device) + + +class Qwen3VLTestModule(torch.nn.Module): + @classmethod + def prepare_model_and_inputs(cls): + raise NotImplementedError + + +class Qwen3VLVisionMLPModel(Qwen3VLTestModule): + def __init__(self, config) -> None: + super().__init__() + self.mlp = Qwen3VLVisionMLP(config.vision_config) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.mlp(hidden_states) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + model = cls(config).eval() + hidden_states = torch.randn(16, config.vision_config.hidden_size) + return model, (hidden_states,) + + +class VisionPatchEmbedModel(Qwen3VLTestModule): + def __init__(self, config) -> None: + super().__init__() + self.patch_embed = Qwen3VLVisionPatchEmbed(config.vision_config) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + return self.patch_embed(pixel_values) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + model = cls(config).eval() + return model, (_make_pixel_values(config, torch.device("cpu")),) + + +class VisionRotaryEmbeddingModel(Qwen3VLTestModule): + def __init__(self, config) -> None: + super().__init__() + head_dim = config.vision_config.hidden_size // config.vision_config.num_heads + self.rotary = Qwen3VLVisionRotaryEmbedding(head_dim // 2) + + def forward(self, max_hw: int) -> torch.Tensor: + return self.rotary(max_hw) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + grid_thw = _make_image_grid_thw(torch.device("cpu")) + max_hw = int(grid_thw[:, 1:].max().item()) + model = cls(config).eval() + return model, (max_hw,) + + +class VisionRotaryApplyModel(Qwen3VLTestModule): + def forward( + self, q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> torch.Tensor: + q_embed, k_embed = apply_rotary_pos_emb_vision(q, k, cos, sin) + return q_embed + k_embed + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + model = cls().eval() + cos, sin = _make_vision_position_embeddings(config, torch.device("cpu")) + head_dim = config.vision_config.hidden_size // config.vision_config.num_heads + q = torch.randn(cos.shape[0], config.vision_config.num_heads, head_dim) + k = torch.randn(cos.shape[0], config.vision_config.num_heads, head_dim) + return model, (q, k, cos, sin) + + +class VisionAttentionModel(Qwen3VLTestModule): + def __init__(self, config) -> None: + super().__init__() + self.attn = Qwen3VLVisionAttention(config.vision_config) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> torch.Tensor: + return self.attn( + hidden_states, + cu_seqlens=cu_seqlens, + position_embeddings=(cos, sin), + ) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + model = cls(config).eval() + hidden_states = model.attn.qkv.weight.new_empty( + 16, config.vision_config.hidden_size + ).normal_() + cos, sin = _make_vision_position_embeddings(config, hidden_states.device) + cu_seqlens = _make_vision_cu_seqlens(hidden_states.device) + return model, (hidden_states, cu_seqlens, cos, sin) + + +class VisionBlockModel(Qwen3VLTestModule): + def __init__(self, config) -> None: + super().__init__() + self.block = Qwen3VLVisionBlock(config.vision_config) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> torch.Tensor: + return self.block( + hidden_states, + cu_seqlens=cu_seqlens, + position_embeddings=(cos, sin), + ) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + model = cls(config).eval() + hidden_states = torch.randn(16, config.vision_config.hidden_size) + cos, sin = _make_vision_position_embeddings(config, hidden_states.device) + cu_seqlens = _make_vision_cu_seqlens(hidden_states.device) + return model, (hidden_states, cu_seqlens, cos, sin) + + +class VisionPatchMergerModel(Qwen3VLTestModule): + def __init__(self, config) -> None: + super().__init__() + self.merger = Qwen3VLVisionPatchMerger(config.vision_config) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.merger(hidden_states) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + model = cls(config).eval() + hidden_states = torch.randn( + config.vision_config.spatial_merge_size**2, + config.vision_config.hidden_size, + ) + return model, (hidden_states,) + + +class TextRMSNormModel(Qwen3VLTestModule): + def __init__(self, config) -> None: + super().__init__() + self.norm = Qwen3VLTextRMSNorm( + config.text_config.hidden_size, eps=config.text_config.rms_norm_eps + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.norm(hidden_states) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + model = cls(config).eval() + hidden_states = torch.randn(2, 8, config.text_config.hidden_size) + return model, (hidden_states,) + + +class TextRotaryEmbeddingModel(Qwen3VLTestModule): + def __init__(self, config) -> None: + super().__init__() + self.rotary = Qwen3VLTextRotaryEmbedding(config.text_config) + + def forward( + self, hidden_states: torch.Tensor, position_ids: torch.Tensor + ) -> torch.Tensor: + cos, sin = self.rotary(hidden_states, position_ids) + return cos + sin + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + model = cls(config).eval() + hidden_states = torch.randn(2, 8, config.text_config.hidden_size) + position_ids = _make_text_position_ids(2, 8, hidden_states.device) + return model, (hidden_states, position_ids) + + +class TextRotaryApplyModel(Qwen3VLTestModule): + def forward( + self, q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> torch.Tensor: + q_embed, k_embed = apply_rotary_pos_emb(q, k, cos, sin) + return q_embed.mean(dim=1) + k_embed.mean(dim=1) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + model = cls().eval() + hidden_states = torch.randn(2, 8, config.text_config.hidden_size) + position_ids = _make_text_position_ids(2, 8, hidden_states.device) + cos, sin = Qwen3VLTextRotaryEmbedding(config.text_config)( + hidden_states, position_ids + ) + q = torch.randn( + 2, + config.text_config.num_attention_heads, + 8, + config.text_config.head_dim, + ) + k = torch.randn( + 2, + config.text_config.num_key_value_heads, + 8, + config.text_config.head_dim, + ) + return model, (q, k, cos, sin) + + +class TextAttentionModel(Qwen3VLTestModule): + def __init__(self, config) -> None: + super().__init__() + self.attn = Qwen3VLTextAttention(config.text_config, layer_idx=0) + self.rotary = Qwen3VLTextRotaryEmbedding(config.text_config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + cos, sin = self.rotary(hidden_states, position_ids) + attn_output, _ = self.attn( + hidden_states=hidden_states, + position_embeddings=(cos, sin), + attention_mask=attention_mask, + ) + return attn_output + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + model = cls(config).eval() + hidden_states = torch.randn(2, 8, config.text_config.hidden_size) + attention_mask = _make_causal_mask(2, 8, hidden_states.device) + position_ids = _make_text_position_ids(2, 8, hidden_states.device) + return model, (hidden_states, attention_mask, position_ids) + + +class QKNormModel(Qwen3VLTestModule): + def __init__(self, config) -> None: + super().__init__() + self.attn = Qwen3VLTextAttention(config.text_config, layer_idx=0) + + def forward(self, q_states: torch.Tensor, k_states: torch.Tensor) -> torch.Tensor: + q_states = self.attn.q_norm(q_states) + k_states = self.attn.k_norm(k_states) + return q_states.mean(dim=(-1, -2)) + k_states.mean(dim=(-1, -2)) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + model = cls(config).eval() + q_states = torch.randn( + 2, + 8, + config.text_config.num_attention_heads, + config.text_config.head_dim, + ) + k_states = torch.randn( + 2, + 8, + config.text_config.num_key_value_heads, + config.text_config.head_dim, + ) + return model, (q_states, k_states) + + +class TextMLPModel(Qwen3VLTestModule): + def __init__(self, config) -> None: + super().__init__() + self.mlp = Qwen3VLTextMLP(config.text_config) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.mlp(hidden_states) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + model = cls(config).eval() + hidden_states = torch.randn(2, 8, config.text_config.hidden_size) + return model, (hidden_states,) + + +class TextDecoderLayerModel(Qwen3VLTestModule): + def __init__(self, config) -> None: + super().__init__() + self.layer = Qwen3VLTextDecoderLayer(config.text_config, layer_idx=0) + self.rotary = Qwen3VLTextRotaryEmbedding(config.text_config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + cos, sin = self.rotary(hidden_states, position_ids) + return self.layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + position_embeddings=(cos, sin), + ) + + @classmethod + def prepare_model_and_inputs(cls): + config = _make_qwen3_vl_2b_instruct_layer_config() + model = cls(config).eval() + hidden_states = torch.randn(2, 8, config.text_config.hidden_size) + attention_mask = _make_causal_mask(2, 8, hidden_states.device) + position_ids = _make_text_position_ids(2, 8, hidden_states.device) + return model, (hidden_states, attention_mask, position_ids) + + +@dataclass(frozen=True) +class Qwen3VLTestCase: + model_cls: type[Qwen3VLTestModule] + transform_passes: tuple = field(default_factory=tuple) + + +TOSA_FP_TEST_CASES: dict[str, Qwen3VLTestCase] = { + "vision_mlp": Qwen3VLTestCase(model_cls=Qwen3VLVisionMLPModel), + "vision_patch_embed": Qwen3VLTestCase(model_cls=VisionPatchEmbedModel), + "vision_rotary_embedding": Qwen3VLTestCase(model_cls=VisionRotaryEmbeddingModel), + "vision_rotary_apply": Qwen3VLTestCase(model_cls=VisionRotaryApplyModel), + "vision_attention": Qwen3VLTestCase(model_cls=VisionAttentionModel), + "vision_block": Qwen3VLTestCase(model_cls=VisionBlockModel), + "vision_patch_merger": Qwen3VLTestCase(model_cls=VisionPatchMergerModel), + "text_rms_norm": Qwen3VLTestCase(model_cls=TextRMSNormModel), + "text_rotary_embedding": Qwen3VLTestCase(model_cls=TextRotaryEmbeddingModel), + "text_rotary_apply": Qwen3VLTestCase(model_cls=TextRotaryApplyModel), + "text_attention": Qwen3VLTestCase(model_cls=TextAttentionModel), + "qk_norm": Qwen3VLTestCase(model_cls=QKNormModel), + "text_mlp": Qwen3VLTestCase(model_cls=TextMLPModel), + "text_decoder_layer": Qwen3VLTestCase(model_cls=TextDecoderLayerModel), +} + +VGF_NO_QUANT_TEST_CASES: dict[str, Qwen3VLTestCase] = TOSA_FP_TEST_CASES + + +@common.parametrize( + "test_case", + TOSA_FP_TEST_CASES, +) +def test_qwen3_vl_tosa_FP(test_case: Qwen3VLTestCase): + model, inputs = test_case.model_cls.prepare_model_and_inputs() + with torch.no_grad(): + pipeline = TosaPipelineFP[input_t]( + model, + inputs, + aten_op=[], + exir_op=[], + transform_passes=list(test_case.transform_passes), + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +@common.parametrize( + "test_case", + VGF_NO_QUANT_TEST_CASES, +) +def test_qwen3_vl_vgf_no_quant(test_case: Qwen3VLTestCase): + model, inputs = test_case.model_cls.prepare_model_and_inputs() + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + model, + inputs, + aten_op=[], + exir_op=[], + quantize=False, + transform_passes=list(test_case.transform_passes), + ) + pipeline.run() diff --git a/backends/arm/test/models/test_llama.py b/backends/arm/test/models/test_llama.py index 704e3e07926..1602aa7b4ba 100644 --- a/backends/arm/test/models/test_llama.py +++ b/backends/arm/test/models/test_llama.py @@ -31,7 +31,11 @@ from executorch.extension.llm.export.config.llm_config import LlmConfig +from transformers import GenerationConfig, LlamaConfig, LlamaForCausalLM +from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM + input_t = Tuple[torch.Tensor] +input_th = Tuple[torch.Tensor, torch.Tensor] # Add project dir to sys path to workaround importlib.import_module() conditions in model_factory.py this_files_dir = os.path.dirname(os.path.abspath(__file__)) @@ -41,6 +45,22 @@ logger = logging.getLogger(__name__) +class HFPositionalAdapter(torch.nn.Module): + def __init__(self, exportable): + super().__init__() + self.inner = exportable + + def forward(self, input_ids, cache_position): + # HF StaticCache eager path requires int64 index tensors, but keeping + # cache_position as int32 during export capture avoids adding an extra + # int64->int32 cast node in the lowered graph. + if torch._dynamo.is_compiling(): + cp = cache_position + else: + cp = cache_position.to(torch.long) + return self.inner(input_ids=input_ids, cache_position=cp) + + class TestLlama: """Test class of Llama models. @@ -51,6 +71,44 @@ class TestLlama: """ + def prepare_model_hf_static(self): + """ + Build a tiny HF LLaMA wrapped with TorchExportableModuleForDecoderOnlyLM (StaticCache) + See https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/executorch.py#L214C17-L214C53 + """ + # Tiny config + cfg = LlamaConfig( + vocab_size=32000, + hidden_size=256, + intermediate_size=512, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + use_cache=True, + ) + base = LlamaForCausalLM(cfg).eval() + + # REQUIRED: generation_config must request a 'static' cache with batch_size & max_cache_len + base.generation_config = GenerationConfig( + use_cache=True, + cache_implementation="static", + cache_config={"batch_size": 1, "max_cache_len": 128}, + ) + + exportable = TorchExportableModuleForDecoderOnlyLM( + model=base, batch_size=1, max_cache_len=128 + ) + + # Positional adapter so the pipeline can call module(*inputs) + model_for_pipeline = HFPositionalAdapter(exportable).eval() + + # The tester will call model(*inputs). Provide (input_ids, cache_position) + input_ids = torch.tensor([[0]], dtype=torch.long) # shape [1, 1] + cache_position = torch.tensor([0], dtype=torch.int32) # shape [1] + inputs = (input_ids, cache_position) + + return model_for_pipeline, inputs, None + def prepare_model(self): checkpoint = None params_file = None @@ -86,6 +144,10 @@ def prepare_model(self): # TODO: Enable key value cache args = [ "--disable_dynamic_shape", + "--max_seq_length", + "4096", + "--max_context_length", + "4096", "-c", checkpoint, "-p", @@ -93,6 +155,7 @@ def prepare_model(self): "--model", model_name, ] + parser = build_args_parser() args = parser.parse_args(args) llm_config = LlmConfig.from_args(args) @@ -123,11 +186,10 @@ def test_llama_tosa_FP(): aten_op=[], exir_op=[], custom_path="llama_tosa_fb", - run_on_tosa_ref_model=False, # Just want to write TOSA FB to disk + run_on_tosa_ref_model=True, use_to_edge_transform_and_lower=True, transform_passes=[InsertInt32CastsAfterInt64PlaceholdersPass()], ) - pipeline.add_stage_after("to_executorch", pipeline.tester.serialize) pipeline.run() @@ -144,12 +206,36 @@ def test_llama_tosa_INT(): aten_op=[], exir_op=[], custom_path="llama_tosa_fb_int", - run_on_tosa_ref_model=False, # Just want to write TOSA FB to disk + run_on_tosa_ref_model=True, use_to_edge_transform_and_lower=True, frobenius_threshold=None, cosine_threshold=None, ) - pipeline.add_stage_after("to_executorch", pipeline.tester.serialize) + pipeline.run() + + +def test_llama_tosa_INT_static(): + llama_model, llama_inputs, _ = TestLlama().prepare_model_hf_static() + if llama_model is None or llama_inputs is None: + pytest.skip("Missing model and/or input files") + + with torch.no_grad(): + pipeline = TosaPipelineINT[input_th]( + llama_model, + llama_inputs, + aten_op=[], + exir_op=[], + custom_path="llama_tosa_hf_static_int", + run_on_tosa_ref_model=True, + use_to_edge_transform_and_lower=True, + fold_quantize=True, + ) + # NOTE: HF StaticCache INT currently keeps two delegated subgraphs + # after partitioning on this path, so expect two delegate calls in EXIR. + pipeline.change_args( + "check_count.exir", + {"torch.ops.higher_order.executorch_call_delegate": 2}, + ) pipeline.run() diff --git a/backends/arm/test/models/test_mobilenet_v3_arm.py b/backends/arm/test/models/test_mobilenet_v3_arm.py index eccdc839e62..da9f99010b1 100644 --- a/backends/arm/test/models/test_mobilenet_v3_arm.py +++ b/backends/arm/test/models/test_mobilenet_v3_arm.py @@ -47,14 +47,17 @@ def test_mv3_tosa_FP(): @pytest.mark.slow def test_mv3_tosa_FP_fp16(): - inputs_fp16 = tuple(t.to(torch.float16) for t in model_inputs) + input_tensor_fp16 = torch.rand( + 1, 3, 232, 232, generator=torch.Generator().manual_seed(0) + ) + inputs_fp16 = (normalize(input_tensor_fp16).to(torch.float16),) pipeline = TosaPipelineFP[input_t]( mv3_fp16, inputs_fp16, aten_op=[], exir_op=[], use_to_edge_transform_and_lower=True, - atol=5e-2, + atol=6e-2, ) pipeline.run() diff --git a/backends/arm/test/models/test_swin2sr_arm.py b/backends/arm/test/models/test_swin2sr_arm.py new file mode 100644 index 00000000000..6bf9b2a18d5 --- /dev/null +++ b/backends/arm/test/models/test_swin2sr_arm.py @@ -0,0 +1,118 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Tuple + +import torch + +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.test_pipeline import ( + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) +from transformers import Swin2SRConfig, Swin2SRForImageSuperResolution + +input_t = Tuple[torch.Tensor] + +exir_ops = [ + "executorch_exir_dialects_edge__ops_aten_add_Tensor", + "executorch_exir_dialects_edge__ops_aten_convolution_default", + "executorch_exir_dialects_edge__ops_aten_layer_norm_default", + "executorch_exir_dialects_edge__ops_aten_matmul_default", + "executorch_exir_dialects_edge__ops_aten_mul_Tensor", + "executorch_exir_dialects_edge__ops_aten_pixel_shuffle_default", + "executorch_exir_dialects_edge__ops_aten_softmax_int", +] + + +class TinySwin2SR(torch.nn.Module): + def __init__(self): + super().__init__() + config = Swin2SRConfig( + image_size=8, + patch_size=1, + num_channels=3, + embed_dim=16, + depths=[1, 1], + num_heads=[1, 1], + window_size=4, + upscale=2, + img_range=1.0, + resi_connection="1conv", + upsampler="pixelshuffle", + ) + self.model = Swin2SRForImageSuperResolution(config).eval() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.model(pixel_values=x, return_dict=True).reconstruction + + +def make_model_and_inputs() -> tuple[torch.nn.Module, input_t]: + model = TinySwin2SR().eval() + inputs = (torch.rand(1, 3, 8, 8),) + return model, inputs + + +def test_swin2sr_tosa_FP(): + model, model_inputs = make_model_and_inputs() + pipeline = TosaPipelineFP[input_t]( + model, + model_inputs, + aten_op=[], + exir_op=exir_ops, + use_to_edge_transform_and_lower=True, + ) + pipeline.pop_stage("check_count.exir") + # TODO: MLETORCH-2134 re-enable once Swin2SR runs on the TOSA ref model. + pipeline.pop_stage("run_method_and_compare_outputs") + pipeline.run() + + +def test_swin2sr_tosa_INT(): + model, model_inputs = make_model_and_inputs() + pipeline = TosaPipelineINT[input_t]( + model, + model_inputs, + aten_op=[], + exir_op=exir_ops, + use_to_edge_transform_and_lower=True, + ) + pipeline.pop_stage("check_count.exir") + # TODO: MLETORCH-2134 re-enable once Swin2SR runs on the TOSA ref model. + pipeline.pop_stage("run_method_and_compare_outputs") + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_swin2sr_vgf_quant(): + model, model_inputs = make_model_and_inputs() + pipeline = VgfPipeline[input_t]( + model, + model_inputs, + aten_op=[], + exir_op=exir_ops, + use_to_edge_transform_and_lower=True, + quantize=True, + ) + pipeline.pop_stage("check_count.exir") + # TODO: MLETORCH-2134 re-enable once Swin2SR runs on the TOSA ref model. + pipeline.pop_stage("run_method_and_compare_outputs") + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_swin2sr_vgf_no_quant(): + model, model_inputs = make_model_and_inputs() + pipeline = VgfPipeline[input_t]( + model, + model_inputs, + aten_op=[], + exir_op=exir_ops, + use_to_edge_transform_and_lower=True, + quantize=False, + ) + pipeline.pop_stage("check_count.exir") + pipeline.run() diff --git a/backends/arm/test/ops/test_bmm.py b/backends/arm/test/ops/test_bmm.py index d21b33cfec4..a1aaa736d41 100644 --- a/backends/arm/test/ops/test_bmm.py +++ b/backends/arm/test/ops/test_bmm.py @@ -1,4 +1,4 @@ -# Copyright 2024-2025 Arm Limited and/or its affiliates. +# Copyright 2024-2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -10,11 +10,12 @@ import torch +from executorch.backends.arm.quantizer import get_symmetric_a16w8_quantization_config from executorch.backends.arm.test import common - from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, + OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, VgfPipeline, @@ -23,9 +24,6 @@ aten_op_bmm = "torch.ops.aten.bmm.default" exir_op_bmm = "executorch_exir_dialects_edge__ops_aten_bmm_default" -aten_op_mm = "torch.ops.aten.matmul.default" -exir_op_mm = "executorch_exir_dialects_edge__ops_aten_matmul_default" - input_t1 = Tuple[torch.Tensor, torch.Tensor] # Input x @@ -191,3 +189,52 @@ def test_bmm_vgf_quant_single_input(test_data: input_t1): quantize=True, ) pipeline.run() + + +a16w8_bmm_test_parameters = { + "rand_same": lambda: (torch.rand(2, 1, 1), torch.rand(2, 1, 1)), + "rand_diff": lambda: (torch.rand(5, 3, 5), torch.rand(5, 5, 2)), + "rand_rect": lambda: (torch.rand(1, 55, 3), torch.rand(1, 3, 44)), + "rand_batch10": lambda: (torch.rand(10, 1, 10), torch.rand(10, 10, 5)), + "rand_neg": lambda: ( + -10 * torch.randn(2, 32, 64), + 5 + 5 * torch.randn(2, 64, 32), + ), +} + + +@common.parametrize("test_data", a16w8_bmm_test_parameters) +@common.XfailIfNoCorstone300 +def test_bmm_a16w8_u55_INT(test_data: input_t1): + """U55 does not support bmm with INT16 inputs. + + Verify bmm is rejected. + + """ + pipeline = OpNotSupportedPipeline[input_t1]( + BMM(), + test_data(), + non_delegated_ops={exir_op_bmm: 1}, + n_expected_delegates=0, + u55_subset=True, + quantize=True, + tosa_extensions=["int16"], + ) + pipeline.quantizer.set_global(get_symmetric_a16w8_quantization_config()) + pipeline.run() + + +@common.parametrize("test_data", a16w8_bmm_test_parameters) +@common.XfailIfNoCorstone320 +def test_bmm_a16w8_u85_INT(test_data: input_t1): + pipeline = EthosU85PipelineINT[input_t1]( + BMM(), + test_data(), + aten_op_bmm, + exir_op_bmm, + a16w8_quantization=True, + symmetric_io_quantization=True, + qtol=1, + epsilon=2**-16, + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_conv1d.py b/backends/arm/test/ops/test_conv1d.py index 486b6b3ce7c..5fcdcf50465 100644 --- a/backends/arm/test/ops/test_conv1d.py +++ b/backends/arm/test/ops/test_conv1d.py @@ -399,3 +399,35 @@ def test_convolution_1d_vgf_quant_a8w4(test_data): get_symmetric_a8w4_quantization_config(is_per_channel=per_channel_quantization) ) pipeline.run() + + +@common.parametrize("test_data", test_data_INT) +@common.XfailIfNoCorstone300 +def test_conv1d_a16w8_u55_INT(test_data): + model, per_channel_quantization = test_data() + pipeline = EthosU55PipelineINT[input_t]( + model, + model.get_inputs(), + aten_op, + exir_op, + a16w8_quantization=True, + symmetric_io_quantization=True, + per_channel_quantization=per_channel_quantization, + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_INT) +@common.XfailIfNoCorstone320 +def test_conv1d_a16w8_u85_INT(test_data): + model, per_channel_quantization = test_data() + pipeline = EthosU85PipelineINT[input_t]( + model, + model.get_inputs(), + aten_op, + exir_op, + a16w8_quantization=True, + symmetric_io_quantization=True, + per_channel_quantization=per_channel_quantization, + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_conv2d.py b/backends/arm/test/ops/test_conv2d.py index 310471a966c..fdb625f5580 100644 --- a/backends/arm/test/ops/test_conv2d.py +++ b/backends/arm/test/ops/test_conv2d.py @@ -124,299 +124,316 @@ def forward(self, x): return x -conv2d_2x2_3x2x14x14_nobias = Conv2d( - in_channels=2, - out_channels=3, - kernel_size=(2, 2), - stride=1, - bias=False, - padding=3, - width=14, - height=14, - batches=2, - padding_mode="circular", -) +def conv2d_2x2_3x2x14x14_nobias(): + return Conv2d( + in_channels=2, + out_channels=3, + kernel_size=(2, 2), + stride=1, + bias=False, + padding=3, + width=14, + height=14, + batches=2, + padding_mode="circular", + ) -conv2d_3x3_1x3x12x12_st1_pd1_reflect = Conv2d( - in_channels=3, - out_channels=4, - kernel_size=(3, 3), - stride=1, - padding=3, - width=12, - height=12, - batches=1, - padding_mode="reflect", -) -conv2d_3x3_1x3x12x12_st1_pd1_replicate = Conv2d( - in_channels=3, - out_channels=4, - kernel_size=(3, 3), - stride=1, - padding=3, - width=12, - height=12, - batches=1, - padding_mode="replicate", -) +def conv2d_3x3_1x3x12x12_st1_pd3_reflect(): + return Conv2d( + in_channels=3, + out_channels=4, + kernel_size=(3, 3), + stride=1, + padding=3, + width=12, + height=12, + batches=1, + padding_mode="reflect", + ) -conv2d_3x3_1x3x24x24_st1 = Conv2d( - in_channels=3, - out_channels=10, - kernel_size=(3, 3), - stride=1, - padding=0, - width=24, - height=24, - batches=1, -) -conv2d_3x3_1x3x12x12_st2_pd1 = Conv2d( - in_channels=3, - out_channels=4, - kernel_size=(3, 3), - stride=2, - padding=1, - width=12, - height=12, - batches=1, -) +def conv2d_3x3_1x3x12x12_st1_pd3_replicate(): + return Conv2d( + in_channels=3, + out_channels=4, + kernel_size=(3, 3), + stride=1, + padding=3, + width=12, + height=12, + batches=1, + padding_mode="replicate", + ) -conv2d_1x1_1x2x16x16_st1 = Conv2d( - in_channels=2, - out_channels=1, - kernel_size=(1, 1), - stride=1, - padding=0, - width=16, - height=16, - batches=1, -) -conv2d_2x2_2x1x14x13_st2 = Conv2d( - in_channels=2, - out_channels=1, - kernel_size=(2, 2), - stride=2, - padding=0, - width=14, - height=13, - batches=1, -) +def conv2d_3x3_1x3x24x24_st1(): + return Conv2d( + in_channels=3, + out_channels=10, + kernel_size=(3, 3), + stride=1, + padding=0, + width=24, + height=24, + batches=1, + ) -conv2d_5x5_3x2x24x24_st1 = Conv2d( - in_channels=2, - out_channels=3, - kernel_size=(5, 5), - stride=1, - padding=0, - width=24, - height=24, - batches=2, -) -conv2d_3x3_1x3x28x28_st2_pd1 = Conv2d( - in_channels=3, - out_channels=16, - kernel_size=(3, 3), - stride=2, - padding=1, - width=28, - height=28, - batches=1, -) +def conv2d_3x3_1x3x12x12_st2_pd1(): + return Conv2d( + in_channels=3, + out_channels=4, + kernel_size=(3, 3), + stride=2, + padding=1, + width=12, + height=12, + batches=1, + ) -conv2d_5x5_1x3x14x15_st3_pd1 = Conv2d( - in_channels=3, - out_channels=16, - kernel_size=(5, 5), - stride=3, - padding=1, - width=14, - height=15, - batches=1, -) -conv2d_7x7_1x3x16x16_st2_pd1_dl2 = Conv2d( - in_channels=3, - out_channels=3, - kernel_size=(7, 7), - stride=2, - padding=1, - dilation=2, - width=16, - height=16, - batches=1, -) +def conv2d_1x1_1x2x16x16_st1(): + return Conv2d( + in_channels=2, + out_channels=1, + kernel_size=(1, 1), + stride=1, + padding=0, + width=16, + height=16, + batches=1, + ) -conv2d_7x7_1x3x15x15_st1_pd0_dl1 = Conv2d( - in_channels=3, - out_channels=3, - kernel_size=(7, 7), - stride=1, - padding=0, - dilation=1, - width=15, - height=15, - batches=1, -) -conv2d_5x5_1x3x14x14_st5_pd0_dl1 = Conv2d( - in_channels=3, - out_channels=3, - kernel_size=(5, 5), - stride=5, - padding=0, - dilation=1, - width=14, - height=14, - batches=1, -) +def conv2d_2x2_2x1x14x13_st2(): + return Conv2d( + in_channels=2, + out_channels=1, + kernel_size=(2, 2), + stride=2, + padding=0, + width=14, + height=13, + batches=1, + ) -conv2d_5x5_1x3x9x9_st5_pd0_dl1 = Conv2d( - in_channels=3, - out_channels=3, - kernel_size=(5, 5), - stride=5, - padding=0, - dilation=1, - width=9, - height=9, - batches=1, -) -conv2d_3x3_1x3x8x9_st3_pd0_dl1 = Conv2d( - in_channels=3, - out_channels=3, - kernel_size=(3, 3), - stride=3, - padding=0, - dilation=1, - width=8, - height=9, - batches=1, -) +def conv2d_5x5_3x2x24x24_st1(): + return Conv2d( + in_channels=2, + out_channels=3, + kernel_size=(5, 5), + stride=1, + padding=0, + width=24, + height=24, + batches=2, + ) -conv2d_3x3_1x3x9x8_st3_pd0_dl1 = Conv2d( - in_channels=3, - out_channels=3, - kernel_size=(3, 3), - stride=3, - padding=0, - dilation=1, - width=8, - height=9, - batches=1, -) -conv2d_3x4_1x3x7x7_st3_pd0_dl1 = Conv2d( - in_channels=3, - out_channels=3, - kernel_size=(3, 4), - stride=3, - padding=0, - dilation=1, - width=7, - height=7, - batches=1, -) +def conv2d_3x3_1x3x28x28_st2_pd1(): + return Conv2d( + in_channels=3, + out_channels=16, + kernel_size=(3, 3), + stride=2, + padding=1, + width=28, + height=28, + batches=1, + ) -conv2d_4x3_1x3x7x7_st3_pd0_dl1 = Conv2d( - in_channels=3, - out_channels=3, - kernel_size=(4, 3), - stride=3, - padding=0, - dilation=1, - width=7, - height=7, - batches=1, -) -two_conv2d_nobias = Conv2d( - nbr_conv=2, - width=32, - height=32, - in_channels=[3, 10], - out_channels=[10, 15], - kernel_size=[(5, 5), (5, 5)], - stride=[1, 1], - padding=[0, 0], - bias=[False, False], - batches=1, -) +def conv2d_5x5_1x3x14x15_st3_pd1(): + return Conv2d( + in_channels=3, + out_channels=16, + kernel_size=(5, 5), + stride=3, + padding=1, + width=14, + height=15, + batches=1, + ) -two_conv2d = Conv2d( - nbr_conv=2, - width=32, - height=32, - in_channels=[3, 10], - out_channels=[10, 15], - kernel_size=[(5, 5), (5, 5)], - stride=[1, 1], - padding=[0, 0], - bias=[True, True], - batches=1, -) -conv2d_groups = Conv2d( - in_channels=12, - out_channels=9, - kernel_size=(3, 3), - stride=1, - padding=0, - dilation=1, - width=7, - height=7, - batches=1, - groups=3, - bias=False, -) +def conv2d_7x7_1x3x16x16_st2_pd1_dl2(): + return Conv2d( + in_channels=3, + out_channels=3, + kernel_size=(7, 7), + stride=2, + padding=1, + dilation=2, + width=16, + height=16, + batches=1, + ) -conv2d_groups_bias = Conv2d( - in_channels=15, - out_channels=5, - kernel_size=(3, 3), - stride=1, - padding=0, - dilation=1, - width=7, - height=7, - batches=1, - groups=5, - bias=True, -) -# Shenanigan to get a nicer output when test fails. With unittest it looks like: -# FAIL: test_convolution_2d_tosa_INT_2_3x3_1x3x12x12_st2_pd1 -test_data_FP = { - "2x2_3x2x14x14_nobias_circular": lambda: conv2d_2x2_3x2x14x14_nobias, - "3x3_1x3x12x12_st1_pd1_reflect": lambda: conv2d_3x3_1x3x12x12_st1_pd1_reflect, - "3x3_1x3x12x12_st1_pd1_replicate": lambda: conv2d_3x3_1x3x12x12_st1_pd1_replicate, - "3x3_1x3x24x24_st1": lambda: conv2d_3x3_1x3x24x24_st1, - "3x3_1x3x12x12_st2_pd1": lambda: conv2d_3x3_1x3x12x12_st2_pd1, - "1x1_1x2x16x16_st1": lambda: conv2d_1x1_1x2x16x16_st1, - "2x2_2x1x14x13_st2_needs_adjust_pass": lambda: conv2d_2x2_2x1x14x13_st2, - "5x5_1x3x14x15_st3_pd1_needs_adjust_pass": lambda: conv2d_5x5_1x3x14x15_st3_pd1, - "7x7_1x3x16x16_st2_pd1_dl2_needs_adjust_pass": lambda: conv2d_7x7_1x3x16x16_st2_pd1_dl2, - "7x7_1x3x15x15_st1_pd0_dl1_needs_adjust_pass": lambda: conv2d_7x7_1x3x15x15_st1_pd0_dl1, - "5x5_1x3x14x14_st5_pd0_dl1_needs_adjust_pass": lambda: conv2d_5x5_1x3x14x14_st5_pd0_dl1, - "5x5_1x3x9x9_st5_pd0_dl1_needs_adjust_pass": lambda: conv2d_5x5_1x3x9x9_st5_pd0_dl1, - "3x3_1x3x9x8_st3_pd0_dl1_needs_adjust_pass": lambda: conv2d_3x3_1x3x9x8_st3_pd0_dl1, - "3x3_1x3x8x9_st3_pd0_dl1_needs_adjust_pass": lambda: conv2d_3x3_1x3x8x9_st3_pd0_dl1, - "3x4_1x3x7x7_st3_pd0_dl1_needs_adjust_pass": lambda: conv2d_3x4_1x3x7x7_st3_pd0_dl1, - "4x3_1x3x7x7_st3_pd0_dl1_needs_adjust_pass": lambda: conv2d_4x3_1x3x7x7_st3_pd0_dl1, - "5x5_3x2x24x24_st1": lambda: conv2d_5x5_3x2x24x24_st1, - "3x3_1x3x28x28_st2_pd1": lambda: conv2d_3x3_1x3x28x28_st2_pd1, - "two_conv2d_nobias": lambda: two_conv2d_nobias, - "two_conv2d": lambda: two_conv2d, - "groups": lambda: conv2d_groups, - "groups_bias": lambda: conv2d_groups_bias, -} +def conv2d_7x7_1x3x15x15_st1_pd0_dl1(): + return Conv2d( + in_channels=3, + out_channels=3, + kernel_size=(7, 7), + stride=1, + padding=0, + dilation=1, + width=15, + height=15, + batches=1, + ) -test_data_FP_bf16 = { - "bf16_3x3": lambda: Conv2d( + +def conv2d_5x5_1x3x14x14_st5_pd0_dl1(): + return Conv2d( + in_channels=3, + out_channels=3, + kernel_size=(5, 5), + stride=5, + padding=0, + dilation=1, + width=14, + height=14, + batches=1, + ) + + +def conv2d_5x5_1x3x9x9_st5_pd0_dl1(): + return Conv2d( + in_channels=3, + out_channels=3, + kernel_size=(5, 5), + stride=5, + padding=0, + dilation=1, + width=9, + height=9, + batches=1, + ) + + +def conv2d_3x3_1x3x8x9_st3_pd0_dl1(): + return Conv2d( + in_channels=3, + out_channels=3, + kernel_size=(3, 3), + stride=3, + padding=0, + dilation=1, + width=8, + height=9, + batches=1, + ) + + +def conv2d_3x3_1x3x9x8_st3_pd0_dl1(): + return Conv2d( + in_channels=3, + out_channels=3, + kernel_size=(3, 3), + stride=3, + padding=0, + dilation=1, + width=9, + height=8, + batches=1, + ) + + +def conv2d_3x4_1x3x7x7_st3_pd0_dl1(): + return Conv2d( + in_channels=3, + out_channels=3, + kernel_size=(3, 4), + stride=3, + padding=0, + dilation=1, + width=7, + height=7, + batches=1, + ) + + +def conv2d_4x3_1x3x7x7_st3_pd0_dl1(): + return Conv2d( + in_channels=3, + out_channels=3, + kernel_size=(4, 3), + stride=3, + padding=0, + dilation=1, + width=7, + height=7, + batches=1, + ) + + +def two_conv2d_nobias(): + return Conv2d( + nbr_conv=2, + width=32, + height=32, + in_channels=[3, 10], + out_channels=[10, 15], + kernel_size=[(5, 5), (5, 5)], + stride=[1, 1], + padding=[0, 0], + bias=[False, False], + batches=1, + ) + + +def two_conv2d(): + return Conv2d( + nbr_conv=2, + width=32, + height=32, + in_channels=[3, 10], + out_channels=[10, 15], + kernel_size=[(5, 5), (5, 5)], + stride=[1, 1], + padding=[0, 0], + bias=[True, True], + batches=1, + ) + + +def conv2d_groups(): + return Conv2d( + in_channels=12, + out_channels=9, + kernel_size=(3, 3), + stride=1, + padding=0, + dilation=1, + width=7, + height=7, + batches=1, + groups=3, + bias=False, + ) + + +def conv2d_groups_bias(): + return Conv2d( + in_channels=15, + out_channels=5, + kernel_size=(3, 3), + stride=1, + padding=0, + dilation=1, + width=7, + height=7, + batches=1, + groups=5, + bias=True, + ) + + +def conv2d_bf16_3x3(): + return Conv2d( height=12, width=12, in_channels=3, @@ -426,8 +443,11 @@ def forward(self, x): padding=(1, 1), bias=True, dtype=torch.bfloat16, - ), - "bf16_1x1": lambda: Conv2d( + ) + + +def conv2d_bf16_1x1(): + return Conv2d( height=8, width=8, in_channels=2, @@ -437,10 +457,11 @@ def forward(self, x): padding=(0, 3), bias=False, dtype=torch.bfloat16, - ), -} -test_data_FP_fp16 = { - "fp16_3x3": lambda: Conv2d( + ) + + +def conv2d_fp16_3x3(): + return Conv2d( height=12, width=12, in_channels=3, @@ -450,8 +471,11 @@ def forward(self, x): padding=(1, 1), bias=True, dtype=torch.float16, - ), - "fp16_1x1": lambda: Conv2d( + ) + + +def conv2d_fp16_1x1(): + return Conv2d( height=8, width=8, in_channels=2, @@ -461,7 +485,43 @@ def forward(self, x): padding=(0, 3), bias=False, dtype=torch.float16, - ), + ) + + +# Shenanigan to get a nicer output when test fails. With unittest it looks like: +# FAIL: test_convolution_2d_tosa_INT_2_3x3_1x3x12x12_st2_pd1 +test_data_FP = { + "2x2_3x2x14x14_nobias_circular": conv2d_2x2_3x2x14x14_nobias, + "3x3_1x3x12x12_st1_pd3_reflect": conv2d_3x3_1x3x12x12_st1_pd3_reflect, + "3x3_1x3x12x12_st1_pd3_replicate": conv2d_3x3_1x3x12x12_st1_pd3_replicate, + "3x3_1x3x24x24_st1": conv2d_3x3_1x3x24x24_st1, + "3x3_1x3x12x12_st2_pd1": conv2d_3x3_1x3x12x12_st2_pd1, + "1x1_1x2x16x16_st1": conv2d_1x1_1x2x16x16_st1, + "2x2_2x1x14x13_st2_needs_adjust_pass": conv2d_2x2_2x1x14x13_st2, + "5x5_1x3x14x15_st3_pd1_needs_adjust_pass": conv2d_5x5_1x3x14x15_st3_pd1, + "7x7_1x3x16x16_st2_pd1_dl2_needs_adjust_pass": conv2d_7x7_1x3x16x16_st2_pd1_dl2, + "7x7_1x3x15x15_st1_pd0_dl1_needs_adjust_pass": conv2d_7x7_1x3x15x15_st1_pd0_dl1, + "5x5_1x3x14x14_st5_pd0_dl1_needs_adjust_pass": conv2d_5x5_1x3x14x14_st5_pd0_dl1, + "5x5_1x3x9x9_st5_pd0_dl1_needs_adjust_pass": conv2d_5x5_1x3x9x9_st5_pd0_dl1, + "3x3_1x3x9x8_st3_pd0_dl1_needs_adjust_pass": conv2d_3x3_1x3x9x8_st3_pd0_dl1, + "3x3_1x3x8x9_st3_pd0_dl1_needs_adjust_pass": conv2d_3x3_1x3x8x9_st3_pd0_dl1, + "3x4_1x3x7x7_st3_pd0_dl1_needs_adjust_pass": conv2d_3x4_1x3x7x7_st3_pd0_dl1, + "4x3_1x3x7x7_st3_pd0_dl1_needs_adjust_pass": conv2d_4x3_1x3x7x7_st3_pd0_dl1, + "5x5_3x2x24x24_st1": conv2d_5x5_3x2x24x24_st1, + "3x3_1x3x28x28_st2_pd1": conv2d_3x3_1x3x28x28_st2_pd1, + "two_conv2d_nobias": two_conv2d_nobias, + "two_conv2d": two_conv2d, + "groups": conv2d_groups, + "groups_bias": conv2d_groups_bias, +} + +test_data_FP_bf16 = { + "bf16_3x3": conv2d_bf16_3x3, + "bf16_1x1": conv2d_bf16_1x1, +} +test_data_FP_fp16 = { + "fp16_3x3": conv2d_fp16_3x3, + "fp16_1x1": conv2d_fp16_1x1, } # Generate a new test set paired with per_channel_quant=True/False. diff --git a/backends/arm/test/ops/test_gelu.py b/backends/arm/test/ops/test_gelu.py index 82e3cb1e53e..84626c70b4d 100644 --- a/backends/arm/test/ops/test_gelu.py +++ b/backends/arm/test/ops/test_gelu.py @@ -6,6 +6,7 @@ from typing import Tuple import torch + from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, @@ -176,3 +177,37 @@ def test_gelu_vgf_quant(test_data: input_t1): quantize=True, ) pipeline.run() + + +@common.parametrize("test_data", Gelu.test_data) +@common.XfailIfNoCorstone300 +def test_gelu_a16w8_u55_INT(test_data: input_t1): + approximate, data = test_data() + pipeline = EthosU55PipelineINT[input_t1]( + Gelu(approximate), + (data,), + Gelu.aten_op, + Gelu.exir_op, + a16w8_quantization=True, + symmetric_io_quantization=True, + qtol=128, + epsilon=2**-16, + ) + pipeline.run() + + +@common.parametrize("test_data", Gelu.test_data) +@common.XfailIfNoCorstone320 +def test_gelu_a16w8_u85_INT(test_data: input_t1): + approximate, data = test_data() + pipeline = EthosU85PipelineINT[input_t1]( + Gelu(approximate), + (data,), + Gelu.aten_op, + Gelu.exir_op, + a16w8_quantization=True, + symmetric_io_quantization=True, + qtol=128, + epsilon=2**-16, + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_mean_dim.py b/backends/arm/test/ops/test_mean_dim.py index 2b60dc0211f..a5931fc3cbe 100644 --- a/backends/arm/test/ops/test_mean_dim.py +++ b/backends/arm/test/ops/test_mean_dim.py @@ -6,6 +6,7 @@ # LICENSE file in the root directory of this source tree. from typing import Callable +import pytest import torch from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( @@ -393,3 +394,68 @@ def test_mean_tosa_INT(test_data): symmetric_io_quantization=True, ) pipeline.run() + + +a16w8_mean_test_parameters = { + "rank_2_keepdim": lambda: (torch.rand(7, 3), (0, 1), True), + "rank_3_keepdim": lambda: (torch.rand(5, 7, 3), (0, 1, 2), True), + "rand_23_keepdim": lambda: (torch.rand(1, 5, 7, 3), (2, 3), True), + "rand_0123_keepdim": lambda: (torch.rand(1, 5, 7, 3), (0, 1, 2, 3), True), + "rand_none_keepdim": lambda: (torch.rand(1, 5, 7, 3), None, True), + "rank_2": lambda: (torch.rand(5, 7), (-2, -1), False), + "rand_123": lambda: (torch.rand(1, 5, 7, 3), (-3, -2, -1), False), +} + +a16w8_mean_test_parameters_u85_xfails = { + "rank_1_keepdim": lambda: (torch.rand(7), 0, True), + "rand_1_keepdim": lambda: (torch.rand(1, 5, 7, 3), (1), True), + "rank_1": lambda: (torch.rand(7), (-1), False), + "rand_3": lambda: (torch.rand(1, 5, 7, 3), (-1), False), +} + + +@common.parametrize( + "test_data", {**a16w8_mean_test_parameters, **a16w8_mean_test_parameters_u85_xfails} +) +@common.XfailIfNoCorstone300 +def test_mean_dim_a16w8_u55_INT(test_data): + test_data, dim, keep_dim = test_data() + pipeline = EthosU55PipelineINT[input_t]( + MeanDim(dim, keep_dim), + (test_data,), + [], + symmetric_io_quantization=True, + a16w8_quantization=True, + ) + pipeline.run() + + +@common.parametrize("test_data", a16w8_mean_test_parameters) +@common.XfailIfNoCorstone320 +def test_mean_dim_a16w8_u85_INT(test_data): + test_data, dim, keep_dim = test_data() + pipeline = EthosU85PipelineINT[input_t]( + MeanDim(dim, keep_dim), + (test_data,), + [], + symmetric_io_quantization=True, + a16w8_quantization=True, + ) + pipeline.run() + + +@common.parametrize("test_data", a16w8_mean_test_parameters_u85_xfails) +@pytest.mark.xfail( + reason="U85 mean_dim a16w8 produces all-zero output for single-dim reductions" +) +@common.XfailIfNoCorstone320 +def test_mean_dim_a16w8_u85_INT_xfail(test_data): + test_data, dim, keep_dim = test_data() + pipeline = EthosU85PipelineINT[input_t]( + MeanDim(dim, keep_dim), + (test_data,), + [], + symmetric_io_quantization=True, + a16w8_quantization=True, + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_pow.py b/backends/arm/test/ops/test_pow.py index 2d007fa7e68..6d304ce0627 100644 --- a/backends/arm/test/ops/test_pow.py +++ b/backends/arm/test/ops/test_pow.py @@ -147,7 +147,6 @@ def test_pow_tensor_tensor_vgf_no_quant(test_data: Pow_TensorTensor.input_t): x_fail_FP = { "exp_two": "TOSA constraints: If x <0 .", - "exp_zero": "MLETORCH-2041 : Invalid inputs.", } diff --git a/backends/arm/test/ops/test_split.py b/backends/arm/test/ops/test_split.py index 6af3362de7a..f655785410c 100644 --- a/backends/arm/test/ops/test_split.py +++ b/backends/arm/test/ops/test_split.py @@ -310,3 +310,38 @@ def test_split_tensor_vgf_quant(test_data: Tuple): quantize=True, ) pipeline.run() + + +a16w8_split_test_parameters = { + "a16w8_1d_split_2": lambda: (torch.rand(10), 2, 0), + "a16w8_2d_split_4": lambda: (torch.rand(8, 4), 4, 0), + "a16w8_3d_split_4": lambda: (torch.rand(4, 4, 8), 4, 2), +} + + +@common.parametrize("test_data", a16w8_split_test_parameters) +@common.XfailIfNoCorstone300 +def test_split_a16w8_u55_INT(test_data: input_t1): + pipeline = EthosU55PipelineINT[input_t1]( + Split(), + test_data(), + aten_ops=[], + exir_ops=exir_op, + a16w8_quantization=True, + symmetric_io_quantization=True, + ) + pipeline.run() + + +@common.parametrize("test_data", a16w8_split_test_parameters) +@common.XfailIfNoCorstone320 +def test_split_a16w8_u85_INT(test_data: input_t1): + pipeline = EthosU85PipelineINT[input_t1]( + Split(), + test_data(), + aten_ops=[], + exir_ops=exir_op, + a16w8_quantization=True, + symmetric_io_quantization=True, + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_squeeze.py b/backends/arm/test/ops/test_squeeze.py index fcef851024b..5d18e8c695c 100644 --- a/backends/arm/test/ops/test_squeeze.py +++ b/backends/arm/test/ops/test_squeeze.py @@ -1,4 +1,4 @@ -# Copyright 2024-2025 Arm Limited and/or its affiliates. +# Copyright 2024-2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -16,6 +16,7 @@ from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, + OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, VgfPipeline, @@ -60,6 +61,12 @@ def forward(self, x: torch.Tensor): return x.squeeze() +unsupported_cases = { + "squeeze_dim_no_effect": lambda: (torch.randn(3, 4, 5), 1), + "squeeze_no_effect": lambda: (torch.randn(3, 4, 5),), +} + + ############## ## Squeeze ### ############## @@ -137,6 +144,16 @@ def test_squeeze_dim_vgf_quant(test_data: Tuple): pipeline.run() +def test_squeeze_no_target_not_delegated() -> None: + pipeline = OpNotSupportedPipeline[input_t1]( + Squeeze(), + unsupported_cases["squeeze_no_effect"](), + {"executorch_exir_dialects_edge__ops_aten_squeeze_copy_dims": 1}, + n_expected_delegates=0, + ) + pipeline.run() + + ################# ## SqueezeDim ### ################# @@ -214,6 +231,16 @@ def test_squeeze_dim_vgf_quant_2(test_data: Tuple): pipeline.run() +def test_squeeze_dim_no_target_not_delegated() -> None: + pipeline = OpNotSupportedPipeline[Tuple[torch.Tensor, int]]( + SqueezeDim(), + unsupported_cases["squeeze_dim_no_effect"](), + {"executorch_exir_dialects_edge__ops_aten_squeeze_copy_dims": 1}, + n_expected_delegates=0, + ) + pipeline.run() + + ################## ## SqueezeDims ### ################## diff --git a/backends/arm/test/ops/test_upsample_bilinear2d.py b/backends/arm/test/ops/test_upsample_bilinear2d.py index f084e0ebe14..705f9ba4d20 100644 --- a/backends/arm/test/ops/test_upsample_bilinear2d.py +++ b/backends/arm/test/ops/test_upsample_bilinear2d.py @@ -551,12 +551,21 @@ def test_upsample_bilinear2d_vec_vgf_no_quant_Upsample(test_data: torch.Tensor): @common.SkipIfNoModelConverter def test_upsample_bilinear2d_vec_vgf_no_quant_Interpolate(test_data: torch.Tensor): data, size, scale_factor, compare = test_data() + match data.dtype: + case torch.float16: + atol = 1e-2 + rtol = 1e-2 + case _: + atol = 1e-3 + rtol = 1e-3 pipeline = VgfPipeline[input_t1]( Interpolate(size, scale_factor), (data,), aten_op, exir_op, quantize=False, + atol=atol, + rtol=rtol, ) if not compare: pipeline.pop_stage(-1) diff --git a/backends/arm/test/ops/test_var.py b/backends/arm/test/ops/test_var.py index a7943bfc19b..74ceb4b557d 100644 --- a/backends/arm/test/ops/test_var.py +++ b/backends/arm/test/ops/test_var.py @@ -32,6 +32,11 @@ class Var(torch.nn.Module): ), } + test_parameters_ethosu = { + "var_4d_keep_dim_0_correction": lambda: (torch.randn(1, 50, 10, 20), True, 0), + "var_4d_keep_dim_1_correction": lambda: (torch.randn(1, 30, 15, 20), True, 1), + } + def __init__(self, keepdim: bool = True, correction: int = 0): super().__init__() self.keepdim = keepdim @@ -141,6 +146,18 @@ def forward( return x.var(dim=self.dim, keepdim=self.keepdim, correction=self.correction) +# Var parameters that pass on Ethos-U hardware (keepdim=True only) +var_test_parameters_hw = { + "var_4d_keep_dim_0_correction": lambda: (torch.randn(1, 50, 10, 20), True, 0), + "var_4d_keep_dim_1_correction": lambda: (torch.randn(1, 30, 15, 20), True, 1), +} + +a16w8_var_test_parameters = { + "var_4d_keep_dim_0_correction": lambda: (torch.randn(1, 50, 10, 20), True, 0), + "var_4d_keep_dim_1_correction": lambda: (torch.randn(1, 30, 15, 20), True, 1), +} + + ########## ## Var ### ########## @@ -170,7 +187,7 @@ def test_var_dim_tosa_INT_no_dim(test_data: Tuple): pipeline.run() -@common.parametrize("test_data", Var.test_parameters) +@common.parametrize("test_data", var_test_parameters_hw) @common.XfailIfNoCorstone300 def test_var_dim_u55_INT_no_dim(test_data: Tuple): test_data, keepdim, correction = test_data() @@ -183,7 +200,7 @@ def test_var_dim_u55_INT_no_dim(test_data: Tuple): pipeline.run() -@common.parametrize("test_data", Var.test_parameters) +@common.parametrize("test_data", var_test_parameters_hw) @common.XfailIfNoCorstone320 def test_var_dim_u85_INT_no_dim(test_data: Tuple): test_data, keepdim, correction = test_data() @@ -224,6 +241,36 @@ def test_var_dim_vgf_quant_no_dim(test_data: Tuple): pipeline.run() +@common.parametrize("test_data", a16w8_var_test_parameters) +@common.XfailIfNoCorstone300 +def test_var_a16w8_u55_INT(test_data: Tuple): + test_data, keepdim, correction = test_data() + pipeline = EthosU55PipelineINT[input_t1]( + Var(keepdim, correction), + (test_data,), + aten_ops=[], + exir_ops=[], + a16w8_quantization=True, + symmetric_io_quantization=True, + ) + pipeline.run() + + +@common.parametrize("test_data", a16w8_var_test_parameters) +@common.XfailIfNoCorstone320 +def test_var_a16w8_u85_INT(test_data: Tuple): + test_data, keepdim, correction = test_data() + pipeline = EthosU85PipelineINT[input_t1]( + Var(keepdim, correction), + (test_data,), + aten_ops=[], + exir_ops=[], + a16w8_quantization=True, + symmetric_io_quantization=True, + ) + pipeline.run() + + ############# ## VarDim ### ############# diff --git a/backends/arm/test/passes/test_decompose_int_pow_pass.py b/backends/arm/test/passes/test_decompose_int_pow_pass.py index 6846392f248..7761c031e2c 100644 --- a/backends/arm/test/passes/test_decompose_int_pow_pass.py +++ b/backends/arm/test/passes/test_decompose_int_pow_pass.py @@ -59,18 +59,18 @@ def get_inputs(self) -> input_t: def test_decompose_int_pow_tosa_FP(data: TestParam) -> None: module_with_inputs, nbr_muls = data module = cast(torch.nn.Module, module_with_inputs) + pow_op = "executorch_exir_dialects_edge__ops_aten_pow_Tensor_Scalar" pipeline = PassPipeline[input_t]( module, module_with_inputs.get_inputs(), quantize=False, ops_before_pass={ - "executorch_exir_dialects_edge__ops_aten_pow_Tensor_Scalar": 1, + pow_op: 1, }, ops_not_before_pass=[], ops_after_pass={ "executorch_exir_dialects_edge__ops_aten_mul_Tensor": nbr_muls, }, - ops_not_after_pass=["executorch_exir_dialects_edge__ops_pow_Tensor_Scalar"], pass_list=[DecomposeIntPowPass], ) pipeline.run() diff --git a/backends/arm/test/quantizer/test_selective_quantization.py b/backends/arm/test/quantizer/test_selective_quantization.py index a59a509ce06..7efa590680a 100644 --- a/backends/arm/test/quantizer/test_selective_quantization.py +++ b/backends/arm/test/quantizer/test_selective_quantization.py @@ -7,6 +7,7 @@ from typing import Dict import torch + from executorch.backends.arm.quantizer import ( get_symmetric_a16w8_quantization_config, get_symmetric_quantization_config, @@ -16,13 +17,18 @@ from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import QuantizationPipeline from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.cortex_m.test.tester import ramp_tensor +from executorch.backends.test.harness.stages import StageType +from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY from torchvision import models, transforms # type: ignore[import-untyped] from torchvision.ops.misc import Conv2dNormActivation # type: ignore[import-untyped] -def get_quantizer(): +def get_quantizer(use_composable_quantizer: bool = False): tosa_spec = TosaSpecification.create_from_string("TOSA-1.0+INT") - quantizer = TOSAQuantizer(tosa_spec) + quantizer = TOSAQuantizer( + tosa_spec, use_composable_quantizer=use_composable_quantizer + ) quantizer.set_global(get_symmetric_quantization_config()) return quantizer @@ -53,6 +59,25 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return x + y +class Cat(torch.nn.Module): + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return torch.cat((x, y), dim=1) + + +class LinearGraphTail(torch.nn.Module): + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 10) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.linear(x) + x = torch.relu(x) + x = torch.sigmoid(x) + return torch.neg(x) + + class AddSoftmaxAdd(torch.nn.Module): module_names = {"add_0": None, "add_1": None} module_types = { @@ -131,11 +156,94 @@ def test_selective_quant_module_type_tosa_INT(model): pipeline.run() +def test_selective_quant_cat_node_target_none_tosa_INT(): + model = Cat() + inputs = (torch.randn(1, 2, 4), torch.randn(1, 3, 4)) + + quantizer = get_quantizer(use_composable_quantizer=True) + quantizer.set_node_target(torch.ops.aten.cat.default, None) + + pipeline = QuantizationPipeline[tuple[torch.Tensor, torch.Tensor]]( + model, + inputs, + quantizer=quantizer, + qspecs={ + "aten.cat.default": { + None: 1, + }, + }, + ) + + pipeline.run() + + +def test_composable_io_none_skips_global_tosa_INT(): + model = Add() + inputs = (torch.randn(1, 10), torch.randn(1, 10)) + + quantizer = get_quantizer(use_composable_quantizer=True) + quantizer.set_io(None) + + pipeline = QuantizationPipeline[tuple[torch.Tensor, torch.Tensor]]( + model, + inputs, + quantizer=quantizer, + input_qspecs={None: 2}, + output_qspecs={None: 1}, + ) + + pipeline.run() + + +def test_composable_global_none_linear_graph_tail_tosa_INT(): + model = LinearGraphTail() + inputs = (torch.randn(1, 10),) + + quantizer = get_quantizer(use_composable_quantizer=True) + quantizer.set_global(None) + + pipeline = QuantizationPipeline[tuple[torch.Tensor]]( + model, + inputs, + quantizer=quantizer, + qspecs={ + "aten.linear.default": {None: 1}, + "aten.relu.default": {None: 1}, + "aten.sigmoid.default": {None: 1}, + "aten.neg.default": {None: 1}, + }, + ) + + pipeline.run() + + graph = pipeline.tester.get_graph(StageType.QUANTIZE) + unannotated_nodes = [ + node.name + for node in graph.nodes + if node.op == "call_function" and Q_ANNOTATION_KEY not in node.meta + ] + assert not unannotated_nodes + + mv3 = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights) mv3.eval() normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) +class SharedBufferEmbeddingLinearConstantFold(torch.nn.Module): + + def __init__(self): + super().__init__() + self.shared = torch.nn.Embedding(4, 4) + self.lm_head = torch.nn.Linear(4, 4, bias=False) + self.lm_head.weight = self.shared.weight + + def forward(self, ids, x): + y0 = self.shared(ids).sum(dim=1) + z = self.lm_head(x) + return y0 + z + + def test_mv3_selective_quant_int16_tosa_INT(): model = mv3 inputs = (normalize(torch.randn(1, 3, 224, 224)),) @@ -209,3 +317,33 @@ def test_mv3_io_quant_tosa_INT(): ) pipeline.run() + + +def test_multiple_folded_get_attr(): + """In torchao/quantization/pt2e/constant_fold.py:constant_fold, get_attr + node targets are deleted as soon as there is one get_attr node w/o users + using the target. + + If there are multiple get_attr nodes refering the same target such as in + this test, the function crashes if no workaround is present. + + """ + + model = SharedBufferEmbeddingLinearConstantFold() + example_inputs = ( + torch.tensor([[0, 1]], dtype=torch.long), + ramp_tensor(-2, 2, (1, 4)), + ) + + quantizer = get_quantizer() + quantizer.set_module_type(torch.nn.Embedding, None) + + pipeline = QuantizationPipeline( + model, + example_inputs, + quantizer=quantizer, + qspecs=None, + input_qspecs=None, + output_qspecs=None, + ) + pipeline.run() diff --git a/backends/arm/test/runner_utils.py b/backends/arm/test/runner_utils.py index 914a95f0c8d..93887fbda6b 100644 --- a/backends/arm/test/runner_utils.py +++ b/backends/arm/test/runner_utils.py @@ -12,6 +12,7 @@ import subprocess # nosec B404 - invoked only for trusted toolchain binaries import sys import tempfile +from collections.abc import Iterable from pathlib import Path from types import NoneType @@ -848,39 +849,98 @@ def vkml_emulation_layer_installed() -> bool: return layers_exists and deploy_exists -def assert_elf_path_exists(elf_path): - if not os.path.exists(elf_path): - raise FileNotFoundError( - f"Did not find build arm_executor_runner or executor_runner in path {elf_path}, \ - run setup_testing.sh or setup_testing_vkml.sh?" - ) - - -def get_elf_path(target_board: str, use_portable_ops: bool = False) -> str: - elf_path = "" +def _elf_search_roots() -> list[Path]: + roots: list[Path] = [] + for env_var in ( + "EXECUTORCH_ROOT", + "GITHUB_WORKSPACE", + "BUILD_WORKSPACE_DIRECTORY", + ): + env_root = os.environ.get(env_var) + if env_root: + roots.append(Path(env_root).expanduser()) + + cwd = Path.cwd().resolve() + search_parents = [cwd, *cwd.parents, *Path(__file__).resolve().parents] + for parent in search_parents: + if (parent / "examples" / "arm").is_dir() or (parent / "arm_test").exists(): + roots.append(parent) + + unique_roots: list[Path] = [] + seen: set[Path] = set() + for root in roots: + resolved = root.resolve() + if resolved not in seen: + unique_roots.append(resolved) + seen.add(resolved) + return unique_roots + + +def _elf_path_candidates( + target_board: str, use_portable_ops: bool = False +) -> list[Path]: if target_board not in VALID_TARGET: raise ValueError(f"Unsupported target: {target_board}") - if use_portable_ops: - portable_ops_str = "portable-ops_" - else: - portable_ops_str = "" - + portable_ops_str = "portable-ops_" if use_portable_ops else "" if target_board in ("corstone-300", "corstone-320"): - elf_path = os.path.join( + build_dir = Path( "arm_test", f"arm_semihosting_executor_runner_{portable_ops_str}{target_board}", - "arm_executor_runner", ) - elif target_board == "vkml_emulation_layer": - elf_path = os.path.join( - f"arm_test/arm_executor_runner_{portable_ops_str}vkml", - "executor_runner", + binary_name = "arm_executor_runner" + else: + build_dir = Path("arm_test", f"arm_executor_runner_{portable_ops_str}vkml") + binary_name = "executor_runner" + + candidates: list[Path] = [] + for root in _elf_search_roots(): + root_build_dir = root / build_dir + candidates.extend( + [ + root_build_dir / binary_name, + root_build_dir / "Release" / binary_name, + root_build_dir / "examples" / "arm" / "executor_runner" / binary_name, + root_build_dir + / "examples" + / "arm" + / "executor_runner" + / "Release" + / binary_name, + ] ) - assert_elf_path_exists(elf_path) - return elf_path + unique_candidates: list[Path] = [] + seen: set[Path] = set() + for candidate in candidates: + resolved = candidate.resolve(strict=False) + if resolved not in seen: + unique_candidates.append(resolved) + seen.add(resolved) + return unique_candidates + + +def _resolve_existing_elf_path(elf_candidates: Iterable[Path]) -> Path: + checked: list[Path] = [] + for elf_path in elf_candidates: + checked.append(elf_path) + if elf_path.exists(): + return elf_path + + checked_paths = ", ".join(str(path) for path in checked) + raise FileNotFoundError( + "Did not find build arm_executor_runner or executor_runner. " + f"Tried: {checked_paths}. " + "Run setup_testing.sh or setup_testing_vkml.sh?" + ) + + +def get_elf_path(target_board: str, use_portable_ops: bool = False) -> str: + elf_path = _resolve_existing_elf_path( + _elf_path_candidates(target_board, use_portable_ops=use_portable_ops) + ) + return str(elf_path) def arm_executor_runner_exists(target_board: str, use_portable_ops: bool = False): diff --git a/backends/arm/test/targets.bzl b/backends/arm/test/targets.bzl index d53e5630808..0a3faa6a074 100644 --- a/backends/arm/test/targets.bzl +++ b/backends/arm/test/targets.bzl @@ -38,6 +38,12 @@ def define_arm_tests(): "ops/test_to_copy.py", "ops/test_exp.py", "ops/test_reciprocal.py", + "ops/test_mean_dim.py", + "ops/test_var.py", + "ops/test_conv1d.py", + "ops/test_gelu.py", + "ops/test_bmm.py", + "ops/test_split.py", ] # Quantization @@ -51,7 +57,7 @@ def define_arm_tests(): "misc/test_compile_spec.py", # "misc/test_evaluate_model.py", "misc/test_pass_pipeline_config.py", - "misc/test_tosa_dialect_resize.py", + "misc/tosa_dialect/test_tosa_resize.py", "misc/test_tosa_spec.py", "misc/test_bn_relu_folding_qat.py", "misc/test_custom_partition.py", diff --git a/backends/arm/test/test_arm_backend.sh b/backends/arm/test/test_arm_backend.sh new file mode 100755 index 00000000000..6a69e94b147 --- /dev/null +++ b/backends/arm/test/test_arm_backend.sh @@ -0,0 +1,359 @@ +#!/bin/bash +# Copyright 2024-2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +set -e + +script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) + +# Executorch root +et_root_dir=$(cd ${script_dir}/../../.. && pwd) +cd "${et_root_dir}" +pwd +scratch_dir=${et_root_dir}/examples/arm/arm-scratch +setup_path_script=${scratch_dir}/setup_path.sh +_setup_msg="please refer to ${et_root_dir}/examples/arm/setup.sh to properly install necessary tools." + + +TEST_SUITE=$1 + +# Source the tools +# This should be prepared by the setup.sh +[[ -f ${setup_path_script} ]] \ + || { echo "Missing ${setup_path_script}. ${_setup_msg}"; exit 1; } +source ${setup_path_script} + +help() { + echo "Usage:" + echo " $0 " + echo " where can be any of:" + # This will list all lines in this file that is starting with test_ remove () { and print it as a list. + # e,g, "test_pytest() { # Test ops and other things" -> test_pytest # Test ops and other things + echo "all # run all tests" + grep "^test_" $0 | sed 's/([^)]*)[[:space:]]*{*//g' + exit +} + +if [[ -z "${TEST_SUITE:-}" ]]; then + echo "Missing test suite name, exiting..." + help +else + echo "Run Arm backend test suite ${TEST_SUITE}" +fi + +TEST_SUITE_NAME="$(basename "$0") ${TEST_SUITE}" + +EXCLUDE_TARGET_EXPR="(not u55) and (not u85) and (not tosa) and (not _vgf_)" +PYTEST_RETRY_ARGS=(--reruns 2 --reruns-delay 1) + +all() { # Run all tests + # This will list all lines in this file that is starting with test_ remove () { and add this script name in + # front of it and execute it in a sub shell + # e.g. from this file: + # + # test_pytest() { # Test ops and other things + # bla bla bla + # } + # test_pytest_ethosu_fvp() { # Same as test_pytest but ... + # bla bla bla + # } + #... + # become a small script: + # ---- + # backends/arm/test/test_arm_backend.sh test_pytest # Test ops and other things + # backends/arm/test/test_arm_backend.sh test_pytest_ethosu_fvp # Same as test_pytest but ... + # ... + # ---- + # That is executed + echo "${TEST_SUITE_NAME}: Run all tests" + grep "^test_" backends/arm/test/test_arm_backend.sh | sed 's/([^)]*)[[:space:]]*{*//g' | sed "s|^|$0 |" | sh -e + echo "${TEST_SUITE_NAME}: PASS" +} + +# ------------------------------------------- +# -------- Non target-specific tests -------- +# ------------------------------------------- +test_pytest_ops_no_target() { + echo "${TEST_SUITE_NAME}: Run pytest ops for target-less tests" + + # Run arm baremetal pytest tests without target + pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=10 backends/arm/test/ --ignore=backends/arm/test/models -k "${EXCLUDE_TARGET_EXPR}" + echo "${TEST_SUITE_NAME}: PASS" +} + +test_pytest_models_no_target() { + echo "${TEST_SUITE_NAME}: Run pytest models for target-less tests" + + # Install model dependencies for pytest + source backends/arm/scripts/install_models_for_test.sh + + # Run arm baremetal pytest tests without FVP + pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=0 backends/arm/test/models -k "${EXCLUDE_TARGET_EXPR}" + echo "${TEST_SUITE_NAME}: PASS" +} + +# ------------------------------------- +# -------- TOSA specific tests -------- +# ------------------------------------- +test_pytest_ops_tosa() { + echo "${TEST_SUITE_NAME}: Run pytest ops for TOSA" + + pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=10 backends/arm/test/ --ignore=backends/arm/test/models -k tosa + echo "${TEST_SUITE_NAME}: PASS" +} + +test_pytest_models_tosa() { + echo "${TEST_SUITE_NAME}: Run pytest models for TOSA" + + # Install model dependencies for pytest + source backends/arm/scripts/install_models_for_test.sh + + pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=0 backends/arm/test/models -k tosa + echo "${TEST_SUITE_NAME}: PASS" +} + +test_run_tosa() { + echo "${TEST_SUITE_NAME}: Test TOSA delegate examples with run.sh" + + echo "${TEST_SUITE_NAME}: Test target TOSA" + examples/arm/run.sh --et_build_root=arm_test/test_run --target=TOSA-1.0+INT --model_name=add + examples/arm/run.sh --et_build_root=arm_test/test_run --target=TOSA-1.0+INT --model_name=mul + + echo "${TEST_SUITE_NAME}: PASS" +} + +# ---------------------------------------------- +# -------- Arm Ethos-U55 specific tests -------- +# ---------------------------------------------- +test_pytest_ops_ethos_u55() { + echo "${TEST_SUITE_NAME}: Run pytest ops for Arm Ethos-U55" + + backends/arm/scripts/build_executorch.sh + backends/arm/test/setup_testing.sh + + pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=10 backends/arm/test/ --ignore=backends/arm/test/models -k u55 + echo "${TEST_SUITE_NAME}: PASS" +} + +test_pytest_models_ethos_u55() { + echo "${TEST_SUITE_NAME}: Run pytest models for Arm Ethos-U55" + + backends/arm/scripts/build_executorch.sh + backends/arm/test/setup_testing.sh + + # Install model dependencies for pytest + source backends/arm/scripts/install_models_for_test.sh + + pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=0 backends/arm/test/models -k u55 + echo "${TEST_SUITE_NAME}: PASS" +} + +test_run_ethos_u55() { + echo "${TEST_SUITE_NAME}: Test ethos-u55 delegate examples with run.sh" + + echo "${TEST_SUITE_NAME}: Test target Ethos-U55" + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-64 --model_name=add + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --bundleio + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-256 --model_name=add --bundleio --etdump + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --etdump + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=mul + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --pte_placement=elf + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-256 --model_name=add --pte_placement=0x38000000 + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=mul --bundleio --pte_placement=elf + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=mul --bundleio --pte_placement=0x38000000 + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --bundleio --pte_placement=0x38000000 + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=examples/arm/example_modules/add.py + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=examples/arm/example_modules/add.py --bundleio + + # Cortex-M op tests + echo "${TEST_SUITE_NAME}: Test target Cortex-M55 (on Ethos-U55)" + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --bundleio --no_delegate --select_ops_list="aten::add.out" + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qadd --bundleio + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qops --bundleio + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qops --bundleio --no_delegate --select_ops_list="aten::sub.out,aten::add.out,aten::mul.out" + + echo "${TEST_SUITE_NAME}: PASS" +} + +# ---------------------------------------------- +# -------- Arm Ethos-U85 specific tests -------- +# ---------------------------------------------- +test_pytest_ops_ethos_u85() { + echo "${TEST_SUITE_NAME}: Run pytest ops for Arm Ethos-U85" + + backends/arm/scripts/build_executorch.sh + backends/arm/test/setup_testing.sh + + # Run arm baremetal pytest tests with FVP + pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=10 backends/arm/test/ --ignore=backends/arm/test/models -k u85 + echo "${TEST_SUITE_NAME}: PASS" +} + +test_pytest_models_ethos_u85() { + echo "${TEST_SUITE_NAME}: Run pytest models for Arm Ethos-U85" + + backends/arm/scripts/build_executorch.sh + backends/arm/test/setup_testing.sh + + # Install model dependencies for pytest + source backends/arm/scripts/install_models_for_test.sh + + pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=0 backends/arm/test/models -k u85 + echo "${TEST_SUITE_NAME}: PASS" +} + +test_run_ethos_u85() { + echo "${TEST_SUITE_NAME}: Test ethos-u85 delegate examples with run.sh" + + echo "${TEST_SUITE_NAME}: Test target Ethos-U85" + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=add + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-256 --model_name=add --bundleio + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-512 --model_name=add --bundleio --etdump + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-1024 --model_name=add --etdump + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-2048 --model_name=mul --pte_placement=elf + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=mul --pte_placement=0x38000000 + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=mul --bundleio --pte_placement=elf + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-256 --model_name=mul --bundleio --pte_placement=0x38000000 + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=examples/arm/example_modules/add.py + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-1024 --model_name=examples/arm/example_modules/add.py --bundleio + + # Cortex-M op tests + echo "${TEST_SUITE_NAME}: Test target Cortex-M55 (on Ethos-U85)" + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=qops --bundleio + + echo "${TEST_SUITE_NAME}: PASS" +} + +# ---------------------------------------------------------- +# -------- Vulkan Graph Format (VGF) specific tests -------- +# ---------------------------------------------------------- +test_pytest_ops_vkml() { + echo "${TEST_SUITE_NAME}: Run pytest operator tests with VKML runtime" + + source backends/arm/test/setup_testing_vkml.sh + + pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=10 backends/arm/test/ \ + --ignore=backends/arm/test/models -k _vgf_ + echo "${TEST_SUITE_NAME}: PASS" +} + +test_pytest_models_vkml() { + echo "${TEST_SUITE_NAME}: Run pytest model tests with VKML runtime" + + source backends/arm/test/setup_testing_vkml.sh + + # Install model dependencies for pytest + source backends/arm/scripts/install_models_for_test.sh + + pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=0 backends/arm/test/models -k _vgf_ + echo "${TEST_SUITE_NAME}: PASS" +} + +test_run_vkml() { + echo "${TEST_SUITE_NAME}: Test VKML delegate examples with run.sh" + + source backends/arm/test/setup_testing_vkml.sh + + echo "${TEST_SUITE_NAME}: Test VKML" + out_folder="arm_test/test_run" + vkml_build_dir="${build_root_test_dir}" + + examples/arm/run.sh --build-dir="${vkml_build_dir}" --et_build_root=${out_folder} --target=vgf --model_name=add --output=${out_folder}/runner + examples/arm/run.sh --build-dir="${vkml_build_dir}" --et_build_root=${out_folder} --target=vgf --model_name=mul --output=${out_folder}/runner + + examples/arm/run.sh --build-dir="${vkml_build_dir}" --et_build_root=${out_folder} --target=vgf --model_name=qadd --output=${out_folder}/runner + examples/arm/run.sh --build-dir="${vkml_build_dir}" --et_build_root=${out_folder} --target=vgf --model_name=qops --output=${out_folder}/runner + + echo "${TEST_SUITE_NAME}: PASS" +} + +# ------------------------------------ +# -------- Miscellaneous tests ------- +# ------------------------------------ +test_model_smollm2_135M() { + echo "${TEST_SUITE_NAME}: Test SmolLM2-135M on Ethos-U85" + + # Build common libs once + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --build_libs + + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u85-128 --model=smollm2 --extra_flags="-DEXECUTORCH_SELECT_OPS_LIST=dim_order_ops::_to_dim_order_copy.out" --specify_ethosu_scratch + + echo "${TEST_SUITE_NAME}: PASS" +} + +test_smaller_stories_llama() { + echo "${TEST_SUITE_NAME}: Test smaller_stories_llama" + + backends/arm/scripts/build_executorch.sh + + mkdir -p stories110M + pushd stories110M + wget -N https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.pt + echo '{"dim": 768, "multiple_of": 32, "n_heads": 12, "n_layers": 12, "norm_eps": 1e-05, "vocab_size": 32000}' > params.json + popd + + # Get path to source directory + pytest \ + -c /dev/null \ + "${PYTEST_RETRY_ARGS[@]}" \ + --verbose \ + --color=yes \ + --numprocesses=auto \ + --junit-xml=stories110M/test-reports/unittest.xml \ + backends/arm/test/models/test_llama.py \ + --llama_inputs stories110M/stories110M.pt stories110M/params.json stories110m + + echo "${TEST_SUITE_NAME}: PASS" +} + +test_memory_allocation() { + echo "${TEST_SUITE_NAME}: Test ethos-u memory allocation with run.sh" + + mkdir -p arm_test/test_run + # Ethos-U85 + echo "${TEST_SUITE_NAME}: Test target Ethos-U85" + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=examples/arm/example_modules/add.py &> arm_test/test_run/full.log + python3 backends/arm/test/test_memory_allocator_log.py --log arm_test/test_run/full.log \ + --require "model_pte_program_size" "<= 3200 B" \ + --require "method_allocator_planned" "<= 64 B" \ + --require "method_allocator_loaded" "<= 1024 B" \ + --require "method_allocator_input" "<= 16 B" \ + --require "Total DRAM used" "<= 0.06 KiB" + echo "${TEST_SUITE_NAME}: PASS" +} + +test_undefinedbehavior_sanitizer() { + echo "${TEST_SUITE_NAME}: Test ethos-u executor_runner with UBSAN" + + mkdir -p arm_test/test_run + # Ethos-U85 + echo "${TEST_SUITE_NAME}: Test target Ethos-U85" + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=examples/arm/example_modules/add.py --build_type=UndefinedSanitizer + echo "${TEST_SUITE_NAME}: PASS" +} + +test_address_sanitizer() { + echo "${TEST_SUITE_NAME}: Test ethos-u executor_runner with ASAN" + + mkdir -p arm_test/test_run + # Ethos-U85 + echo "${TEST_SUITE_NAME}: Test target Ethos-U85" + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=examples/arm/example_modules/add.py --build_type=AddressSanitizer + echo "${TEST_SUITE_NAME}: PASS" +} + + +if [[ -z "${TEST_SUITE}" ]]; then + echo "Missing test suite name." + exit 1 +elif [[ ! "${TEST_SUITE}" =~ ^[a-zA-Z_][a-zA-Z0-9_]*$ ]]; then + echo "Invalid test suite name: ${TEST_SUITE}" + exit 1 +elif ! declare -F -- "${TEST_SUITE}" > /dev/null; then + echo "Unknown test suite: ${TEST_SUITE}" + exit 1 +fi + +"${TEST_SUITE}" diff --git a/backends/arm/test/test_arm_baremetal.sh b/backends/arm/test/test_arm_baremetal.sh deleted file mode 100755 index ad8cd8b7d3a..00000000000 --- a/backends/arm/test/test_arm_baremetal.sh +++ /dev/null @@ -1,347 +0,0 @@ -#!/bin/bash -# Copyright 2024-2026 Arm Limited and/or its affiliates. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# TODO: Rename this script - -set -e - -script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) - -# Executorch root -et_root_dir=$(cd ${script_dir}/../../.. && pwd) -cd "${et_root_dir}" -pwd -scratch_dir=${et_root_dir}/examples/arm/arm-scratch -setup_path_script=${scratch_dir}/setup_path.sh -_setup_msg="please refer to ${et_root_dir}/examples/arm/setup.sh to properly install necessary tools." - - -TEST_SUITE=$1 - -# Source the tools -# This should be prepared by the setup.sh -[[ -f ${setup_path_script} ]] \ - || { echo "Missing ${setup_path_script}. ${_setup_msg}"; exit 1; } -source ${setup_path_script} - -help() { - echo "Usage:" - echo " $0 " - echo " where can be any of:" - # This will list all lines in this file that is starting with test_ remove () { and print it as a list. - # e,g, "test_pytest() { # Test ops and other things" -> test_pytest # Test ops and other things - echo "all # run all tests" - grep "^test_" $0 | sed 's/([^)]*)[[:space:]]*{*//g' - exit -} - -if [[ -z "${TEST_SUITE:-}" ]]; then - echo "Missing test suite name, exiting..." - help -else - echo "Run Arm baremetal test suite ${TEST_SUITE}" -fi - -TEST_SUITE_NAME="$(basename "$0") ${TEST_SUITE}" - -EXCLUDE_TARGET_EXPR="(not u55) and (not u85) and (not tosa) and (not _vgf_)" -PYTEST_RETRY_ARGS=(--reruns 2 --reruns-delay 1) - -all() { # Run all tests - # This will list all lines in this file that is starting with test_ remove () { and add this script name in - # front of it and execute it in a sub shell - # e.g. from this file: - # - # test_pytest() { # Test ops and other things - # bla bla bla - # } - # test_pytest_ethosu_fvp() { # Same as test_pytest but ... - # bla bla bla - # } - #... - # become a small script: - # ---- - # backends/arm/test/test_arm_baremetal.sh test_pytest # Test ops and other things - # backends/arm/test/test_arm_baremetal.sh test_pytest_ethosu_fvp # Same as test_pytest but ... - # ... - # ---- - # That is executed - echo "${TEST_SUITE_NAME}: Run all tests" - grep "^test_" backends/arm/test/test_arm_baremetal.sh | sed 's/([^)]*)[[:space:]]*{*//g' | sed "s|^|$0 |" | sh - echo "${TEST_SUITE_NAME}: PASS" -} - -# ------------------------------------------- -# -------- Non target-specific tests -------- -# ------------------------------------------- -test_pytest_ops_no_target() { - echo "${TEST_SUITE_NAME}: Run pytest ops for target-less tests" - - # Run arm baremetal pytest tests without target - pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=10 backends/arm/test/ --ignore=backends/arm/test/models -k "${EXCLUDE_TARGET_EXPR}" - echo "${TEST_SUITE_NAME}: PASS" -} - -test_pytest_models_no_target() { - echo "${TEST_SUITE_NAME}: Run pytest models for target-less tests" - - # Install model dependencies for pytest - source backends/arm/scripts/install_models_for_test.sh - - # Run arm baremetal pytest tests without FVP - pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=0 backends/arm/test/models -k "${EXCLUDE_TARGET_EXPR}" - echo "${TEST_SUITE_NAME}: PASS" -} - -# ------------------------------------- -# -------- TOSA specific tests -------- -# ------------------------------------- -test_pytest_ops_tosa() { - echo "${TEST_SUITE_NAME}: Run pytest ops for TOSA" - - pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=10 backends/arm/test/ --ignore=backends/arm/test/models -k tosa - echo "${TEST_SUITE_NAME}: PASS" -} - -test_pytest_models_tosa() { - echo "${TEST_SUITE_NAME}: Run pytest models for TOSA" - - # Install model dependencies for pytest - source backends/arm/scripts/install_models_for_test.sh - - pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=0 backends/arm/test/models -k tosa - echo "${TEST_SUITE_NAME}: PASS" -} - -test_run_tosa() { - echo "${TEST_SUITE_NAME}: Test TOSA delegate examples with run.sh" - - echo "${TEST_SUITE_NAME}: Test target TOSA" - examples/arm/run.sh --et_build_root=arm_test/test_run --target=TOSA-1.0+INT --model_name=add - examples/arm/run.sh --et_build_root=arm_test/test_run --target=TOSA-1.0+INT --model_name=mul - - echo "${TEST_SUITE_NAME}: PASS" -} - -# ---------------------------------------------- -# -------- Arm Ethos-U55 specific tests -------- -# ---------------------------------------------- -test_pytest_ops_ethos_u55() { - echo "${TEST_SUITE_NAME}: Run pytest ops for Arm Ethos-U55" - - backends/arm/scripts/build_executorch.sh - backends/arm/test/setup_testing.sh - - pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=10 backends/arm/test/ --ignore=backends/arm/test/models -k u55 - echo "${TEST_SUITE_NAME}: PASS" -} - -test_pytest_models_ethos_u55() { - echo "${TEST_SUITE_NAME}: Run pytest models for Arm Ethos-U55" - - backends/arm/scripts/build_executorch.sh - backends/arm/test/setup_testing.sh - - # Install model dependencies for pytest - source backends/arm/scripts/install_models_for_test.sh - - pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=0 backends/arm/test/models -k u55 - echo "${TEST_SUITE_NAME}: PASS" -} - -test_run_ethos_u55() { - echo "${TEST_SUITE_NAME}: Test ethos-u55 delegate examples with run.sh" - - echo "${TEST_SUITE_NAME}: Test target Ethos-U55" - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-64 --model_name=add - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --bundleio - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-256 --model_name=add --bundleio --etdump - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --etdump - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=mul - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --pte_placement=elf - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-256 --model_name=add --pte_placement=0x38000000 - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=mul --bundleio --pte_placement=elf - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=mul --bundleio --pte_placement=0x38000000 - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --bundleio --pte_placement=0x38000000 - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=examples/arm/example_modules/add.py - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=examples/arm/example_modules/add.py --bundleio - - # Cortex-M op tests - echo "${TEST_SUITE_NAME}: Test target Cortex-M55 (on Ethos-U55)" - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --bundleio --no_delegate --select_ops_list="aten::add.out" - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qadd --bundleio - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qops --bundleio - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qops --bundleio --no_delegate --select_ops_list="aten::sub.out,aten::add.out,aten::mul.out" - - echo "${TEST_SUITE_NAME}: PASS" -} - -# ---------------------------------------------- -# -------- Arm Ethos-U85 specific tests -------- -# ---------------------------------------------- -test_pytest_ops_ethos_u85() { - echo "${TEST_SUITE_NAME}: Run pytest ops for Arm Ethos-U85" - - backends/arm/scripts/build_executorch.sh - backends/arm/test/setup_testing.sh - - # Run arm baremetal pytest tests with FVP - pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=10 backends/arm/test/ --ignore=backends/arm/test/models -k u85 - echo "${TEST_SUITE_NAME}: PASS" -} - -test_pytest_models_ethos_u85() { - echo "${TEST_SUITE_NAME}: Run pytest models for Arm Ethos-U85" - - backends/arm/scripts/build_executorch.sh - backends/arm/test/setup_testing.sh - - # Install model dependencies for pytest - source backends/arm/scripts/install_models_for_test.sh - - pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=0 backends/arm/test/models -k u85 - echo "${TEST_SUITE_NAME}: PASS" -} - -test_run_ethos_u85() { - echo "${TEST_SUITE_NAME}: Test ethos-u85 delegate examples with run.sh" - - echo "${TEST_SUITE_NAME}: Test target Ethos-U85" - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=add - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-256 --model_name=add --bundleio - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-512 --model_name=add --bundleio --etdump - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-1024 --model_name=add --etdump - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-2048 --model_name=mul --pte_placement=elf - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=mul --pte_placement=0x38000000 - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=mul --bundleio --pte_placement=elf - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-256 --model_name=mul --bundleio --pte_placement=0x38000000 - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=examples/arm/example_modules/add.py - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-1024 --model_name=examples/arm/example_modules/add.py --bundleio - - # Cortex-M op tests - echo "${TEST_SUITE_NAME}: Test target Cortex-M55 (on Ethos-U85)" - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=qops --bundleio - - echo "${TEST_SUITE_NAME}: PASS" -} - -# ---------------------------------------------------------- -# -------- Vulkan Graph Format (VGF) specific tests -------- -# ---------------------------------------------------------- -test_pytest_ops_vkml() { - echo "${TEST_SUITE_NAME}: Run pytest operator tests with VKML runtime" - - source backends/arm/test/setup_testing_vkml.sh - - pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=10 backends/arm/test/ \ - --ignore=backends/arm/test/models -k _vgf_ - echo "${TEST_SUITE_NAME}: PASS" -} - -test_pytest_models_vkml() { - echo "${TEST_SUITE_NAME}: Run pytest model tests with VKML runtime" - - source backends/arm/test/setup_testing_vkml.sh - - # Install model dependencies for pytest - source backends/arm/scripts/install_models_for_test.sh - - pytest "${PYTEST_RETRY_ARGS[@]}" --verbose --color=yes --numprocesses=auto --durations=0 backends/arm/test/models -k _vgf_ - echo "${TEST_SUITE_NAME}: PASS" -} - -test_run_vkml() { - echo "${TEST_SUITE_NAME}: Test VKML delegate examples with run.sh" - - echo "${TEST_SUITE_NAME}: Test VKML" - out_folder="arm_test/test_run" - - examples/arm/run.sh --et_build_root=${out_folder} --target=vgf --model_name=add --output=${out_folder}/runner - examples/arm/run.sh --et_build_root=${out_folder} --target=vgf --model_name=mul --output=${out_folder}/runner - - examples/arm/run.sh --et_build_root=${out_folder} --target=vgf --model_name=qadd --output=${out_folder}/runner - examples/arm/run.sh --et_build_root=${out_folder} --target=vgf --model_name=qops --output=${out_folder}/runner - - echo "${TEST_SUITE_NAME}: PASS" -} - -# ------------------------------------ -# -------- Miscelaneous tests -------- -# ------------------------------------ -test_model_smollm2_135M() { - echo "${TEST_SUITE_NAME}: Test SmolLM2-135M on Ethos-U85" - - # Build common libs once - python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --build_libs - - python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u85-128 --model=smollm2 --extra_flags="-DEXECUTORCH_SELECT_OPS_LIST=dim_order_ops::_to_dim_order_copy.out" --specify_ethosu_scratch - - echo "${TEST_SUITE_NAME}: PASS" -} - -test_smaller_stories_llama() { - echo "${TEST_SUITE_NAME}: Test smaller_stories_llama" - - backends/arm/scripts/build_executorch.sh - - mkdir -p stories110M - pushd stories110M - wget -N https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.pt - echo '{"dim": 768, "multiple_of": 32, "n_heads": 12, "n_layers": 12, "norm_eps": 1e-05, "vocab_size": 32000}' > params.json - popd - - # Get path to source directory - pytest \ - -c /dev/null \ - "${PYTEST_RETRY_ARGS[@]}" \ - --verbose \ - --color=yes \ - --numprocesses=auto \ - --junit-xml=stories110M/test-reports/unittest.xml \ - backends/arm/test/models/test_llama.py \ - --llama_inputs stories110M/stories110M.pt stories110M/params.json stories110m - - echo "${TEST_SUITE_NAME}: PASS" -} - -test_memory_allocation() { - echo "${TEST_SUITE_NAME}: Test ethos-u memory allocation with run.sh" - - mkdir -p arm_test/test_run - # Ethos-U85 - echo "${TEST_SUITE_NAME}: Test target Ethos-U85" - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=examples/arm/example_modules/add.py &> arm_test/test_run/full.log - python3 backends/arm/test/test_memory_allocator_log.py --log arm_test/test_run/full.log \ - --require "model_pte_program_size" "<= 3200 B" \ - --require "method_allocator_planned" "<= 64 B" \ - --require "method_allocator_loaded" "<= 1024 B" \ - --require "method_allocator_input" "<= 16 B" \ - --require "Total DRAM used" "<= 0.06 KiB" - echo "${TEST_SUITE_NAME}: PASS" -} - -test_undefinedbehavior_sanitizer() { - echo "${TEST_SUITE_NAME}: Test ethos-u executor_runner with UBSAN" - - mkdir -p arm_test/test_run - # Ethos-U85 - echo "${TEST_SUITE_NAME}: Test target Ethos-U85" - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=examples/arm/example_modules/add.py --build_type=UndefinedSanitizer - echo "${TEST_SUITE_NAME}: PASS" -} - -test_address_sanitizer() { - echo "${TEST_SUITE_NAME}: Test ethos-u executor_runner with ASAN" - - mkdir -p arm_test/test_run - # Ethos-U85 - echo "${TEST_SUITE_NAME}: Test target Ethos-U85" - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=examples/arm/example_modules/add.py --build_type=AddressSanitizer - echo "${TEST_SUITE_NAME}: PASS" -} - - -${TEST_SUITE} diff --git a/backends/arm/test/test_arm_baremetal.sh b/backends/arm/test/test_arm_baremetal.sh new file mode 120000 index 00000000000..21637b9850b --- /dev/null +++ b/backends/arm/test/test_arm_baremetal.sh @@ -0,0 +1 @@ +test_arm_backend.sh \ No newline at end of file diff --git a/backends/arm/test/test_arm_ootb.sh b/backends/arm/test/test_arm_ootb.sh index e97dfc10ddb..078566445f9 100755 --- a/backends/arm/test/test_arm_ootb.sh +++ b/backends/arm/test/test_arm_ootb.sh @@ -22,7 +22,7 @@ if [[ "$1" == "-h" || "$1" == "--help" ]]; then fi if [[ $# -eq 0 ]]; then - TEST_SUITES=(run_ootb_tests_ethos_u run_ootb_tests_tosa run_ootb_tests_vgf run_deit_e2e_ethos_u) + TEST_SUITES=(run_ootb_tests_ethos_u run_ootb_tests_tosa run_ootb_tests_vgf run_deit_e2e_ethos_u run_swin2sr_e2e_vgf) else TEST_SUITES=("$1") fi @@ -66,7 +66,7 @@ run_deit_e2e_ethos_u() { local image_path="${work_root}/dog.bmp" local pte_path="${export_dir}/deit_tiny_smoke.pte" local toolchain_file="${et_root_dir}/examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake" - echo "${FUNCNAME}: Work root is ${work_root}; existing artifacts will be reused if present" + echo "${FUNCNAME}: Work directory: ${work_root}; existing artifacts will be reused if present" mkdir -p "${model_dir}" "${export_dir}" "${build_dir}" @@ -150,6 +150,135 @@ run_deit_e2e_ethos_u() { echo "${FUNCNAME}: PASS" } +run_swin2sr_e2e_vgf() { + echo "$FUNCNAME: Prepare demo assets, export FP/INT8, build, and run the Swin2SR VGF e2e test" + + local script_dir + script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) + et_root_dir=$(cd "${script_dir}/../../.." && pwd) + local example_dir="${et_root_dir}/examples/arm/super_resolution_example_vgf" + local work_root="${et_root_dir}/arm_test/swin2sr_vgf_ootb_smoke" + local demo_dir="${work_root}/demo_assets" + local runtime_dir="${demo_dir}/runtime" + local runner_path="${work_root}/executor_runner" + local input_image="${runtime_dir}/demo_lr_64.png" + local fp_pte_path="${demo_dir}/swin2sr_x2_vgf_fp.pte" + local int8_pte_path="${demo_dir}/swin2sr_x2_vgf_int8.pte" + local fp_output_image="${runtime_dir}/demo_fp_128.png" + local int8_output_image="${runtime_dir}/demo_int8_128.png" + local checkpoint_id="caidas/swin2SR-classical-sr-x2-64" + local checkpoint_revision="cee1c923c6a37361c6e5650b65dcf4be821e5d52" + echo "${FUNCNAME}: Work directory: ${work_root}; existing artifacts will be reused if present" + + mkdir -p "${demo_dir}" "${runtime_dir}" + + setup_path_script=${et_root_dir}/examples/arm/arm-scratch/setup_path.sh + source ${setup_path_script} + + echo "${FUNCNAME}: Installing example requirements" + pip install -r "${example_dir}/requirements.txt" + + echo "${FUNCNAME}: Preparing deterministic demo assets" + python3 "${example_dir}/model_export/prepare_demo_assets.py" \ + --output-dir "${demo_dir}" + + echo "${FUNCNAME}: Building VKML executor_runner" + "${et_root_dir}/backends/arm/scripts/build_executor_runner_vkml.sh" \ + --output="${work_root}" + + if [[ ! -f "${runner_path}" ]]; then + runner_path=$(find "${work_root}" -name executor_runner -type f | head -n 1) + fi + [[ -f "${runner_path}" ]] || { + echo "${FUNCNAME}: Missing executor_runner under ${work_root}" + return 1 + } + + echo "${FUNCNAME}: Exporting FP Swin2SR model" + python3 "${example_dir}/model_export/export_super_resolution.py" \ + --model-name swin2sr \ + --checkpoint "${checkpoint_id}" \ + --checkpoint-revision "${checkpoint_revision}" \ + --input-height 64 \ + --input-width 64 \ + --quantization-mode none \ + --eval-lr-dir "${demo_dir}/eval/lr" \ + --eval-hr-dir "${demo_dir}/eval/hr" \ + --num-eval-samples 2 \ + --output-path "${fp_pte_path}" + + for artifact in \ + "${fp_pte_path}" \ + "${demo_dir}/swin2sr_x2_vgf_fp.json" \ + "${demo_dir}/swin2sr_x2_vgf_fp_delegation.txt" \ + "${demo_dir}/swin2sr_x2_vgf_fp_metrics.json"; do + [[ -f "${artifact}" ]] || { + echo "${FUNCNAME}: Missing FP export artifact ${artifact}" + return 1 + } + done + + echo "${FUNCNAME}: Exporting INT8 Swin2SR model" + python3 "${example_dir}/model_export/export_super_resolution.py" \ + --model-name swin2sr \ + --checkpoint "${checkpoint_id}" \ + --checkpoint-revision "${checkpoint_revision}" \ + --input-height 64 \ + --input-width 64 \ + --quantization-mode int8 \ + --calibration-lr-dir "${demo_dir}/calibration/lr" \ + --eval-lr-dir "${demo_dir}/eval/lr" \ + --eval-hr-dir "${demo_dir}/eval/hr" \ + --num-calibration-samples 4 \ + --num-eval-samples 2 \ + --output-path "${int8_pte_path}" + + for artifact in \ + "${int8_pte_path}" \ + "${demo_dir}/swin2sr_x2_vgf_int8.json" \ + "${demo_dir}/swin2sr_x2_vgf_int8_delegation.txt" \ + "${demo_dir}/swin2sr_x2_vgf_int8_metrics.json"; do + [[ -f "${artifact}" ]] || { + echo "${FUNCNAME}: Missing INT8 export artifact ${artifact}" + return 1 + } + done + + echo "${FUNCNAME}: Running FP runtime smoke" + python3 "${example_dir}/runtime/run_super_resolution.py" \ + --model-path "${fp_pte_path}" \ + --runner "${runner_path}" \ + --input-image "${input_image}" \ + --output-image "${fp_output_image}" \ + --working-dir "${runtime_dir}/fp_work" + + [[ -f "${fp_output_image}" ]] || { + echo "${FUNCNAME}: Missing FP runtime output ${fp_output_image}" + return 1 + } + + if [[ "$(uname -s)" == "Linux" ]]; then + echo "${FUNCNAME}: Running INT8 runtime smoke" + python3 "${example_dir}/runtime/run_super_resolution.py" \ + --model-path "${int8_pte_path}" \ + --runner "${runner_path}" \ + --input-image "${input_image}" \ + --output-image "${int8_output_image}" \ + --working-dir "${runtime_dir}/int8_work" + + [[ -f "${int8_output_image}" ]] || { + echo "${FUNCNAME}: Missing INT8 runtime output ${int8_output_image}" + return 1 + } + else + # TODO: MLETORCH-2105 remove this once the next ML SDK release supports + # quantized VKML runtime validation on Darwin. + echo "${FUNCNAME}: Skipping INT8 runtime on $(uname -s); quantized VKML runtime validation is Linux-only" + fi + + echo "${FUNCNAME}: PASS" +} + for suite in "${TEST_SUITES[@]}"; do "${suite}" done diff --git a/backends/arm/test/test_github_trunk_zephyr.sh b/backends/arm/test/test_github_trunk_zephyr.sh index 36d4e40411c..c4ab0e5cfbc 100755 --- a/backends/arm/test/test_github_trunk_zephyr.sh +++ b/backends/arm/test/test_github_trunk_zephyr.sh @@ -9,7 +9,7 @@ # github/workflows/trunk.yml as a bash script, (ignoring the conda commands) # This is mostly useful for testing this before upstreaming or for debugging CI issues. # Or a way to setup zephyrOS so you can play with it. -# Parse optional argument to set TARGET_LIST environment variable +# Parse optional arguments to select README and target list. # Target list (--target-list) adds a comma separated list of targets you want to test # matching the zephyr/README.md tags: @@ -20,22 +20,33 @@ set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" -ROOT_DIR="$(cd "${SCRIPT_DIR}" && pwd)" +ROOT_DIR="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +TARGET_LIST_OVERRIDDEN=0 +if [[ -n "${TARGET_LIST:-}" ]]; then + TARGET_LIST_OVERRIDDEN=1 +fi TARGET_LIST="${TARGET_LIST:-ethos-u55,cortex-m55,ethos-u85}" +README_PATH="${README_PATH:-}" +HELLO_README_PATH="zephyr/samples/hello-executorch/README.md" +MV2_README_PATH="zephyr/samples/mv2-ethosu/README.md" +DEFAULT_MV2_TARGET_LIST="ethos-u55,ethos-u85" usage() { cat <&2 + usage >&2 + exit 2 +fi -SCRIPT_CONTENT="$(python - <<'PY' -from ruamel.yaml import YAML -with open(".github/workflows/trunk.yml") as f: - data = YAML().load(f) -script = data["jobs"]["test-arm-backend-zephyr"]["with"]["script"] -filtered = [] -con_lines = ("CONDA_ENV=", "conda activate") -for line in script.splitlines(): - if any(line.strip().startswith(prefix) for prefix in con_lines): - continue - filtered.append(line) -print("\n".join(filtered)) -PY -)" +cd "${ROOT_DIR}" + +run_zephyr_readme() { + local readme_path="$1" + local targets="$2" + + echo "Running ${readme_path} targets: ${targets}" + .ci/scripts/test_zephyr.sh \ + --zephyr-samples-readme-path "${readme_path}" \ + --targets "${targets}" +} -printf "%s\n" "${SCRIPT_CONTENT}" | /bin/bash +if [[ -n "${README_PATH}" ]]; then + run_zephyr_readme "${README_PATH}" "${TARGET_LIST}" +else + run_zephyr_readme "${HELLO_README_PATH}" "${TARGET_LIST}" + run_zephyr_readme "${MV2_README_PATH}" "${DEFAULT_MV2_TARGET_LIST}" +fi diff --git a/backends/arm/tosa/partitioner.py b/backends/arm/tosa/partitioner.py index 97ce7f75248..07548eb5d69 100644 --- a/backends/arm/tosa/partitioner.py +++ b/backends/arm/tosa/partitioner.py @@ -104,6 +104,15 @@ def _is_noop_expand(node: torch.fx.node.Node) -> bool: return all(m == 1 for m in multiples) and not changes_rank +def _is_noop_squeeze(node: torch.fx.Node) -> bool: + if node.target != exir_ops.edge.aten.squeeze_copy.dims: + return False + else: + input_tensor = get_first_fake_tensor(ensure_type(torch.fx.Node, node.args[0])) + output_tensor = get_first_fake_tensor(node) + return input_tensor.shape == output_tensor.shape + + def _is_view_copy(node: torch.fx.node.Node) -> bool: return node.target == exir_ops.edge.aten.view_copy.default @@ -388,6 +397,7 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: or _is_noop_expand(node) or _is_noop_detach_copy(node) or _is_noop_to_dim_order_copy(node) + or _is_noop_squeeze(node) or _is_view_copy(node) or _is_noop_as_strided_copy(node) or node.target in Q_OPS diff --git a/backends/cadence/aot/compiler_funcs.py b/backends/cadence/aot/compiler_funcs.py index bf7f79127a0..02dcde7fd39 100644 --- a/backends/cadence/aot/compiler_funcs.py +++ b/backends/cadence/aot/compiler_funcs.py @@ -35,6 +35,7 @@ def trace( model.eval() decomp_table = torch.export.default_decompositions() + ops_to_keep = [*(ops_to_keep or []), torch.ops.aten._safe_softmax.default] # pyre-fixme[6]: For 1st argument expected `Dict[typing.Callable[..., typing.Any remove_decompositions(decomp_table, ops_to_keep) program = torch.export.export(model, inputs, strict=strict).run_decompositions( diff --git a/backends/cadence/aot/passes.py b/backends/cadence/aot/passes.py index 8a03d72420e..f43ac3e4d2c 100644 --- a/backends/cadence/aot/passes.py +++ b/backends/cadence/aot/passes.py @@ -33,6 +33,7 @@ from executorch.backends.cadence.aot.replace_ops import ( CadenceReplaceOpsInGraph, ReplaceMulTensorWithMulAndFullOpsPass, + ReplaceSafeSoftmaxWithSoftmax, ) from executorch.backends.cadence.aot.simplify_ops import CadenceSimplifyOpsInGraph from executorch.backends.cadence.aot.type_dispatch import CompileTimeTypeDispatchPass @@ -131,7 +132,8 @@ def apply_torch_ops_passes(expo_program: ExportedProgram) -> ExportedProgram: """ aten_passes: List[Callable[[torch.fx.GraphModule], Optional[PassResult]]] = [ - ReplaceMulTensorWithMulAndFullOpsPass() + ReplaceSafeSoftmaxWithSoftmax(), + ReplaceMulTensorWithMulAndFullOpsPass(), ] # TODO(T230417247): Use PassResult which is currently ignored. PassManager(aten_passes)(expo_program.graph_module) diff --git a/backends/cadence/aot/remove_ops.py b/backends/cadence/aot/remove_ops.py index e532d088e5c..c221c3a5a18 100644 --- a/backends/cadence/aot/remove_ops.py +++ b/backends/cadence/aot/remove_ops.py @@ -603,16 +603,16 @@ def maybe_remove_or_replace(self, node: Node) -> bool: @register_cadence_pass(CadencePassAttribute(opt_level=2)) class RemovePermutesAroundElementwiseOps(_SharedRemovePermutesAroundElementwiseOps): - permutable_ops: set[EdgeOpOverload] = ( - _SharedRemovePermutesAroundElementwiseOps.permutable_ops - | { - exir_ops.edge.cadence.quantize_per_tensor.default, - exir_ops.edge.cadence.dequantize_per_tensor.default, - exir_ops.edge.cadence.quantized_relu.per_tensor, - exir_ops.edge.cadence.requantize.per_tensor, - exir_ops.edge.cadence.quantized_add.per_tensor, - } - ) + def __init__(self) -> None: + super().__init__( + extra_permutable_ops={ + exir_ops.edge.cadence.quantize_per_tensor.default, + exir_ops.edge.cadence.dequantize_per_tensor.default, + exir_ops.edge.cadence.quantized_relu.per_tensor, + exir_ops.edge.cadence.requantize.per_tensor, + exir_ops.edge.cadence.quantized_add.per_tensor, + } + ) @register_cadence_pass(CadencePassAttribute(opt_level=2)) diff --git a/backends/cortex_m/CMakeLists.txt b/backends/cortex_m/CMakeLists.txt index 8c8255b7b1b..876c65982e6 100644 --- a/backends/cortex_m/CMakeLists.txt +++ b/backends/cortex_m/CMakeLists.txt @@ -50,6 +50,26 @@ else() FetchContent_MakeAvailable(cmsis_nn) endif() +if(TARGET cmsis-nn) + if(CMSIS_NN_LOCAL_PATH AND EXISTS "${CMSIS_NN_LOCAL_PATH}") + set(cmsis_nn_source_dir "${CMSIS_NN_LOCAL_PATH}") + else() + set(cmsis_nn_source_dir "${cmsis_nn_SOURCE_DIR}") + endif() + if(cmsis_nn_source_dir) + set(cmsis_nn_include_dir "${cmsis_nn_source_dir}/Include") + set_target_properties( + cmsis-nn + PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES + "$;$" + ) + install(DIRECTORY "${cmsis_nn_include_dir}/" + DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/cmsis-nn" + ) + endif() +endif() + # Cortex-M ops kernel sources set(_cortex_m_kernels__srcs ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_dequantize_per_tensor.cpp diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index f95587a00d3..f0326ec76c4 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -8,6 +8,7 @@ from typing import Any, Optional, Type from executorch.backends.arm._passes import ( + DeduplicateGetAttrPass, FoldAndAnnotateQParamsPass, ScalarsToAttributePass, ) @@ -52,6 +53,7 @@ class CortexMPassManager(PassManager): ReplaceScalarWithTensorArgPass, ClampHardswishPass, DecomposeMeanPass, + DeduplicateGetAttrPass, ] def __init__( diff --git a/backends/cortex_m/quantizer/pattern_matcher.py b/backends/cortex_m/quantizer/pattern_matcher.py index 3694fd22a6c..6e09fdbe58f 100644 --- a/backends/cortex_m/quantizer/pattern_matcher.py +++ b/backends/cortex_m/quantizer/pattern_matcher.py @@ -113,15 +113,25 @@ def _get_match(self, node_queue: List[Node]) -> List[Node]: return [] def _get_matches( - self, node_queue: List[Node], quantization_config: QuantizationConfig + self, node_queue: List[Node], quantization_config: Optional[QuantizationConfig] ) -> List[PatternMatchResult]: """Returns the longest accepted match starting at the first node of the queue as well as longer rejected matches. """ + # Annotating with None means rejecting quantization - this is always supported. + if quantization_config is None: + node = node_queue[0] + if node.meta.get(self.Q_PATTERN_MATCHED_KEY, False): + return [ + PatternMatchResult([node], False, self.REJECT_PREVIOUSLY_ANNOTATED) + ] + + node.meta[self.Q_PATTERN_MATCHED_KEY] = True + return [PatternMatchResult([node], True)] + matches: list[PatternMatchResult] = [] accepted = False max_match_length = len(node_queue) - while max_match_length > 0 and not accepted: match = self._get_match(node_queue[:max_match_length]) max_match_length = ( @@ -136,7 +146,7 @@ def _get_matches( return matches def _dequeue_and_get_matches( - self, node_queue: List[Node], quantization_config: QuantizationConfig + self, node_queue: List[Node], quantization_config: Optional[QuantizationConfig] ) -> List[PatternMatchResult]: """Dequeues the longest accepted match starting at the first node of the queue, and returns all potential matches that were checked (rejected @@ -160,7 +170,7 @@ def _dequeue_and_get_matches( return potential_matches def find_pattern_matches( - self, nodes: Iterator[Node], quantization_config: QuantizationConfig + self, nodes: Iterator[Node], quantization_config: Optional[QuantizationConfig] ) -> Iterator[PatternMatchResult]: """Match all given patterns in the graph and return match results with acceptance/rejection status. Each node can only be part of one match, diff --git a/backends/mlx/custom_ops.py b/backends/mlx/custom_ops.py index d7d6288ba8f..c03db05d918 100644 --- a/backends/mlx/custom_ops.py +++ b/backends/mlx/custom_ops.py @@ -228,8 +228,16 @@ def rope( # final angles: [1, 1, T, half] angles = (pos_range * inv_freq) * float(scale) else: - # assume freqs is already per-position, just reshape to [1,1,T,half] - angles = freqs.to(torch.float32).view(1, 1, T, half) + if freqs.ndim == 1: + # 1D raw frequencies: compute angles = positions * (1/freqs) + inv_freq = (1.0 / freqs.to(torch.float32)).view(1, 1, 1, half) + pos_range = torch.arange( + pos, pos + T, device=x.device, dtype=torch.float32 + ).view(1, 1, T, 1) + angles = (pos_range * inv_freq) * float(scale) + else: + # 2D per-position angles: reshape to [1,1,T,half] + angles = freqs.to(torch.float32).view(1, 1, T, half) cos = angles.cos().to(x.dtype) # [1,1,T,half] sin = angles.sin().to(x.dtype) # [1,1,T,half] diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index 1f961459d22..fb6597d171e 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -242,6 +242,11 @@ inline void exec_rope(const RopeNode& n, ExecutionState& st, StreamOrDevice s) { freqs_arr = st.const_tensor_ref(*n.freqs); } + // MLX requires exactly one of base or freqs — when freqs is provided, + // base must be nullopt. + std::optional base = + freqs_arr ? std::nullopt : std::optional(n.base); + // MLX has two overloads: rope(..., int offset, ...) and rope(..., const // array& offset, ...) Call the appropriate one based on is_vid if (n.offset.is_vid) { @@ -250,14 +255,14 @@ inline void exec_rope(const RopeNode& n, ExecutionState& st, StreamOrDevice s) { st.set_tensor( n.out, fast::rope( - x, n.dims, n.traditional, n.base, n.scale, offset, freqs_arr, s)); + x, n.dims, n.traditional, base, n.scale, offset, freqs_arr, s)); } else { // Tensor offset from Tid const array& offset = st.const_tensor_ref(n.offset.tid); st.set_tensor( n.out, fast::rope( - x, n.dims, n.traditional, n.base, n.scale, offset, freqs_arr, s)); + x, n.dims, n.traditional, base, n.scale, offset, freqs_arr, s)); } } diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py index afc45adcc93..4471610519e 100644 --- a/backends/mlx/test/test_ops.py +++ b/backends/mlx/test/test_ops.py @@ -1803,6 +1803,82 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]: return (q, k, pos_tensor) +class RopeCustomFreqsModel(nn.Module): + """Model that applies RoPE with custom 1D frequencies (partial rotary).""" + + def __init__(self, dims: int = 32, head_dim: int = 64): + super().__init__() + self.dims = dims + self.head_dim = head_dim + # Simulate proportional RoPE: compute freqs for rotary dims only + inv_freq = 1.0 / ( + 500000.0 ** (torch.arange(0, dims, 2, dtype=torch.float32) / head_dim) + ) + self.register_buffer("freqs", 1.0 / inv_freq, persistent=False) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + pos_tensor: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + pos = pos_tensor.item() + q_rot = torch.ops.mlx.rope(q, self.dims, pos, False, 0.0, 1.0, self.freqs) + k_rot = torch.ops.mlx.rope(k, self.dims, pos, False, 0.0, 1.0, self.freqs) + return q_rot, k_rot + + +@register_test +class RopeCustomFreqsTest(OpTestCase): + """Test RoPE with custom 1D frequencies (partial rotary, like Gemma 4).""" + + name = "rope_custom_freqs" + rtol = 1e-4 + atol = 1e-4 + + def __init__( + self, + batch_size: int = 1, + num_heads: int = 8, + seq_len: int = 4, + head_dim: int = 64, + dims: int = 32, + pos: int = 0, + ): + self.batch_size = batch_size + self.num_heads = num_heads + self.seq_len = seq_len + self.head_dim = head_dim + self.dims = dims + self.pos = pos + self.name = "rope_custom_freqs" + + @classmethod + def get_test_configs(cls) -> List["RopeCustomFreqsTest"]: + configs = [ + cls(), + cls(pos=10), + cls(head_dim=128, dims=64), + ] + for cfg in configs: + parts = ["rope_custom_freqs"] + if cfg.pos > 0: + parts.append(f"pos{cfg.pos}") + if cfg.head_dim != 64: + parts.append(f"hd{cfg.head_dim}") + cfg.name = "_".join(parts) + return configs + + def create_model(self) -> nn.Module: + return RopeCustomFreqsModel(dims=self.dims, head_dim=self.head_dim) + + def create_inputs(self) -> Tuple[torch.Tensor, ...]: + q = torch.randn(self.batch_size, self.num_heads, self.seq_len, self.head_dim) + k = torch.randn(self.batch_size, self.num_heads, self.seq_len, self.head_dim) + pos_tensor = torch.tensor(self.pos, dtype=torch.int64) + return (q, k, pos_tensor) + + from executorch.backends.mlx.llm.cache import KVCache diff --git a/backends/nxp/aten_passes/convert_1d_conv_to_2d.py b/backends/nxp/aten_passes/convert_1d_conv_to_2d.py index 6963ac6bc3e..f95c271f2fa 100644 --- a/backends/nxp/aten_passes/convert_1d_conv_to_2d.py +++ b/backends/nxp/aten_passes/convert_1d_conv_to_2d.py @@ -8,6 +8,7 @@ try_get_tensor_constant_from_node, ) from executorch.backends.nxp.backend.graph_utils import is_batch_norm +from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec from torch._subclasses import FakeTensor, FakeTensorMode from torch.ao.quantization.fx.utils import get_new_attr_name_with_prefix from torch.export.unflatten import _assign_attr, _AttrKind @@ -26,18 +27,21 @@ class ConvertConv1dToConv2dPass(PassBase): r""" The NXP backend supports only 2D convolutions. Rewrite 1D convolutions into an equivalent 2D form by inserting a singleton spatial dimension and then remove it again. - If batch norm is present after the convolution, it is also converted from 1D to 2D. + If batch norm and/or a fusable activation (as defined by the NeutronTargetSpec) follow the convolution, + they are also kept in 2D (before the squeeze) so the partitioner can fuse them with the convolution. - Without batch norm: + Without batch norm or activation: + + Without batch norm or activation: x W x W - [N, C1, H] [I/O, I/O, k] [N, C1, H] [I/O, I/O, 1, k] + [N, C, H] [I/O, I/O, k] [N, C, H] [I/O, I/O, 1, k] │ │ │ │ │ │ ┌─────────▼──────────┐ │ │ │ │ unsqueeze(x, -2) │ │ │ │ └─────────▼──────────┘ │ │ │ │ │ - │ │ [N, C1, 1, H] │ + │ │ [N, C, 1, H] │ │ │ │ │ └────────┐ ┌────────┘ └──────────┐ ┌──────────┘ │ │ │ │ @@ -46,26 +50,26 @@ class ConvertConv1dToConv2dPass(PassBase): │ (1D/transposed 1D) │ ────────────────► │ (2D/transposed 2D) │ └────────────┬───────────┘ with └───────────┬───────────┘ │ │ - │ [N, C2, 1, H] + │ [N, C, 1, H] │ │ │ ┌─────────▼──────────┐ │ │ squeeze(x, -2) │ │ └─────────┬──────────┘ │ │ ▼ ▼ - [N, C2, H] [N, C2, H] + [N, C, H] [N, C, H] y y With batch norm: x W x W - [N, C1, H] [I/O, I/O, k] [N, C1, H] [I/O, I/O, 1, k] + [N, C, H] [I/O, I/O, k] [N, C, H] [I/O, I/O, 1, k] │ │ │ │ │ │ ┌─────────▼──────────┐ │ │ │ │ unsqueeze(x, -2) │ │ │ │ └─────────▼──────────┘ │ │ │ │ │ - │ │ [N, C1, 1, H] │ + │ │ [N, C, 1, H] │ │ │ │ │ └────────┐ ┌────────┘ └──────────┐ ┌──────────┘ │ │ │ │ @@ -74,24 +78,102 @@ class ConvertConv1dToConv2dPass(PassBase): │ (1D/transposed 1D) │ ────────────────► │ (2D/transposed 2D) │ └────────────┬───────────┘ with └───────────┬───────────┘ │ │ - [N, C2, H] [N, C2, 1, H] + [N, C, H] [N, C, 1, H] │ │ ┌───────▼───────┐ ┌───────▼───────┐ │ batch_norm │ │ batch_norm │ │ (1D) │ │ (2D) │ └───────┬───────┘ └───────┬───────┘ │ │ - │ [N, C3, 1, H] + │ [N, C, 1, H] │ │ │ ┌───────▼────────┐ │ │ squeeze(-2) │ │ └───────┬────────┘ │ │ ▼ ▼ - [N, C3, H] [N, C3, H] + [N, C, H] [N, C, H] + y y + + With activation (e.g. relu): + + x W x W + [N, C, H] [I/O, I/O, k] [N, C, H] [I/O, I/O, 1, k] + │ │ │ │ + │ │ ┌─────────▼──────────┐ │ + │ │ │ unsqueeze(x, -2) │ │ + │ │ └─────────▼──────────┘ │ + │ │ │ │ + │ │ [N, C, 1, H] │ + │ │ │ │ + └────────┐ ┌────────┘ └──────────┐ ┌──────────┘ + │ │ │ │ + ┌────────▼───────▼───────┐ ┌────────▼─────▼────────┐ + │ convolution ◄──B [O] replace │ convolution ◄──B [O] + │ (1D/transposed 1D) │ ────────────────► │ (2D/transposed 2D) │ + └────────────┬───────────┘ with └───────────┬───────────┘ + │ │ + [N, C, H] [N, C, 1, H] + │ │ + ┌───────▼───────┐ ┌───────▼───────┐ + │ relu │ │ relu │ + └───────┬───────┘ └───────┬───────┘ + │ │ + │ [N, C, 1, H] + │ │ + │ ┌───────▼────────┐ + │ │ squeeze(-2) │ + │ └───────┬────────┘ + │ │ + ▼ ▼ + [N, C, H] [N, C, H] + y y + + With batch norm and activation: + + x W x W + [N, C, H] [I/O, I/O, k] [N, C, H] [I/O, I/O, 1, k] + │ │ │ │ + │ │ ┌─────────▼──────────┐ │ + │ │ │ unsqueeze(x, -2) │ │ + │ │ └─────────▼──────────┘ │ + │ │ │ │ + │ │ [N, C, 1, H] │ + │ │ │ │ + └────────┐ ┌────────┘ └──────────┐ ┌──────────┘ + │ │ │ │ + ┌────────▼───────▼───────┐ ┌────────▼─────▼────────┐ + │ convolution ◄──B [O] replace │ convolution ◄──B [O] + │ (1D/transposed 1D) │ ────────────────► │ (2D/transposed 2D) │ + └────────────┬───────────┘ with └───────────┬───────────┘ + │ │ + [N, C, H] [N, C, 1, H] + │ │ + ┌───────▼───────┐ ┌───────▼───────┐ + │ batch_norm │ │ batch_norm │ + │ (1D) │ │ (2D) │ + └───────┬───────┘ └───────┬───────┘ + │ │ + [N, C, H] [N, C, 1, H] + │ │ + ┌───────▼───────┐ ┌───────▼───────┐ + │ relu │ │ relu │ + └───────┬───────┘ └───────┬───────┘ + │ │ + │ [N, C, 1, H] + │ │ + │ ┌───────▼────────┐ + │ │ squeeze(-2) │ + │ └───────┬────────┘ + │ │ + ▼ ▼ + [N, C, H] [N, C, H] y y """ + def __init__(self, neutron_target_spec: NeutronTargetSpec): + self.neutron_target_spec = neutron_target_spec + @staticmethod def _is_conv_1d(node: Node) -> bool: return node.target == torch.ops.aten.conv1d.default @@ -204,12 +286,12 @@ def _create_some_conv_2d_node(self, target, *conv_args): return some_conv_node - def _create_sq_or_unsq_node(self, target, *sq_or_unsq_args) -> Node: - sq_or_unsq_node = self.graph_module.graph.call_function(target, sq_or_unsq_args) + def _create_generic_node_by_target(self, target, *args) -> Node: + new_node = self.graph_module.graph.call_function(target, args) - sq_or_unsq_node.meta["source_fn_stack"] = [(sq_or_unsq_node.name, target)] + new_node.meta["source_fn_stack"] = [(new_node.name, target)] with FakeTensorMode() as mode: - inp_node = sq_or_unsq_args[0] + inp_node = args[0] fake_input = FakeTensor.from_tensor( torch.empty( self._get_node_shape(inp_node), dtype=self._get_node_dtype(inp_node) @@ -217,12 +299,12 @@ def _create_sq_or_unsq_node(self, target, *sq_or_unsq_args) -> Node: mode, ) - output = target(fake_input, *sq_or_unsq_args[1:]) - sq_or_unsq_node.meta["val"] = FakeTensor.from_tensor( + output = target(fake_input, *args[1:]) + new_node.meta["val"] = FakeTensor.from_tensor( torch.empty(output.shape, dtype=output.dtype), mode ) - return sq_or_unsq_node + return new_node @staticmethod def _get_conv_1d_transp_args(node: Node): @@ -299,7 +381,7 @@ def _convert_node_1d_args_to_2d(self, old_1d_node: Node): # input = [n, c, h] => [n, c, 1, h] unsqueeze_target = torch.ops.aten.unsqueeze.default inp_unsq_args = (input_node, -2) - inp_unsq_node = self._create_sq_or_unsq_node( + inp_unsq_node = self._create_generic_node_by_target( unsqueeze_target, *inp_unsq_args ) @@ -357,35 +439,47 @@ def call(self, graph_module: GraphModule) -> PassResult: ) old_1d_conv_users = list(old_1d_node.users.keys()) + last_4d_node = new_2d_node + node_to_replace = old_1d_node + nodes_to_erase = [] + + old_1d_bn_users = old_1d_conv_users + if len(old_1d_conv_users) == 1 and is_batch_norm(old_1d_conv_users[0]): bn_1d_node = old_1d_conv_users[0] - - # also convert batch_norm 1d to 2d - with self.graph_module.graph.inserting_after(new_2d_node): + with self.graph_module.graph.inserting_after(last_4d_node): bn_2d_args = (new_2d_node,) + bn_1d_node.args[1:] bn_2d_node = self._create_batch_norm_2d_node(*bn_2d_args) - - with self.graph_module.graph.inserting_after(bn_2d_node): - squeeze_target = torch.ops.aten.squeeze.dim - - out_sq_args = (bn_2d_node, -2) - out_sq_node = self._create_sq_or_unsq_node( - squeeze_target, *out_sq_args - ) - - bn_1d_node.replace_all_uses_with(out_sq_node) - self.graph_module.graph.erase_node(bn_1d_node) - - else: - with self.graph_module.graph.inserting_after(new_2d_node): - squeeze_target = torch.ops.aten.squeeze.dim - - out_sq_args = (new_2d_node, -2) - out_sq_node = self._create_sq_or_unsq_node( - squeeze_target, *out_sq_args + last_4d_node = bn_2d_node + node_to_replace = bn_1d_node + nodes_to_erase.append(bn_1d_node) + old_1d_bn_users = list(bn_1d_node.users.keys()) + + if len( + old_1d_bn_users + ) == 1 and self.neutron_target_spec.neutron_target_info.is_supported_fused_activation__aten( + old_1d_bn_users[0] + ): + act_1d_node = old_1d_bn_users[0] + with self.graph_module.graph.inserting_after(last_4d_node): + act_2d_args = (last_4d_node,) + act_1d_node.args[1:] + act_2d_node = self._create_generic_node_by_target( + act_1d_node.target, *act_2d_args ) + last_4d_node = act_2d_node + node_to_replace = act_1d_node + nodes_to_erase.append(act_1d_node) + + with self.graph_module.graph.inserting_after(last_4d_node): + squeeze_target = torch.ops.aten.squeeze.dim + out_sq_args = (last_4d_node, -2) + out_sq_node = self._create_generic_node_by_target( + squeeze_target, *out_sq_args + ) - old_1d_node.replace_all_uses_with(out_sq_node) + node_to_replace.replace_all_uses_with(out_sq_node) + for n in reversed(nodes_to_erase): + self.graph_module.graph.erase_node(n) graph_module.graph.erase_node(old_1d_node) made_changes = True diff --git a/backends/nxp/aten_passes/neutron_aten_pass_manager.py b/backends/nxp/aten_passes/neutron_aten_pass_manager.py index 4f1ff2648aa..d00b45be582 100644 --- a/backends/nxp/aten_passes/neutron_aten_pass_manager.py +++ b/backends/nxp/aten_passes/neutron_aten_pass_manager.py @@ -52,7 +52,7 @@ def _get_default_passes(neutron_target_spec, qat_mode: bool = False) -> list[Pas FuseLinearAndAddPass(), MoveActivationBeforeConcat(neutron_target_spec), ConvertDivToMulPass(), - ConvertConv1dToConv2dPass(), + ConvertConv1dToConv2dPass(neutron_target_spec), ] if not qat_mode: diff --git a/backends/nxp/backend/ir/converter/builder/model_builder.py b/backends/nxp/backend/ir/converter/builder/model_builder.py index d3547acb67f..f97a194ce87 100755 --- a/backends/nxp/backend/ir/converter/builder/model_builder.py +++ b/backends/nxp/backend/ir/converter/builder/model_builder.py @@ -742,7 +742,10 @@ def _validate_new_tensor_name(self, name: str) -> str: return new_name def op_code_index_for_op_type( - self, op_type: BuiltinOperator, version: int = 1, custom_code: str = None + self, + op_type: BuiltinOperator | int, + version: int = 1, + custom_code: str | None = None, ): """ Return the index to the 'operator_codes' vector in the TFLite model for the operator diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/constant_pad_nd_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/constant_pad_nd_converter.py index efab4fb95c7..ca59eae811c 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/constant_pad_nd_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/constant_pad_nd_converter.py @@ -7,6 +7,7 @@ from typing import Collection import numpy as np +import torch from executorch.backends.nxp.backend.data_format import NXP_NODE_FORMAT @@ -41,19 +42,33 @@ def _is_supported_on_target( parameters_mapping: dict[str, Parameter], custom_delegation_options: CustomDelegationOptions, ) -> bool: - paddings = node.args[1] - if node.meta[NXP_NODE_FORMAT].is_channels_first(): - # Dim `1` will end up being the channels. It is padded by paddings[4:6]. - if len(paddings) > 4 and paddings[4:6] != [0, 0]: - # Attempt to Pad channels dimension -> currently not supported - return False - else: - # Dim `-1` will end up being the channels. It is padded by paddings[:2]. - if len(paddings) > 0 and paddings[:2] != [0, 0]: - # Attempt to Pad channels dimension -> currently not supported + if custom_delegation_options.use_new_flow_neutron_c: + # Requirements specified by the new Neutron flow documentation. + + if not NodeConverter.uses_quantization_type_for_io( + node, + supported_types=[torch.int8, torch.uint8], + input_indices=[0], + output_indices=[0], + ): return False - return True + return True + + else: + paddings = node.args[1] + if node.meta[NXP_NODE_FORMAT].is_channels_first(): + # Dim `1` will end up being the channels. It is padded by paddings[4:6]. + if len(paddings) > 4 and paddings[4:6] != [0, 0]: + # Attempt to Pad channels dimension -> currently not supported + return False + else: + # Dim `-1` will end up being the channels. It is padded by paddings[:2]. + if len(paddings) > 0 and paddings[:2] != [0, 0]: + # Attempt to Pad channels dimension -> currently not supported + return False + + return True @staticmethod def _is_supported_in_IR( @@ -110,7 +125,14 @@ def _convert_paddings_to_tflite( return paddings def convert(self, node: Node): - """Convert the `aten.constant_pad_nd` operator to TFLite `PadV2`.""" + """Convert the `aten.constant_pad_nd` operator to NeutronIR `PadV2`. + The ExecuTorch schema is: + constant_pad_nd( + Tensor self, + SymInt[] pad, + Scalar value=0 + ) -> Tensor + """ self.assert_convertible(node) t_op = self._create_tflite_op_with_io_tensors(node) diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/leaky_relu_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/leaky_relu_converter.py index e6fcf0e5110..ab778631f74 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/leaky_relu_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/leaky_relu_converter.py @@ -3,6 +3,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import torch + from executorch.backends.nxp.backend.ir.converter.node_converter import ( CustomDelegationOptions, NodeConverter, @@ -10,6 +12,8 @@ from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.leaky_relu_options import ( LeakyRelu, ) + +from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec from torch.fx import Node from torch.nn import Parameter @@ -24,6 +28,29 @@ def _is_supported_in_IR( ) -> bool: return True + @staticmethod + def _is_supported_on_target( + node: Node, + neutron_target_spec: NeutronTargetSpec, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + if custom_delegation_options.use_new_flow_neutron_c: + # Requirements specified by the new Neutron flow documentation. + + if not NodeConverter.uses_quantization_type_for_io( + node, + supported_types=[torch.int8, torch.uint8], + input_indices=[0], + output_indices=[0], + ): + return False + + return True + else: + + return True + def convert(self, node: Node): """Convert the `aten.leaky_relu.default` operator to Neutron IR `LeakyRelu`. The schema is: diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool2d_with_indices_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool2d_with_indices_converter.py index e300d6bbe9f..975aaf57625 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool2d_with_indices_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool2d_with_indices_converter.py @@ -83,14 +83,10 @@ def _is_supported_on_target( ): return False - maximum_supported_kernel_size = 4096 # If there is no padding, Neutron allows maximum stride of 4096. Otherwise, it's 32. But the converter # always inserts a `Pad` operator to add the padding, so the `MaxPool` never pads it's input itself, so # 4096 is always the limit. And similarly, the `MaxPool` input padding limitation does not apply either. maximum_supported_stride = 4096 - - if any(k > maximum_supported_kernel_size for k in kernel_size): - return False if any(s > maximum_supported_stride for s in stride): return False diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/sigmoid_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/sigmoid_converter.py index 96e4655d011..7be2ce180c3 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/sigmoid_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/sigmoid_converter.py @@ -1,8 +1,10 @@ -# Copyright 2025 NXP +# Copyright 2025-2026 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import torch + from executorch.backends.nxp.backend.ir.converter.node_converter import ( CustomDelegationOptions, NodeConverter, @@ -10,6 +12,8 @@ from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import ( BuiltinOperator, ) + +from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec from torch.fx import Node from torch.nn import Parameter @@ -24,7 +28,37 @@ def _is_supported_in_IR( ) -> bool: return True + @staticmethod + def _is_supported_on_target( + node: Node, + neutron_target_spec: NeutronTargetSpec, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + if custom_delegation_options.use_new_flow_neutron_c: + # Requirements specified by the new Neutron flow documentation. + + if not NodeConverter.uses_quantization_type_for_io( + node, + supported_types=[torch.int8, torch.uint8], + input_indices=[0], + output_indices=[0], + ): + return False + + return True + + else: + # Requirements of the old Neutron flow. + return True + def convert(self, node: Node): + """Convert the `aten.sigmoid.default` node to NeutronIR `Logistic` operator. + The ExecuTorch schema is: + sigmoid( + Tensor self + ) -> Tensor + """ self.assert_convertible(node) t_op = self._create_tflite_op_with_io_tensors(node) diff --git a/backends/nxp/requirements-eiq.txt b/backends/nxp/requirements-eiq.txt index 61e5c882c40..5fe425aa4ef 100644 --- a/backends/nxp/requirements-eiq.txt +++ b/backends/nxp/requirements-eiq.txt @@ -1,3 +1,3 @@ --index-url https://eiq.nxp.com/repository -eiq-neutron-sdk==3.1.0 +eiq-neutron-sdk==3.1.1 eiq_nsys diff --git a/backends/nxp/tests/BUCK b/backends/nxp/tests/BUCK index f664e0f1faf..c16d6267425 100644 --- a/backends/nxp/tests/BUCK +++ b/backends/nxp/tests/BUCK @@ -92,6 +92,26 @@ fbcode_target(_kind = python_pytest, ], ) +fbcode_target(_kind = python_pytest, + name = "test_convert_1d_conv_to_2d", + srcs = [ + "test_convert_1d_conv_to_2d.py", + ], + env = { + "PYTEST_ADDOPTS": "--ignore-glob=*full_pipeline* -k 'not full_pipeline'", + }, + deps = [ + "//caffe2:torch", + "//executorch/backends/nxp:aten_passes", + "//executorch/backends/nxp:neutron_backend", + ":executorch_pipeline", + ":models", + "fbsource//third-party/pypi/numpy:numpy", + "fbsource//third-party/pypi/pytest:pytest", + "fbsource//third-party/pypi/pytest-mock:pytest-mock", # @manual + ], +) + fbcode_target(_kind = python_pytest, name = "test_integration", srcs = [ diff --git a/backends/nxp/tests/config_importer.py b/backends/nxp/tests/config_importer.py index 286e4fd0f73..dd424584cde 100644 --- a/backends/nxp/tests/config_importer.py +++ b/backends/nxp/tests/config_importer.py @@ -8,7 +8,7 @@ logger = logging.getLogger(__name__) try: - import test.python.config as test_config # noqa F401 + import config as test_config # noqa F401 logger.debug("Importing from executorch-integration") except ImportError: diff --git a/backends/nxp/tests/generic_tests/test_split_group_convolution.py b/backends/nxp/tests/generic_tests/test_split_group_convolution.py index 804b27e910a..12d2f193f57 100644 --- a/backends/nxp/tests/generic_tests/test_split_group_convolution.py +++ b/backends/nxp/tests/generic_tests/test_split_group_convolution.py @@ -161,7 +161,8 @@ def test_split_group_convolution__1d( # `ConvertConv1dToConv2dPass` is needed to convert `conv1d` to `conv2d`. # The 1d variant is not supported. modified_module = NeutronAtenPassManager( - neutron_target_spec, [SplitGroupConvolution(), ConvertConv1dToConv2dPass()] + neutron_target_spec, + [SplitGroupConvolution(), ConvertConv1dToConv2dPass(neutron_target_spec)], )(graph_module).graph_module # Verify that the behavior has not changed. diff --git a/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py index 26d615e156f..2c73ccd8092 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py @@ -29,6 +29,9 @@ ToNHWCPreprocess, ) from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.model_output_comparator import ( + NumericalStatsOutputComparator, +) from executorch.backends.nxp.tests.models import AvgPool2dConvModule, AvgPool2dModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare @@ -314,6 +317,23 @@ def test__basic_nsys_inference(self, mocker): model, input_shape, graph_verifier, use_new_flow_neutron_c=True ) + def test__basic_nsys_inference_qat(self, mocker): + input_shape = (2, 9, 6, 15) + model = AvgPool2dModule(False, 0) + comparator = NumericalStatsOutputComparator() + graph_verifier = DetailedGraphVerifier( + mocker, expected_delegated_ops={AvgPool2D: 1}, expected_non_delegated_ops={} + ) + + lower_run_compare( + model, + input_shape, + graph_verifier, + output_comparator=comparator, + use_new_flow_neutron_c=True, + use_qat=True, + ) + def test__kernel_size_limit(self, mocker): kernel_size = (1, 4096) input_shape = (1, 4) + kernel_size diff --git a/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py index 097b8720169..13a81c16715 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py @@ -4,9 +4,15 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch + from executorch.backends.nxp.backend.ir.conversion_config import ConversionConfig +from executorch.backends.nxp.backend.ir.converter.builder.model_builder import ( + ModelBuilder, +) from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.constant_pad_nd_converter import ( ConstantPadNDConverter, ) @@ -17,16 +23,18 @@ from executorch.backends.nxp.tests.executors import ( convert_run_compare, graph_contains_any_of_ops, + OverrideTargetSupportCheck, ToNCHWPreprocess, ToNHWCPreprocess, ) +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.models import ( ConstantPadNDConvModule, ConstantPadNDModule, ) +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.ops_aliases import ConstantPadND, Convolution from executorch.backends.nxp.tests.use_qat import * # noqa F403 -from executorch.backends.nxp.tests.executors import OverrideTargetSupportCheck -from executorch.exir.dialects._ops import ops as exir_ops @pytest.fixture(autouse=True) @@ -158,9 +166,8 @@ def test_constant_pad_nd__unsupported_paddings(input_shape, paddings, use_qat): model, input_shape, use_qat=use_qat ).exported_program() - nodes = list(exec_program.graph.nodes) # There is at least one non-delegated Pad node - assert any(node.name == "aten_constant_pad_nd_default" for node in nodes) + assert graph_contains_any_of_ops(exec_program.graph, [ConstantPadND]) def test_constant_pad_nd__delegation__formatless__supported_padding(use_qat): @@ -172,9 +179,7 @@ def test_constant_pad_nd__delegation__formatless__supported_padding(use_qat): ).exported_program() # Make sure the `pad` was delegated. - assert not graph_contains_any_of_ops( - exec_program.graph, [exir_ops.edge.aten.constant_pad_nd.default] - ) + assert not graph_contains_any_of_ops(exec_program.graph, [ConstantPadND]) def test_constant_pad_nd__delegation__formatless__unsupported_padding(use_qat): @@ -186,9 +191,7 @@ def test_constant_pad_nd__delegation__formatless__unsupported_padding(use_qat): ).exported_program() # Make sure the `pad` was NOT delegated. - assert graph_contains_any_of_ops( - exec_program.graph, [exir_ops.edge.aten.constant_pad_nd.default] - ) + assert graph_contains_any_of_ops(exec_program.graph, [ConstantPadND]) def test_constant_pad_nd__delegation__channels_first__supported_padding(use_qat): @@ -200,9 +203,7 @@ def test_constant_pad_nd__delegation__channels_first__supported_padding(use_qat) ).exported_program() # Make sure the `pad` was delegated. - assert not graph_contains_any_of_ops( - exec_program.graph, [exir_ops.edge.aten.constant_pad_nd.default] - ) + assert not graph_contains_any_of_ops(exec_program.graph, [ConstantPadND]) def test_constant_pad_nd__delegation__channels_first__unsupported_padding(use_qat): @@ -214,6 +215,122 @@ def test_constant_pad_nd__delegation__channels_first__unsupported_padding(use_qa ).exported_program() # Make sure the `pad` was NOT delegated. - assert graph_contains_any_of_ops( - exec_program.graph, [exir_ops.edge.aten.constant_pad_nd.default] + assert graph_contains_any_of_ops(exec_program.graph, [ConstantPadND]) + + +class TestConstantPadNDNewNeutronFlow: + """The PyTorch padding is added to the individual dimensions from the back (slightly confusing), see: + https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html#torch.nn.functional.pad + """ + + # noinspection PyMethodMayBeStatic + def assert_delegated(self, model, input_shape, mocker, use_qat=False): + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={ConstantPadND: 1}, + expected_non_delegated_ops={}, + ) + + lower_run_compare( + model, + input_shape, + graph_verifier, + use_qat=use_qat, + use_new_flow_neutron_c=True, + ) + + def assert_delegated_and_output_shape_equals( + self, model, input_shape, expected_output_shape, mocker + ): + model_builder_spy = mocker.spy(ModelBuilder, "finish") + + self.assert_delegated(model, input_shape, mocker) + + neutron_ir_subgraph = model_builder_spy.call_args[0][0].get_sub_graph() + assert neutron_ir_subgraph.outputs.tmp_outputs[0].shape.vector == list( + expected_output_shape + ) + + @pytest.mark.parametrize( + "input_shape, paddings", + [ + pytest.param((2,), tuple(range(2)), id="1D, padding H"), + pytest.param((2, 4), tuple(range(2)), id="2D, padding H"), + pytest.param((2, 4), tuple(range(4)), id="2D, padding N, H"), + pytest.param((2, 4, 6), tuple(range(2)), id="3D, padding H"), + pytest.param((2, 4, 6), tuple(range(4)), id="3D, padding C, H"), + pytest.param((2, 4, 6, 8), tuple(range(2)), id="4D, padding W"), + pytest.param((2, 4, 6, 8), tuple(range(4)), id="4D, padding H, W"), + pytest.param((1, 2, 3, 4, 5), tuple(range(2)), id="5D, padding D"), + pytest.param((1, 2, 3, 4, 5), tuple(range(4)), id="5D, padding W, D"), + ], + ) + def test__basic_nsys_inference(self, mocker, input_shape, paddings, use_qat): + # These test cases are also supported by the old flow. + model = ConstantPadNDModule(paddings) + self.assert_delegated(model, input_shape, mocker, use_qat) + + def test__channels_padding(self, mocker): + input_shape = (2, 4, 6) + # These paddings will be applied to the last dimension, which is the channels as the input is formatless. + paddings = (1, 1) + expected_output_shape = (2, 4, 8) # Padded channels. + model = ConstantPadNDModule(paddings) + + self.assert_delegated_and_output_shape_equals( + model, input_shape, expected_output_shape, mocker + ) + + def test__batch_padding(self, mocker): + input_shape = (2, 4, 6) + paddings = (0, 0, 0, 0, 1, 1) # Padding applied to the batch dimension. + expected_output_shape = (4, 4, 6) # Padded batch. + model = ConstantPadNDModule(paddings) + + self.assert_delegated_and_output_shape_equals( + model, input_shape, expected_output_shape, mocker + ) + + @pytest.mark.parametrize("constant", [0.0, -13.37]) + def test__specific_constant(self, mocker, constant): + input_shape = (2, 4, 6) + paddings = (1, 1) + model = ConstantPadNDModule(paddings, constant) + self.assert_delegated(model, input_shape, mocker) + + @pytest.mark.parametrize( + "input_shape, paddings", + [ + pytest.param((1, 4, 6, 8), tuple(range(2)), id="4D, padding W"), + pytest.param((1, 4, 6, 8), tuple(range(4)), id="4D, padding H, W"), + ], ) + def test__channels_first(self, mocker, input_shape, paddings): + model = ConstantPadNDConvModule(paddings) + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={ConstantPadND: 1, Convolution: 1}, + expected_non_delegated_ops={}, + ) + + lower_run_compare( + model, input_shape, graph_verifier, use_new_flow_neutron_c=True + ) + + @pytest.mark.xfail( + strict=True, + raises=RuntimeError, + reason="Known issue in Neutron: https://jira.sw.nxp.com/browse/AIR-14624", # @lint-ignore + ) + def test__bugged_channels_first_case(self, mocker): + input_shape, paddings = (1, 2, 6, 8), (0, 1, 2, 3, 1, 1) + model = ConstantPadNDConvModule(paddings) + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={ConstantPadND: 1, Convolution: 1}, + expected_non_delegated_ops={}, + ) + + lower_run_compare( + model, input_shape, graph_verifier, use_new_flow_neutron_c=True + ) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_convert_upsample_nearest2d.py b/backends/nxp/tests/ir/converter/node_converter/test_convert_upsample_nearest2d.py index 7473a5c98b0..3d9ec84dec9 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_convert_upsample_nearest2d.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_convert_upsample_nearest2d.py @@ -50,6 +50,7 @@ def forward(self, x): pytest.param((1, 8, 3, 3), 12, id="4x upscale, 8 channels, scalar size"), ], ) +@pytest.mark.xfail(strict=True, reason="EIEX-881") def test_convert_upsample_nearest2d__size(mocker, input_shape, size): model = UpsampleNearestModule(size=size) @@ -92,6 +93,7 @@ def test_convert_upsample_nearest2d__size(mocker, input_shape, size): pytest.param((1, 8, 2, 3), (4, 4), id="4x upscale, 8 channels, tuple scale"), ], ) +@pytest.mark.xfail(strict=True, reason="EIEX-881") def test_convert_upsample_nearest2d__scale_factor(mocker, input_shape, scale_factor): model = UpsampleNearestModule(scale=scale_factor) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_leaky_relu_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_leaky_relu_converter.py index 35b58c88608..9adfe992d06 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_leaky_relu_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_leaky_relu_converter.py @@ -4,18 +4,24 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) +from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import ( convert_run_compare, graph_contains_any_of_ops, ) -from executorch.exir.dialects._ops import ops as exir_ops +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.ops_aliases import ExecutorchDelegateCall, LeakyRelu +from executorch.backends.nxp.tests.use_qat import * # noqa F403 @pytest.fixture(autouse=True) @@ -24,17 +30,13 @@ def reseed_model_per_test_run(): np.random.seed(23) -ExecutorchDelegateCall = torch.ops.higher_order.executorch_call_delegate -LeakyRelu2D = exir_ops.edge.aten.leaky_relu.default - - def _assert_successful_delegation(model, input_shape, mocker, atol=0): converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") delegated_ep = to_quantized_edge_program(model, input_shape).exported_program() # Make sure the `leaky_relu` was delegated. assert graph_contains_any_of_ops(delegated_ep.graph, [ExecutorchDelegateCall]) - assert not graph_contains_any_of_ops(delegated_ep.graph, [LeakyRelu2D]) + assert not graph_contains_any_of_ops(delegated_ep.graph, [LeakyRelu]) # Verify correct behavior of the converted NeutronIR model. intermediate_ep = converter_spy.call_args.args[1] @@ -45,7 +47,7 @@ def _assert_successful_delegation(model, input_shape, mocker, atol=0): ).astype(np.int8) # Make sure the tested program contains the `leaky_relu`. - assert graph_contains_any_of_ops(intermediate_ep.graph, [LeakyRelu2D]) + assert graph_contains_any_of_ops(intermediate_ep.graph, [LeakyRelu]) convert_run_compare( intermediate_ep, tfl_model=neutron_ir_model, input_data=input_data, atol=atol @@ -121,3 +123,62 @@ def test_convert_leaky_relu__ranks(mocker, input_shape: tuple[int, ...]): mocker, atol=1, # Common quantization rounding error. ) + + +class TestLeakyReluNewNeutronFlow: + # noinspection PyMethodMayBeStatic + def assert_delegated(self, model, input_shape, mocker, use_qat=False): + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={LeakyRelu: 1}, + expected_non_delegated_ops={}, + ) + + # Create a RandomDatasetCreator that covers also negative numbers to properly test the operator. + dataset_creator = RandomDatasetCreator(low=-2, high=2) + + lower_run_compare( + model, + input_shape, + graph_verifier, + dataset_creator, + use_qat=use_qat, + use_new_flow_neutron_c=True, # Use the new flow. + ) + + @pytest.mark.parametrize( + "input_shape", + [ + (2,), + (2, 3), + (2, 3, 4), + (2, 3, 4, 5), + (2, 3, 4, 5, 6), + ], + ids=lambda shape: f"{len(shape)}D", + ) + def test__default_alpha__input_shapes(self, mocker, input_shape): + model = LeakyReluModule() + self.assert_delegated(model, input_shape, mocker) + + def test__default_alpha__qat(self, mocker, use_qat): + model = LeakyReluModule() + input_shape = (23,) + self.assert_delegated(model, input_shape, mocker, use_qat) + + @pytest.mark.parametrize( + "alpha", + [0.01, 3.14159, 0, 1, float("inf")], + ids=lambda alpha: f"alpha = {alpha}", + ) + def test__specific_alpha(self, mocker, alpha): + model = LeakyReluModule(negative_slope=alpha) + self.assert_delegated(model, (23,), mocker) + + def test__inplace(self, mocker): + model = LeakyReluModule(inplace=True) + self.assert_delegated( + model, + (23,), + mocker, + ) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_max_pool_2d_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_max_pool_2d_converter.py index 7a1c798caa3..583dc2bfd04 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_max_pool_2d_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_max_pool_2d_converter.py @@ -17,6 +17,9 @@ ToChannelLastPreprocess, ) from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.model_output_comparator import ( + NumericalStatsOutputComparator, +) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare from executorch.backends.nxp.tests.ops_aliases import ( ExecutorchDelegateCall, @@ -280,17 +283,30 @@ def test__basic_nsys_inference(self, mocker): model = MaxPool2dModule() self.assert_delegated(model, input_shape, mocker) - def test__kernel_size_limit(self, mocker): - kernel_size = (1, 4096) - input_shape = (1, 4) + kernel_size - model = MaxPool2dModule(kernel_size) - self.assert_delegated(model, input_shape, mocker) + def test__basic_nsys_inference_qat(self, mocker): + input_shape = (2, 11, 7, 16) # The old flow limited the batch size to 1. + model = MaxPool2dModule() + comparator = NumericalStatsOutputComparator() + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={MaxPool2DWithIndices: 1, GetItem: 1}, + expected_non_delegated_ops={}, + ) - def test__kernel_size_limit_exceeded(self): - kernel_size = (1, 4097) # Exceeds the kernel size limit. + lower_run_compare( + model, + input_shape, + graph_verifier, + output_comparator=comparator, + use_new_flow_neutron_c=True, + use_qat=True, + ) + + def test__large_kernel_size(self, mocker): + kernel_size = (1, 5000) input_shape = (1, 4) + kernel_size - model = MaxPool2dModule(kernel_size) - self.assert_not_delegated(model, input_shape) + model = MaxPool2dModule(kernel_size, stride=1) + self.assert_delegated(model, input_shape, mocker) def test__stride_limit__no_padding(self, mocker): stride = 4096 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py index 4e298bc9407..927af47bbf5 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py @@ -20,7 +20,10 @@ ToChannelFirstPreprocess, ToChannelLastPreprocess, ) -from executorch.backends.nxp.tests.graph_verifier import BaseGraphVerifier +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.model_output_comparator import ( + NumericalStatsOutputComparator, +) from executorch.backends.nxp.tests.models import ( MulTensorConvModule, MulTensorModule, @@ -229,19 +232,42 @@ class TestMulTensorNewNeutronFlow: pytest.param((1, 4, 8, 8), id="4D."), ], ) - def test__basic_nsys_inference(self, x_input_shape): + def test__basic_nsys_inference(self, x_input_shape, mocker): + x_input_spec = ModelInputSpec(x_input_shape) + model = MulTensorModule() + graph_verifier = DetailedGraphVerifier( + mocker, expected_delegated_ops={MulTensor: 1}, expected_non_delegated_ops={} + ) + + lower_run_compare( + model, + [x_input_spec, x_input_spec], + graph_verifier, + use_new_flow_neutron_c=True, + ) + + @pytest.mark.parametrize( + "x_input_shape", + [ + pytest.param((1, 4, 8), id="3D."), + pytest.param((1, 4, 8, 8), id="4D."), + ], + ) + def test__basic_nsys_inference_qat(self, x_input_shape, mocker): x_input_spec = ModelInputSpec(x_input_shape) model = MulTensorModule() - graph_verifier = BaseGraphVerifier( - exp_num_delegate_call_nodes=1, - exp_non_delegated_nodes=[], + comparator = NumericalStatsOutputComparator() + graph_verifier = DetailedGraphVerifier( + mocker, expected_delegated_ops={MulTensor: 1}, expected_non_delegated_ops={} ) lower_run_compare( model, [x_input_spec, x_input_spec], graph_verifier, + output_comparator=comparator, use_new_flow_neutron_c=True, + use_qat=True, ) @pytest.mark.parametrize( @@ -259,11 +285,10 @@ def test__basic_nsys_inference(self, x_input_shape): ), ], ) - def test__correct_broadcast(self, input_spec): + def test__correct_broadcast(self, input_spec, mocker): model = MulTensorModule() - graph_verifier = BaseGraphVerifier( - exp_num_delegate_call_nodes=1, - exp_non_delegated_nodes=[], + graph_verifier = DetailedGraphVerifier( + mocker, expected_delegated_ops={MulTensor: 1}, expected_non_delegated_ops={} ) lower_run_compare( @@ -308,16 +333,17 @@ def test__incorrect_broadcast(self, input_spec): ), ], ) - def test__w_conv(self, x_input_shape): + def test__w_conv(self, x_input_shape, mocker): model = MulTensorConvModule() n, c, h, w = x_input_shape y_input_spec = ModelInputSpec((n, 8, h, w)) x_input_spec = ModelInputSpec(x_input_shape) - graph_verifier = BaseGraphVerifier( - exp_num_delegate_call_nodes=1, - exp_non_delegated_nodes=[], + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={MulTensor: 1, Convolution: 1}, + expected_non_delegated_ops={}, ) lower_run_compare( diff --git a/backends/nxp/tests/ir/converter/node_converter/test_sigmoid_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_sigmoid_converter.py index ad03aa18ded..fd7f2ba6a9d 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_sigmoid_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_sigmoid_converter.py @@ -1,23 +1,33 @@ -# Copyright 2025 NXP +# Copyright 2025-2026 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) +from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner +from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import ( convert_run_compare, ToNCHWPreprocess, ToNHWCPreprocess, ) +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.model_output_comparator import ( + AllCloseOutputComparator, +) from executorch.backends.nxp.tests.models import ConvWithSigmoid +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.ops_aliases import DequantizePerTensor, Sigmoid from torch import nn from torch.export import ExportedProgram from executorch.backends.nxp.tests.use_qat import * # noqa F403 @@ -76,3 +86,60 @@ def test_sigmoid_only(mocker, use_qat, input_shape): convert_run_compare( exported_program, tfl_model=tflite_flatbuffers_model, input_data=input_data ) + + +class TestSigmoidNewNeutronFlow: + # noinspection PyMethodMayBeStatic + def assert_delegated(self, model, input_shape, mocker, use_qat=False, atol=None): + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={Sigmoid: 1}, + expected_non_delegated_ops={}, + ) + + # Create a RandomDatasetCreator that covers also negative numbers to properly test the operator. + dataset_creator = RandomDatasetCreator(low=-2, high=2) + + kwargs = {"atol": atol} if atol is not None else {} + output_comparator = AllCloseOutputComparator(**kwargs) + + lower_run_compare( + model, + input_shape, + graph_verifier, + dataset_creator, + output_comparator, + use_qat=use_qat, + use_new_flow_neutron_c=True, # Use the new flow. + ) + + def test__basic_nsys_inference__qat(self, mocker, use_qat): + input_shape = (23,) + model = nn.Sigmoid() + self.assert_delegated(model, input_shape, mocker, use_qat=use_qat) + + @pytest.mark.parametrize( + "input_shape", + [ + (2,), + (2, 3), + (2, 3, 4), + (2, 3, 4, 5), + (2, 3, 4, 5, 6), + ], + ids=lambda shape: f"{len(shape)}D", + ) + def test__input_shapes(self, mocker, input_shape): + model = nn.Sigmoid() + + output_scale = 1.0 / 256.0 + lowering_spy = mocker.spy(NeutronPartitioner, "partition") + self.assert_delegated( + model, input_shape, mocker, atol=output_scale + ) # Allow single bit error. + + # Verify that the `atol` is indeed equal to the output scale. + # In the near future, we would like to add support for testing with int8 IO, where this check will be trivial. + nodes = list(lowering_spy.spy_return.tagged_exported_program.graph.nodes) + assert nodes[-2].target == DequantizePerTensor + assert nodes[-2].args[1] == output_scale diff --git a/backends/nxp/tests/model_output_comparator.py b/backends/nxp/tests/model_output_comparator.py index 4efa01011b8..f0dd7cd2d60 100644 --- a/backends/nxp/tests/model_output_comparator.py +++ b/backends/nxp/tests/model_output_comparator.py @@ -91,7 +91,13 @@ def compare_sample(self, sample_dir, cpu_output_tensors, npu_output_tensors): assert np.any( cpu_tensor ), "Output tensor contains only zeros. This is suspicious." - assert np.allclose(cpu_tensor, npu_tensor, atol=self.atol) + all_close = np.allclose(cpu_tensor, npu_tensor, atol=self.atol) + if not all_close: + max_diff = np.abs(cpu_tensor - npu_tensor).max() + print( + f"NPU output doesn't match reference. Maximum absolute difference: {max_diff}" + ) + assert all_close def _default_postprocess_fn(outputs: np.ndarray, _: str): diff --git a/backends/nxp/tests/nsys_testing.py b/backends/nxp/tests/nsys_testing.py index 9d25c309d25..636e1a28a44 100644 --- a/backends/nxp/tests/nsys_testing.py +++ b/backends/nxp/tests/nsys_testing.py @@ -608,6 +608,7 @@ def _get_caller_name(): def execute_cmd(cmd, cwd="."): env = environ.copy() # Copy the current environment env["LD_LIBRARY_PATH"] = str(NSYS_PATH.parent) + logger.debug(f"Running command: {cmd}") with subprocess.Popen( cmd, diff --git a/backends/nxp/tests/ops_aliases.py b/backends/nxp/tests/ops_aliases.py index f190ca91e1f..584ce1ea2e6 100644 --- a/backends/nxp/tests/ops_aliases.py +++ b/backends/nxp/tests/ops_aliases.py @@ -14,6 +14,7 @@ Abs = exir_ops.edge.aten.abs.default AvgPool2D = exir_ops.edge.aten.avg_pool2d.default Bmm = exir_ops.edge.aten.bmm.default +ConstantPadND = exir_ops.edge.aten.constant_pad_nd.default Convolution = exir_ops.edge.aten.convolution.default DequantizePerChannel = exir_ops.edge.quantized_decomposed.dequantize_per_channel.default DequantizePerTensor = exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default @@ -21,11 +22,13 @@ GetItem = operator.getitem HardTanh = exir_ops.edge.aten.hardtanh.default HardTanh_ = exir_ops.edge.aten.hardtanh_.default +LeakyRelu = exir_ops.edge.aten.leaky_relu.default MaxPool2DWithIndices = exir_ops.edge.aten.max_pool2d_with_indices.default MulTensor = exir_ops.edge.aten.mul.Tensor QuantizePerChannel = exir_ops.edge.quantized_decomposed.quantize_per_channel.default QuantizePerTensor = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default Relu = exir_ops.edge.aten.relu.default +Sigmoid = exir_ops.edge.aten.sigmoid.default Slice = exir_ops.edge.aten.slice.Tensor SliceCopy = exir_ops.edge.aten.slice_copy.Tensor Softmax = exir_ops.edge.aten._softmax.default diff --git a/backends/nxp/tests/outputs_dir_importer.py b/backends/nxp/tests/outputs_dir_importer.py index c018123c949..425bf44a31e 100644 --- a/backends/nxp/tests/outputs_dir_importer.py +++ b/backends/nxp/tests/outputs_dir_importer.py @@ -8,7 +8,7 @@ logger = logging.getLogger(__name__) try: - import test.python.outputs_dir as outputs_dir # noqa F401 + import outputs_dir as outputs_dir # noqa F401 logger.debug("Importing from executorch-integration") except ImportError: diff --git a/backends/nxp/tests/test_convert_1d_conv_to_2d.py b/backends/nxp/tests/test_convert_1d_conv_to_2d.py index 9a1d4d3f91f..27e16a0358c 100644 --- a/backends/nxp/tests/test_convert_1d_conv_to_2d.py +++ b/backends/nxp/tests/test_convert_1d_conv_to_2d.py @@ -24,6 +24,7 @@ ) from executorch.backends.nxp.tests.models import Conv1dModule, ConvTranspose1dModule from executorch.exir.dialects._ops import ops as exir_ops +from torch import nn from torch.export import ExportedProgram @@ -39,6 +40,11 @@ def reseed_model_per_test_run(): AtenConvTranspose2d = torch.ops.aten.conv_transpose2d.input AtenSqueeze = torch.ops.aten.squeeze.dim AtenUnsqueeze = torch.ops.aten.unsqueeze.default +AtenRelu = torch.ops.aten.relu.default +AtenSigmoid = torch.ops.aten.sigmoid.default +AtenTanh = torch.ops.aten.tanh.default +AtenHardtanh = torch.ops.aten.hardtanh.default +AtenBatchNorm = torch.ops.aten.batch_norm.default EdgeConvolution = exir_ops.edge.aten.convolution.default ExecutorchDelegateCall = torch.ops.higher_order.executorch_call_delegate @@ -99,9 +105,9 @@ def test_convert_conv_1d_to_conv2d( outputs_before = [o.detach().numpy() for o in exir_program_aten(example_input)] # Apply the optimization. - NeutronAtenPassManager(neutron_target_spec, [ConvertConv1dToConv2dPass()])( - exir_program_aten - ) + NeutronAtenPassManager( + neutron_target_spec, [ConvertConv1dToConv2dPass(neutron_target_spec)] + )(exir_program_aten) # Make sure no `aten.conv1d` nodes are in the model. assert not graph_contains_any_of_ops( @@ -207,9 +213,9 @@ def test_convert_conv_1d_transp_to_conv2d_transp( outputs_before = [o.detach().numpy() for o in exir_program_aten(example_input)] # Apply the optimization. - NeutronAtenPassManager(neutron_target_spec, [ConvertConv1dToConv2dPass()])( - exir_program_aten - ) + NeutronAtenPassManager( + neutron_target_spec, [ConvertConv1dToConv2dPass(neutron_target_spec)] + )(exir_program_aten) # Make sure no `aten.conv_transpose1d` nodes are in the model. assert not graph_contains_any_of_ops( @@ -393,3 +399,140 @@ def test_convert_conv_1d_to_conv2d_transp_full_pipeline( input_data=example_input, tfl_model=neutron_ir_model, ) + + +class Conv1dActivationModule(nn.Module): + def __init__( + self, in_channels, out_channels, kernel_size, activation, stride=1, padding=0 + ): + super().__init__() + self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding) + self.activation = activation + + def forward(self, x): + return self.activation(self.conv(x)) + + +class Conv1dBNActivationModule(nn.Module): + def __init__( + self, in_channels, out_channels, kernel_size, activation, stride=1, padding=0 + ): + super().__init__() + self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding) + self.bn = nn.BatchNorm1d(out_channels) + self.activation = activation + + def forward(self, x): + return self.activation(self.bn(self.conv(x))) + + +class Conv1dHardtanhUnsupportedModule(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0): + super().__init__() + self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding) + self.hardtanh = nn.Hardtanh(min_val=-1.0, max_val=1.0) + + def forward(self, x): + return self.hardtanh(self.conv(x)) + + +@pytest.mark.parametrize( + "activation, expected_act_target, has_bn", + [ + pytest.param(nn.ReLU(), AtenRelu, False, id="conv1d_relu"), + pytest.param(nn.ReLU(), AtenRelu, True, id="conv1d_bn_relu"), + pytest.param(nn.Sigmoid(), AtenSigmoid, False, id="conv1d_sigmoid"), + pytest.param(nn.Sigmoid(), AtenSigmoid, True, id="conv1d_bn_sigmoid"), + pytest.param(nn.Tanh(), AtenTanh, False, id="conv1d_tanh"), + pytest.param(nn.Tanh(), AtenTanh, True, id="conv1d_bn_tanh"), + pytest.param( + nn.Hardtanh(min_val=0.0, max_val=6.0), + AtenHardtanh, + False, + id="conv1d_relu6", + ), + pytest.param( + nn.Hardtanh(min_val=0.0, max_val=6.0), + AtenHardtanh, + True, + id="conv1d_bn_relu6", + ), + ], +) +def test_convert_conv_1d_to_conv2d_keeps_activation_in_4d( + activation, expected_act_target, has_bn +): + input_shape = (3, 7, 23) + model_cls = Conv1dBNActivationModule if has_bn else Conv1dActivationModule + model = model_cls( + in_channels=7, out_channels=14, kernel_size=3, activation=activation, padding=1 + ) + example_input = torch.rand(input_shape) + + exir_program_aten = torch.export.export(model, (example_input,)).module() + + assert graph_contains_any_of_ops(exir_program_aten.graph, [AtenConv1d]) + outputs_before = [o.detach().numpy() for o in exir_program_aten(example_input)] + + NeutronAtenPassManager( + neutron_target_spec, [ConvertConv1dToConv2dPass(neutron_target_spec)] + )(exir_program_aten) + + assert not graph_contains_any_of_ops(exir_program_aten.graph, [AtenConv1d]) + + nodes = list(exir_program_aten.graph.nodes) + conv_nodes = [i for i, n in enumerate(nodes) if n.target == AtenConv2d] + assert len(conv_nodes) == 1 + i = conv_nodes[0] + + assert nodes[i - 1].target == AtenUnsqueeze + assert nodes[i].target == AtenConv2d + + if has_bn: + assert nodes[i + 1].target == AtenBatchNorm + assert nodes[i + 2].target == expected_act_target + assert nodes[i + 3].target == AtenSqueeze + else: + assert nodes[i + 1].target == expected_act_target + assert nodes[i + 2].target == AtenSqueeze + + outputs_after = [o.detach().numpy() for o in exir_program_aten(example_input)] + + assert len(outputs_before) == len(outputs_after) + for j in range(len(outputs_before)): + assert np.allclose(outputs_before[j], outputs_after[j]) + + +def test_convert_conv_1d_to_conv2d_unsupported_hardtanh_not_fused(): + input_shape = (3, 7, 23) + model = Conv1dHardtanhUnsupportedModule( + in_channels=7, out_channels=14, kernel_size=3, padding=1 + ) + example_input = torch.rand(input_shape) + + exir_program_aten = torch.export.export(model, (example_input,)).module() + + assert graph_contains_any_of_ops(exir_program_aten.graph, [AtenConv1d]) + outputs_before = [o.detach().numpy() for o in exir_program_aten(example_input)] + + NeutronAtenPassManager( + neutron_target_spec, [ConvertConv1dToConv2dPass(neutron_target_spec)] + )(exir_program_aten) + + assert not graph_contains_any_of_ops(exir_program_aten.graph, [AtenConv1d]) + + nodes = list(exir_program_aten.graph.nodes) + conv_nodes = [i for i, n in enumerate(nodes) if n.target == AtenConv2d] + assert len(conv_nodes) == 1 + i = conv_nodes[0] + + assert nodes[i - 1].target == AtenUnsqueeze + assert nodes[i].target == AtenConv2d + assert nodes[i + 1].target == AtenSqueeze + assert nodes[i + 2].target == AtenHardtanh + + outputs_after = [o.detach().numpy() for o in exir_program_aten(example_input)] + + assert len(outputs_before) == len(outputs_after) + for j in range(len(outputs_before)): + assert np.allclose(outputs_before[j], outputs_after[j]) diff --git a/backends/qualcomm/_passes/__init__.py b/backends/qualcomm/_passes/__init__.py index fc59ce3d262..f82157d3cf0 100644 --- a/backends/qualcomm/_passes/__init__.py +++ b/backends/qualcomm/_passes/__init__.py @@ -32,6 +32,7 @@ from .decompose_remainder import DecomposeRemainder from .decompose_roll import DecomposeRoll from .decompose_silu import DecomposeSilu +from .decompose_tan import DecomposeTan from .decompose_threshold import DecomposeThreshold from .decompose_triu import DecomposeTriu from .decompose_trunc import DecomposeTrunc @@ -88,6 +89,7 @@ DecomposeRemainder, DecomposeRoll, DecomposeSilu, + DecomposeTan, DecomposeThreshold, DecomposeTriu, DecomposeTrunc, diff --git a/backends/qualcomm/_passes/decompose_tan.py b/backends/qualcomm/_passes/decompose_tan.py new file mode 100644 index 00000000000..b75cf9ff2df --- /dev/null +++ b/backends/qualcomm/_passes/decompose_tan.py @@ -0,0 +1,71 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.pass_base import ExportPass, PassResult + +from .utils import copy_meta + + +class DecomposeTan(ExportPass): + """ + Decompose tan(x) = sin(x) / cos(x) + """ + + def __init__(self): + super(DecomposeTan, self).__init__() + self.targets = { + torch.ops.aten.tan.default, + exir_ops.edge.aten.tan.default, + } + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + graph = graph_module.graph + + for node in list(graph.nodes): + if node.op == "call_function" and node.target in self.targets: + is_edge = isinstance(node.target, EdgeOpOverload) + + sin_op = ( + exir_ops.edge.aten.sin.default + if is_edge + else torch.ops.aten.sin.default + ) + cos_op = ( + exir_ops.edge.aten.cos.default + if is_edge + else torch.ops.aten.cos.default + ) + div_op = ( + exir_ops.edge.aten.div.Tensor + if is_edge + else torch.ops.aten.div.Tensor + ) + + with graph.inserting_before(node): + sin_node = graph.create_node( + "call_function", sin_op, (node.args[0],) + ) + sin_node.meta = copy_meta(node.meta) + + cos_node = graph.create_node( + "call_function", cos_op, (node.args[0],) + ) + cos_node.meta = copy_meta(node.meta) + + div_node = graph.create_node( + "call_function", div_op, (sin_node, cos_node) + ) + div_node.meta = copy_meta(node.meta) + + for user in node.users.copy(): + user.replace_input_with(node, div_node) + + graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/layout_transform.py b/backends/qualcomm/_passes/layout_transform.py index 9422051addd..5b9c13e6ef4 100644 --- a/backends/qualcomm/_passes/layout_transform.py +++ b/backends/qualcomm/_passes/layout_transform.py @@ -113,11 +113,14 @@ class LayoutTransform(ExportPass): exir_ops.edge.aten.neg.default, exir_ops.edge.aten.pow.Tensor_Scalar, exir_ops.edge.aten.prelu.default, + exir_ops.edge.aten.rand.default, + exir_ops.edge.aten.randn.default, exir_ops.edge.aten.reflection_pad1d.default, exir_ops.edge.aten.reflection_pad2d.default, exir_ops.edge.aten.repeat.default, exir_ops.edge.aten.relu.default, exir_ops.edge.aten.round.default, + exir_ops.edge.aten.scatter.src, exir_ops.edge.aten.sigmoid.default, exir_ops.edge.aten.sign.default, exir_ops.edge.aten.slice_copy.Tensor, diff --git a/backends/qualcomm/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index 57354af11de..b0913bbefd9 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/backends/qualcomm/_passes/qnn_pass_manager.py @@ -37,6 +37,7 @@ DecomposeRemainder, DecomposeRoll, DecomposeSilu, + DecomposeTan, DecomposeThreshold, DecomposeTriu, DecomposeTrunc, @@ -112,6 +113,7 @@ def get_capture_program_passes(): (DecomposeMinMaxDim, True), (DecomposePad, True), (DecomposeRemainder, True), + (DecomposeTan, True), (DecomposeTrunc, True), (ExpandBroadcastTensorShape, True), (FixedLinearKeepDim, True), @@ -236,6 +238,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule): self.add_pass(DecomposeScaledDotProductAttention()) self.add_pass(DecomposeRoll()) self.add_pass(DecomposeSilu()) + self.add_pass(DecomposeTan()) self.add_pass(DecomposeThreshold()) self.add_pass(DecomposeTriu()) self.add_pass(DecomposeTrunc()) diff --git a/backends/qualcomm/_passes/utils.py b/backends/qualcomm/_passes/utils.py index 04371d61e1c..542fa1115a6 100755 --- a/backends/qualcomm/_passes/utils.py +++ b/backends/qualcomm/_passes/utils.py @@ -74,6 +74,7 @@ def get_passes_dependency_for_capture_program(): DecomposeMaxPool3d, DecomposePad, DecomposeRemainder, + DecomposeTan, DecomposeTrunc, ExpandBroadcastTensorShape, FixedLinearKeepDim, @@ -107,6 +108,7 @@ def get_passes_dependency_for_capture_program(): DecomposeMaxPool3d: [RemoveRedundancy], DecomposePad: [RemoveRedundancy], DecomposeRemainder: [RemoveRedundancy], + DecomposeTan: [RemoveRedundancy], DecomposeTrunc: [RemoveRedundancy], ExpandBroadcastTensorShape: [FoldQDQ], FixedLinearKeepDim: [FoldQDQ], diff --git a/backends/qualcomm/builders/README.md b/backends/qualcomm/builders/README.md index d71aeb27be6..a443ed0905c 100644 --- a/backends/qualcomm/builders/README.md +++ b/backends/qualcomm/builders/README.md @@ -368,7 +368,7 @@ Please help update following table if you are contributing new operators: + 🚫 = Deprecated, supported with other QNN Ops -| Operators | HTP - 98/119 Enabled | +| Operators | HTP - 100/120 Enabled | |-----------|---------| | Argmax | ✓ | | Argmin | ✓ | @@ -457,7 +457,8 @@ Please help update following table if you are contributing new operators: | PoolMax2d | ✓ | | Prelu | ✓ | | Quantize | ✓ | -| Rand | ✓ | +| RandomUniformLike | ✓ | +| RandomNormalLike | ✓ | | ReduceMax | ✓ | | ReduceMean | ✓ | | ReduceMin | ✓ | @@ -472,7 +473,7 @@ Please help update following table if you are contributing new operators: | ResizeNearestNeighbor | ✓ | | RoiAlign | ✗ | | RmsNorm | ✓ | -| ScatterElements | ✗ | +| ScatterElements | ✓ | | ScatterNd | ✓ | | Sigmoid | ✓ | | Softmax | ✓ | @@ -517,6 +518,7 @@ The following PyTorch operators are supported through decomposition or annotatio | `aten.remainder.Scalar`, `aten.remainder.Tensor` | `DecomposeRemainder` | | `aten.roll` | `DecomposeRoll` | | `aten.silu` | `DecomposeSilu` | +| `aten.tan` | `DecomposeTan` | | `aten.threshold` | `DecomposeThreshold` | | `aten.triu` | `DecomposeTriu` | | `aten.trunc` | `DecomposeTrunc` | diff --git a/backends/qualcomm/builders/__init__.py b/backends/qualcomm/builders/__init__.py index a897dfa53bd..dd69515b1c1 100644 --- a/backends/qualcomm/builders/__init__.py +++ b/backends/qualcomm/builders/__init__.py @@ -81,6 +81,7 @@ op_prelu, op_quantize, op_rand, + op_randn, op_relu, op_repeat, op_reshape, @@ -89,6 +90,7 @@ op_round, op_rsqrt, op_scalar_tensor, + op_scatter_elements, op_select_copy, op_sigmoid, op_sign, @@ -194,6 +196,7 @@ op_prelu, op_quantize, op_rand, + op_randn, op_relu, op_repeat, op_reshape, @@ -202,6 +205,7 @@ op_round, op_rsqrt, op_scalar_tensor, + op_scatter_elements, op_select_copy, op_sigmoid, op_sign, diff --git a/backends/qualcomm/builders/node_visitor.py b/backends/qualcomm/builders/node_visitor.py index 7380f7a8191..c206950c140 100644 --- a/backends/qualcomm/builders/node_visitor.py +++ b/backends/qualcomm/builders/node_visitor.py @@ -418,12 +418,12 @@ def get_tensor_name( elif is_graph_output(node): tensor_name = f"output_{tensor_name}" - # Save this for intermediate debugger - # Needs idx since node like topk has 2 outputs - if QCOM_TENSOR_NAME in node.meta: - node.meta[QCOM_TENSOR_NAME][wrapper_idx] = tensor_name - else: - node.meta[QCOM_TENSOR_NAME] = {wrapper_idx: tensor_name} + # Only add qcom_tensor_name when enable tensor dump. + # Only runs in qnn_preprocess (not op validation) since that's when + # tensor names are finalized and enable_tensor_dump is True. + if self.enable_tensor_dump: + node.meta.setdefault(QCOM_TENSOR_NAME, {})[wrapper_idx] = tensor_name + return tensor_name def define_custom_tensor_wrapper( diff --git a/backends/qualcomm/builders/op_full.py b/backends/qualcomm/builders/op_full.py index 5ac2e95c57b..7a109ff0637 100644 --- a/backends/qualcomm/builders/op_full.py +++ b/backends/qualcomm/builders/op_full.py @@ -25,8 +25,9 @@ def define_node( node: torch.fx.Node, nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper], ) -> PyQnnManager.PyQnnOpWrapper: + tensor_shape = list(self.get_tensor(node, node).shape) out_tensor = torch.full( - node.args[0], node.args[1], dtype=node.meta["val"].dtype + tensor_shape, node.args[1], dtype=node.meta["val"].dtype ) # since we can derive the constant value of current op in AoT stage diff --git a/backends/qualcomm/builders/op_full_like.py b/backends/qualcomm/builders/op_full_like.py index 66f80ecc80a..69a03d66b13 100644 --- a/backends/qualcomm/builders/op_full_like.py +++ b/backends/qualcomm/builders/op_full_like.py @@ -25,8 +25,8 @@ def define_node( node: torch.fx.Node, nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper], ) -> PyQnnManager.PyQnnOpWrapper: - in_tensor = node.args[0].meta["val"] - ref_tensor = torch.zeros(in_tensor.shape, dtype=in_tensor.dtype) + in_tensor = self.get_tensor(node, node) + ref_tensor = torch.zeros(list(in_tensor.shape), dtype=in_tensor.dtype) out_tensor = torch.full_like(ref_tensor, node.args[1]) # since we can derive the constant value of current op in AoT stage diff --git a/backends/qualcomm/builders/op_randn.py b/backends/qualcomm/builders/op_randn.py new file mode 100644 index 00000000000..6160fc79609 --- /dev/null +++ b/backends/qualcomm/builders/op_randn.py @@ -0,0 +1,79 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Dict + +import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManager + +import numpy as np +import torch +from executorch.backends.qualcomm.utils.constants import QCOM_DATA + +from .node_visitor import NodeVisitor +from .node_visitor_manager import register_node_visitor +from .qnn_constants import OpRandomNormalLike, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class Randn(NodeVisitor): + target = ["aten.randn.default", "aten.randn_like.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper], + ) -> PyQnnManager.PyQnnOpWrapper: + output_tensor = node.meta["val"] + output_shape = list(output_tensor.shape) + + shape_data = np.array(output_shape, dtype=np.uint32) + shape_dims = [len(output_shape)] + + shape_tensor_wrapper = PyQnnManager.TensorWrapper( + f"{node.name}_shape", + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC, + PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_UINT_32, + PyQnnManager.Qnn_QuantizationEncoding_t.QNN_QUANTIZATION_ENCODING_UNDEFINED, + {}, + len(shape_dims), + shape_dims, + [], + shape_data, + True, + ) + + output_tensor_wrapper = self.define_tensor( + node, + node, + output_tensor, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + randn_op = PyQnnManager.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpRandomNormalLike.op_name, + ) + + randn_op.AddInputTensors([shape_tensor_wrapper]) + randn_op.AddOutputTensors([output_tensor_wrapper]) + + randn_op.AddScalarParam( + OpRandomNormalLike.param_mean, + PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_FLOAT_32, + {QCOM_DATA: np.float32(0.0)}, + ) + + randn_op.AddScalarParam( + OpRandomNormalLike.param_scale, + PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_FLOAT_32, + {QCOM_DATA: np.float32(1.0)}, + ) + + return randn_op diff --git a/backends/qualcomm/builders/op_scatter_elements.py b/backends/qualcomm/builders/op_scatter_elements.py new file mode 100644 index 00000000000..4bcf4572803 --- /dev/null +++ b/backends/qualcomm/builders/op_scatter_elements.py @@ -0,0 +1,103 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Dict + +import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManager + +import numpy as np +import torch +from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER, QCOM_DATA + +from .node_visitor import NodeVisitor +from .node_visitor_manager import register_node_visitor +from .qnn_constants import OpScatterElements, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class ScatterElements(NodeVisitor): + target = ["aten.scatter.src"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper], + ) -> PyQnnManager.PyQnnOpWrapper: + input_node = self.get_node(node.args[0]) + input_tensor = self.get_tensor(input_node, node) + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + index_node = self.get_node(node.args[2]) + index_tensor = self.get_tensor(index_node, node) + index_tensor_wrapper = self.define_tensor( + index_node, + node, + index_tensor.to(torch.int32), + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + updates_node = self.get_node(node.args[3]) + updates_tensor = self.get_tensor(updates_node, node) + updates_tensor_wrapper = self.define_tensor( + updates_node, + node, + updates_tensor, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + output_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + output_tensor, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + dim = node.args[1] + if dim < 0: + dim = dim % len(input_tensor.shape) + + if QCOM_AXIS_ORDER in node.meta: + dim = node.meta[QCOM_AXIS_ORDER].index(dim) + + scatter_op = PyQnnManager.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpScatterElements.op_name, + ) + scatter_op.AddInputTensors( + [ + input_tensor_wrapper, + index_tensor_wrapper, + updates_tensor_wrapper, + ] + ) + scatter_op.AddOutputTensors([output_tensor_wrapper]) + + scatter_op.AddScalarParam( + OpScatterElements.param_axis, + PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_UINT_32, + {QCOM_DATA: np.uint32(dim)}, + ) + + scatter_op.AddScalarParam( + OpScatterElements.param_reduction, + PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_UINT_32, + {QCOM_DATA: np.uint32(OpScatterElements.Reduction.NONE)}, + ) + + return scatter_op diff --git a/backends/qualcomm/builders/qnn_constants.py b/backends/qualcomm/builders/qnn_constants.py index d7ec30fddc0..d1f0d3fff00 100644 --- a/backends/qualcomm/builders/qnn_constants.py +++ b/backends/qualcomm/builders/qnn_constants.py @@ -504,6 +504,13 @@ class OpQuantize: op_name: str = "Quantize" +@dataclass(init=False, frozen=True) +class OpRandomNormalLike: + op_name: str = "RandomNormalLike" + param_mean: str = "mean" + param_scale: str = "scale" + + @dataclass(init=False, frozen=True) class OpRandomUniformLike: op_name: str = "RandomUniformLike" @@ -587,6 +594,17 @@ class OpRmsNorm: param_axes: str = "axes" +@dataclass(init=False, frozen=True) +class OpScatterElements: + op_name: str = "ScatterElements" + param_axis: str = "axis" + param_reduction: str = "reduction" + + @unique + class Reduction(IntEnum): + NONE = 0 + + @dataclass(init=False, frozen=True) class OpScatterNd: op_name: str = "ScatterNd" diff --git a/backends/qualcomm/debugger/README.md b/backends/qualcomm/debugger/README.md index fb8f9a1c662..8300920d1d5 100644 --- a/backends/qualcomm/debugger/README.md +++ b/backends/qualcomm/debugger/README.md @@ -50,7 +50,7 @@ Generate optrace and QHAS files using QNN tools under $QNN_SDK_ROOT. After finis adb = SimpleADB( qnn_config=qnn_config, pte_path=f"{args.artifact}/{pte_filename}.pte", - workspace=f"/data/local/tmp/executorch/{pte_filename}, + workspace=f"/data/local/tmp/executorch/{pte_filename}", ) binaries_trace = generate_optrace( args, adb, f"{args.artifact}/{pte_filename}.pte", example_input @@ -78,7 +78,7 @@ qairt_visualizer.view(reports=[optrace, qhas]) - `model`: Path to your QNN model file (e.g., `path_to_your_model.dlc`). - **`reports`**: List of report file paths, including the optrace (`optrace.json`) and QHAS (`optrace_qnn_htp_analysis_summary.json`). -Note: Files ending with `.bin ` do not support graph visualization in qairt_visualizer. +Note: Files ending with `.bin` do not support graph visualization in qairt_visualizer. ## Demo @@ -121,24 +121,24 @@ flowchart TB; debug --> output["Output Results"] ``` -## Instructions - -### 1. Setup +## Prerequisites 1. Follow the [tutorial](https://pytorch.org/executorch/main/getting-started-setup) to set up ExecuTorch. 2. Follow the [tutorial](https://pytorch.org/executorch/stable/build-run-qualcomm-ai-engine-direct-backend.html) to build Qualcomm AI Engine Direct Backend. -### 2. Enable Flag +## Instructions -When executing the script, please add the flag `--dump_intermediate_outputs`. This tells QNN to dump all intermediate tensors during execution. +### 1. Initialize debugger and build binary + +Create a `QNNIntermediateDebugger` with a sample input and pass it to `build_executorch_binary`. The `--dump_intermediate_outputs` flag tells QNN to dump all intermediate tensors during execution. -### 3. Add debugger to the example script -Initialize a `QNNIntermediateDebugger`. Please pass initialized `QNNIntermediateDebugger` and the `args.dump_intermediate_outputs` to `build_executorch_binary` method as well. -#### Example: ```python from executorch.backends.qualcomm.export_utils import build_executorch_binary -from executorch.backends.qualcomm.debugger.qnn_intermediate_debugger import QNNIntermediateDebugger +from executorch.backends.qualcomm.debugger.qnn_intermediate_debugger import ( + OutputFormat, + QNNIntermediateDebugger, +) -qnn_intermediate_debugger = QNNIntermediateDebugger() +qnn_intermediate_debugger = QNNIntermediateDebugger(sample_input=inputs[0]) build_executorch_binary( model=MyModel(), qnn_config=qnn_config, @@ -148,27 +148,61 @@ build_executorch_binary( ) ``` -### 4. Set data num to 1 -It is perfectly fine for users to pass the desired amount of datasets to `build_executorch_binary`, which helps achieve better quantization results. However, after `build_executorch_binary` is called, we need to ensure that we only perform one inference during execution. Please ensure that CPU and QNN is using the same input during execution; otherwise, the debugging results might not be accurate. +After `build_executorch_binary()`, the debugger holds: +- `edge_ep` — edge `ExportedProgram` for CPU golden inference. +- `etrecord_file_path` — path to the generated ET record. + +### 2. Execute on device + +Ensure `dump_intermediate_outputs` is enabled in your `QnnConfig` (or pass `--dump_intermediate_outputs` via CLI). Only run **one inference** for debugging — multiple executions are not supported. + +```python +from executorch.examples.qualcomm.utils import SimpleADB + +adb = SimpleADB( + qnn_config=qnn_config, + pte_path=f"{args.artifact}/{pte_filename}.pte", + workspace=f"/data/local/tmp/executorch/{pte_filename}", +) +adb.push(inputs=inputs) +adb.execute() +``` + +### 3. Pull results and compare + +After execution, pull `etdump.etdp` and `debug_output.bin` from the device. Use `setup_inspector()` to create the `Inspector`, then create comparators and generate results. + +Before comparing per-layer outputs, it is highly recommended to verify that the edge program's final output aligns with the original `nn.Module`. The debugger uses the edge program as the CPU golden reference, so if the edge graph itself has diverged (e.g., due to weights quantization or pass transformations), per-layer comparisons against it may be misleading. -### 5: Pull and process the results. -After QNN execution with the runner, if the previous steps are done correctly, we should be able to get two files: `etdump.etdp` and `debug_output.bin`. -The following example pulls the files back and calls a callback function to process the results. In this callback function, we create the `Inspector`. Then we perform CPU inference to get CPU intermediate results. Now, we have both QNN and CPU intermediate results, we can start generating results to compare the accuracy. Taking the following example, we should be able to get `debug_graph.svg` as an output in the current directory. -#### Example: ```python -from executorch.backends.qualcomm.debugger.qnn_intermediate_debugger import OutputFormat +from executorch.backends.qualcomm.debugger.qcom_numerical_comparator_sample import ( + QcomCosineSimilarityComparator, QcomMSEComparator, +) + def validate_intermediate_tensor(): - inspector = Inspector( + qnn_intermediate_debugger.setup_inspector( etdump_path=f"{args.artifact}/etdump.etdp", debug_buffer_path=f"{args.artifact}/debug_output.bin", ) - qnn_intermediate_debugger.intermediate_output_module(*(inputs[0])) + + # Verify edge program output aligns with the original nn.Module. + # This ensures the edge graph is a reliable golden reference. + edge_result = qnn_intermediate_debugger.edge_ep.module()(*(inputs[0])) + with torch.no_grad(): + source_result = source_model(*(inputs[0])) + score = torch.nn.functional.cosine_similarity( + edge_result.flatten(), source_result.flatten(), dim=0 + ).item() + print("Cosine similarity between nn.Module and edge CPU:", score) + + cos_comparator = qnn_intermediate_debugger.create_comparator( + QcomCosineSimilarityComparator, threshold=0.9 + ) qnn_intermediate_debugger.generate_results( - title="debug_graph", - path=".", - output_format=OutputFormat.SVG_GRAPHS, - inspector=inspector, - evaluator=CosineSimilarityEvaluator(0.9), + title="debug_cos_similarity", + path=args.artifact, + output_format=OutputFormat.SVG_GRAPH, + comparator=cos_comparator, ) adb.pull_debug_output( @@ -176,53 +210,135 @@ adb.pull_debug_output( ) ``` -#### Additional Options -The above example sets output formats as SVG and evaluation metrics using Cosine Similarity. Based on different needs, users can choose other output formats as shown in the `OutputFormat` class under [qnn_intermediate_debugger](./qnn_intermediate_debugger.py) +## Comparators + +Create comparators via the `create_comparator()` factory, which automatically injects the `edge_ep`. A couple sample comparators are provided under [qcom_numerical_comparator_sample.py](./qcom_numerical_comparator_sample.py): + ```python -class OutputFormat(IntEnum): - SVG_GRAPHS = 0 - CSV_FILES = 1 - DUMP_RAW = 2 +cos = qnn_intermediate_debugger.create_comparator(QcomCosineSimilarityComparator, threshold=0.9) +mse = qnn_intermediate_debugger.create_comparator(QcomMSEComparator, threshold=0.1) ``` -For evaluation metrics, if users would like to implement their own metrics, we have provided the option to implement [MetricEvaluatorBase](./metrics_evaluator.py). The following shows how to define custom metrics. +### Custom comparators + +Users can also define their own comparator by implementing a derived class from [QcomNumericalComparatorBase](./qcom_numerical_comparator_base.py). Inside the derived class, users will need to implement `metric_name()`, `is_valid_score()`, and `element_compare()`. The base class handles QNN-specific preprocessing (dequantization, layout conversion) internally — `preprocessing` cannot be overridden. ```python -class RootMeanSquaredErrorEvaluator(MetricEvaluatorBase): - def __init__(self, threshold=0.02): +from executorch.backends.qualcomm.debugger.qcom_numerical_comparator_base import ( + QcomNumericalComparatorBase, +) + +class MyComparator(QcomNumericalComparatorBase): + def __init__(self, edge_ep, threshold=0.5): + super().__init__(edge_ep) self.threshold = threshold def metric_name(self) -> str: - return "Root Mean Squared Error" - - def evaluate( - self, qnn_output: torch.Tensor, cpu_output: torch.Tensor - ) -> Tuple[Any, bool]: - mse = F.mse_loss(qnn_output, cpu_output) - rmse = torch.sqrt(mse) - valid = rmse < self.threshold - return rmse, valid - -qnn_intermediate_debugger.generate_results( - title="my_metric", - path=".", - output_format=OutputFormat.SVG_GRAPHS, - inspector=inspector, - evaluator=RootMeanSquaredErrorEvaluator(), -) + return "my_metric" + + def is_valid_score(self, score: float) -> bool: + return score >= self.threshold + + def element_compare(self, a, b) -> float: + # your comparison logic here + ... +``` + +## Output formats + +| Format | Enum | Output | +|--------|------|--------| +| SVG graph | `OutputFormat.SVG_GRAPH` | Color-coded computation graph (green=pass, red=fail) | +| CSV file | `OutputFormat.CSV_FILE` | Per-node tabular results | + +## Example Script + +An Inception_V3 demo script is provided at [qnn_intermediate_debugger_demo.py](../../../examples/qualcomm/util_scripts/qnn_intermediate_debugger_demo.py). + +Before running, ensure the dataset is downloaded. An example dataset can be retrieved [here](https://www.kaggle.com/datasets/ifigotin/imagenetmini-1000). + +```bash +python -m examples.qualcomm.util_scripts.qnn_intermediate_debugger_demo -b build-android -s $DEVICE_SERIAL -m $SOC_MODEL -d path/to/imagenet/val --dump_intermediate_outputs ``` -### Example Script -We have provided an inception_v3 demo script to help users better understand how to apply the debugger to their scripts. Please refer to [qnn_intermediate_debugger_demo.py](../../../examples/qualcomm/util_scripts/qnn_intermediate_debugger_demo.py) for the example script. +## Limitations +1. Only one execution per debug session — multiple executions may cause unknown behavior. +2. If you have decided to write your own runner (instead of `qnn_executor_runner`), follow the [tutorial](https://pytorch.org/executorch/stable/etdump.html) on how to implement etdump. +3. Does not support graphs with partitions (partial delegation). +4. Does not support LLM models. +5. Does not support graphs with multiple methods. + + +## ExecuTorch QNN HTP Heap Profiling + +Measures DSP memory usage when using context binary models on the HTP backend. + +### Introduction + +DSP heap profiling is available for `QnnContext_createFromBinary` use-cases. It captures total DSP heap usage at two checkpoints: -Before running the example script, please ensure that dataset is downloaded. Example dataset can be retrieved [here](https://www.kaggle.com/datasets/ifigotin/imagenetmini-1000). +- **Before the first context is created** (`before_context_created`) +- **After the last context is freed** (`after_context_freed`) + +The difference between the two values represents heap consumed during context execution. The value after freeing is typically equal to or greater than before creation. + +### Instructions + +#### Run the example test -To execute the model: ```bash -python examples/qualcomm/util_scripts/qnn_intermediate_debugger_demo.py -b build-android -m ${SOC_MODEL} --device ${SERIAL_NUM} --dataset ${PATH_TO_DATASET} --dump_intermediate_outputs +python backends/qualcomm/tests/test_qnn_delegate.py \ + TestQNNQuantizedUtils.test_qnn_backend_runtime_option_heap_profile \ + -b build-android -H ${HOST} -s ${SN} -m ${SOC_MODEL} +``` + +See [test_qnn_delegate.py](../tests/test_qnn_delegate.py) for the full test implementation. + +#### Setting + +```python +from executorch.backends.qualcomm.utils.utils import generate_htp_compiler_spec +from executorch.backends.qualcomm.utils.utils import generate_qnn_executorch_compiler_spec + +backend_options = generate_htp_compiler_spec( + use_multi_contexts=True, +) + +compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=self.chipset_table[TestQNN.soc_model], + backend_options=backend_options, + profile_level=2, +) + +# ... + +self.verify_output( + module, + sample_input, + exec_prog, + save_heap_result=True, +) ``` -### Limitation -1. The current debugger only supports performing one execution. Multiple executions may cause unknown behavior and are not recommended. -2. Please ignore this if you are using `qnn_executor_runner`. If you have decided to write your own runner, please follow the [tutorial](https://pytorch.org/executorch/stable/etdump.html) on how to implement etdump into your own runner. -3. The current debugger does not support graph with partitions. (WIP) -4. The current debugger does not support LLM models. (WIP) +#### Output file format + +The result is written to a text file (default: `htp_heap_usage.txt`) with two lines: + +``` +DSP:before_context_created (bytes), +DSP:after_context_freed (bytes), +``` + +#### Reference result + +Measured on SM8850. A difference of 0 means no additional heap is consumed during context binary execution. + +```console +First value (before_context_created): 928212 bytes +Second value (after_context_freed): 928212 bytes +difference: 0.00 bytes +``` + +### Limitations + +1. Only supported HTP backend on Android and QNX platforms. +2. By enabling this feature, initialization and cleanup time might be impacted. diff --git a/backends/qualcomm/debugger/TARGETS b/backends/qualcomm/debugger/TARGETS index 6a7732231fc..3389dfdfcd9 100644 --- a/backends/qualcomm/debugger/TARGETS +++ b/backends/qualcomm/debugger/TARGETS @@ -15,7 +15,8 @@ runtime.python_library( name = "qnn_intermediate_debugger", srcs = [ "format_outputs.py", - "metrics_evaluator.py", + "qcom_numerical_comparator_base.py", + "qcom_numerical_comparator_sample.py", "qnn_intermediate_debugger.py", ], deps = [ diff --git a/backends/qualcomm/debugger/format_outputs.py b/backends/qualcomm/debugger/format_outputs.py index 7388eef8223..d0dd165b186 100644 --- a/backends/qualcomm/debugger/format_outputs.py +++ b/backends/qualcomm/debugger/format_outputs.py @@ -5,21 +5,30 @@ # LICENSE file in the root directory of this source tree. import csv +import logging import os +import subprocess from typing import Any +import executorch.exir as exir +import pandas import pydot import torch +from executorch.backends.qualcomm.debugger.qcom_numerical_comparator_base import ( + QcomNumericalComparatorBase, +) from executorch.backends.qualcomm.utils.constants import ( QCOM_QUANT_ATTRS, QCOM_SCALE, QCOM_SCALES, - QCOM_TENSOR_NAME, QCOM_ZERO_POINT, QCOM_ZERO_POINTS, ) +from executorch.exir.debug_handle_utils import DEBUG_HANDLE_KEY -from .metrics_evaluator import MetricEvaluatorBase +FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" +logging.basicConfig(level=logging.INFO, format=FORMAT) +logging.getLogger().setLevel(logging.INFO) # Copied from site-packages/torch/fx/passes/graph_drawer.py @@ -39,63 +48,37 @@ def typename(target: Any) -> str: return ret.replace("{", r"\{").replace("}", r"\}") -def retrieve_node_info(evaluator, node, node_tensor_map): - - node_info = {} - node_info["name"] = node.name - node_info["op_code"] = node.op - node_info["target"] = typename(node.target) - node_info["num_users"] = len(node.users) - - if "val" in node.meta: - if isinstance(node.meta["val"], torch.Tensor): - node_info["pytorch_layout"] = node.meta["val"].shape - elif isinstance(node.meta["val"], (list, tuple)): - shape_list = [] - for i in range(len(node.meta["val"])): - shape_list.append(node.meta["val"][i].shape) - node_info["pytorch_layout"] = shape_list - +def get_scale_zero_point(node: torch.fx.node.Node): + scale_zero_point = {"scale(s)": None, "zero_point(s)": None} if quant_attrs := node.meta.get(QCOM_QUANT_ATTRS): - node_info["scale(s)"] = ( + scale_zero_point["scale(s)"] = ( quant_attrs.get(QCOM_SCALES) if QCOM_SCALES in quant_attrs else quant_attrs.get(QCOM_SCALE) ) - node_info["zero_point(s)"] = ( + scale_zero_point["zero_point(s)"] = ( quant_attrs.get(QCOM_ZERO_POINTS) if QCOM_ZERO_POINTS in quant_attrs else quant_attrs.get(QCOM_ZERO_POINT) ) + return scale_zero_point - if node.name in node_tensor_map: - qnn_output, cpu_output, meta = node_tensor_map[node.name] - node_info[QCOM_TENSOR_NAME] = meta.get(QCOM_TENSOR_NAME) - node_info[evaluator.metric_name()], node_info["is_valid_score"] = ( - evaluator.evaluate(qnn_output, cpu_output) - ) - # The values in meta are directly retrieved from the node during the forward hook, which means the values should be the same for meta and node.meta. - # Storing these data during the forward hook helps us compare QNN tensors with CPU tensors without traversing the graph. - # We only check "scale" and not "scales" since the forward hook only stores the node's output, which should always be per tensor. - if QCOM_QUANT_ATTRS in node.meta: - assert ( - node_info["scale(s)"] == node.meta[QCOM_QUANT_ATTRS][QCOM_SCALE] - ), "node meta scale should be same as scale retrieve during forward hook" - assert ( - node_info["zero_point(s)"] - == node.meta[QCOM_QUANT_ATTRS][QCOM_ZERO_POINT] - ), "node meta zero_point should be same as zero_point retrieve during forward hook" +def get_pytorch_layout_info(node: torch.fx.node.Node): + val = node.meta.get("val") + if val is None: + return None + if isinstance(val, torch.Tensor): + return val.shape + return [v.shape for v in val if isinstance(v, torch.Tensor)] - return node_info - -def export_svg( +def export_svg( # noqa: C901 title: str, path: str, - evaluator: MetricEvaluatorBase, - edge_module: torch.fx.GraphModule, - node_tensor_map: dict, + edge_ep: exir.ExirExportedProgram, + numeric_results: pandas.core.frame.DataFrame, + comparator: QcomNumericalComparatorBase, ): def get_node_style(is_valid_score: bool): template = { @@ -117,37 +100,46 @@ def get_node_style(is_valid_score: bool): node_map = {} # Create node - for node in edge_module.graph.nodes: + for node in edge_ep.graph_module.graph.nodes: # These are just nodes before fold_quant and still there if len(node.users) == 0 and node.op == "placeholder": continue - node_info = retrieve_node_info( - evaluator=evaluator, node=node, node_tensor_map=node_tensor_map - ) + + pytorch_layout = get_pytorch_layout_info(node) + scale_zero_point = get_scale_zero_point(node) + scale = scale_zero_point["scale(s)"] + zero_point = scale_zero_point["zero_point(s)"] node_label = "{" - node_label += f"name=%{node_info.get('name')}" + r"\n" - node_label += f"|op_code={node_info.get('op_code')}" + r"\n" - node_label += f"|qnn_tensor_name={node_info.get('qnn_tensor_name')}" + r"\n" - node_label += f"|target={node_info.get('target')}" + r"\n" - node_label += f"|num_users={node_info.get('num_users')}" + r"\n" - node_label += f"|pytorch_layout={node_info.get('pytorch_layout')}" + r"\n" - node_label += f"|scale(s)={node_info.get('scale(s)')}" + r"\n" - node_label += f"|zero_point(s)={node_info.get('zero_point(s)')}" + r"\n" - node_label += ( - f"|{evaluator.metric_name()}={node_info.get(evaluator.metric_name())}" - + r"\n" - ) - node_label += f"|is_valid_score={node_info.get('is_valid_score')}" + r"\n" + node_label += f"name=%{node.name}" + r"\n" + node_label += f"|op_code={node.op}" + r"\n" + node_label += f"|target={typename(node.target)}" + r"\n" + node_label += f"|num_users={len(node.users)}" + r"\n" + node_label += f"|pytorch_layout={pytorch_layout}" + r"\n" + node_label += f"|scale(s)={scale}" + r"\n" + node_label += f"|zero_point(s)={zero_point}" + r"\n" + + is_valid_score = None + if debug_handle := node.meta.get(DEBUG_HANDLE_KEY, None): + node_label += f"|debug_handle={debug_handle}" + r"\n" + debug_handle = (debug_handle,) + if debug_handle in numeric_results.index: + score = numeric_results.loc[[debug_handle], "gap"].iat[0][0] + assert isinstance( + score, float + ), f"Expecting QcomNumericalComparatorBase element_compare to return float, but get {type(score)}." + node_label += f"|{comparator.metric_name()}={score:.3f}" + r"\n" + is_valid_score = comparator.is_valid_score(score) + node_label += f"|is_valid_score={is_valid_score}" + r"\n" node_label += "}" - template = get_node_style(node_info.get("is_valid_score")) + template = get_node_style(is_valid_score) pydot_node = pydot.Node(node.name, label=node_label, **template) node_map[node.name] = pydot_node pydot_graph.add_node(pydot_node) # Create edge - for node in edge_module.graph.nodes: + for node in edge_ep.graph_module.graph.nodes: if len(node.users) == 0 and node.op == "placeholder": continue cur_pydot_node = node_map[node.name] @@ -157,28 +149,68 @@ def get_node_style(is_valid_score: bool): pydot_graph.add_edge( pydot.Edge(cur_pydot_node, user_pydot_node, dir="forward") ) + dot_file_path = os.path.join(path, f"{title}.dot") + pydot_graph.write_raw(dot_file_path) + logging.info(f"Intermediate debugger dot graph saved at: {dot_file_path}") svg_file_path = os.path.join(path, f"{title}.svg") - pydot_graph.write_svg(svg_file_path) - print(f"Intermediate debugger graph saved at: {svg_file_path}") + try: + subprocess.run( + ["dot", "-Tsvg", dot_file_path, "-o", svg_file_path], + timeout=5, + check=True, + ) + logging.info(f"Intermediate debugger SVG graph saved at: {svg_file_path}.") + except subprocess.TimeoutExpired: + logging.warning( + f"SVG generation timed out after 5s, skipping. " + f"Only saving the dot file: {dot_file_path}." + ) + except (subprocess.CalledProcessError, FileNotFoundError) as e: + logging.warning(f"SVG generation failed ({e}), skipping.") def export_csv( title: str, path: str, - evaluator: MetricEvaluatorBase, - edge_module: torch.fx.GraphModule, - node_tensor_map: dict, + edge_ep: exir.ExirExportedProgram, + numeric_results: pandas.core.frame.DataFrame, + comparator: QcomNumericalComparatorBase, ): node_info_list = [] - for node in edge_module.graph.nodes: + for node in edge_ep.graph_module.graph.nodes: # These are just nodes before fold_quant and still there if len(node.users) == 0 and node.op == "placeholder": continue - node_info = retrieve_node_info( - evaluator=evaluator, node=node, node_tensor_map=node_tensor_map + + pytorch_layout = get_pytorch_layout_info(node) + scale_zero_point = get_scale_zero_point(node) + scale = scale_zero_point["scale(s)"] + zero_point = scale_zero_point["zero_point(s)"] + score = None + is_valid_score = None + if debug_handle := node.meta.get(DEBUG_HANDLE_KEY, None): + if (debug_handle,) in numeric_results.index: + score = numeric_results.loc[[(debug_handle,)], "gap"].iat[0][0] + assert isinstance( + score, float + ), f"Expecting QcomNumericalComparatorBase element_compare to return float, but get {type(score)}." + is_valid_score = comparator.is_valid_score(score) + + node_info_list.append( + { + "name": node.name, + "op_code": node.op, + "target": typename(node.target), + "num_users": len(node.users), + "pytorch_layout": pytorch_layout, + "scale(s)": scale, + "zero_point(s)": zero_point, + "debug_handle": debug_handle, + comparator.metric_name(): score, + "is_valid_score": is_valid_score, + } ) - node_info_list.append(node_info) # Writing to a CSV file csv_file_path = os.path.join(path, f"{title}.csv") @@ -186,13 +218,13 @@ def export_csv( fieldnames = [ "name", "op_code", - "qnn_tensor_name", "target", "num_users", "pytorch_layout", "scale(s)", "zero_point(s)", - f"{evaluator.metric_name()}", + "debug_handle", + comparator.metric_name(), "is_valid_score", ] writer = csv.DictWriter(csv_file, fieldnames=fieldnames) @@ -201,23 +233,3 @@ def export_csv( writer.writerows(node_info_list) print(f"Intermediate debugger csv saved at: {csv_file_path}") - - -def export_raw( - path: str, - edge_module: torch.fx.GraphModule, - node_tensor_map: dict, -): - for node in edge_module.graph.nodes: - # These are just unused nodes before fold_quant and still there - if len(node.users) == 0 and node.op == "placeholder": - continue - if paired_event := node_tensor_map.get(node.name): - qnn_output, cpu_output, meta = paired_event - qnn_tensor_name = meta[QCOM_TENSOR_NAME] - qnn_output_path = os.path.join(path, qnn_tensor_name + "_qnn.raw") - cpu_output_path = os.path.join(path, qnn_tensor_name + "_cpu.raw") - qnn_output.numpy().tofile(qnn_output_path) - cpu_output.numpy().tofile(cpu_output_path) - - print(f"Intermediate debugger raw files saved at: {path}") diff --git a/backends/qualcomm/debugger/metrics_evaluator.py b/backends/qualcomm/debugger/metrics_evaluator.py deleted file mode 100644 index 55c8b92b034..00000000000 --- a/backends/qualcomm/debugger/metrics_evaluator.py +++ /dev/null @@ -1,90 +0,0 @@ -# Copyright (c) Qualcomm Innovation Center, Inc. -# All rights reserved -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from abc import ABC, abstractmethod -from typing import Any, Tuple - -import torch - - -class MetricEvaluatorBase(ABC): - @abstractmethod - def metric_name(self) -> str: - """ - A name for this metric evaluation - - Returns: - str: name of the metric evaluation - """ - ... - - @abstractmethod - def evaluate( - self, qnn_output: torch.Tensor, cpu_output: torch.Tensor, **kwargs - ) -> Tuple[Any, bool]: - """ - This abstract method should accept both QNN and CPU outputs for a single layer. - Define your own logic to compare the results. - - Args: - qnn_output (torch.Tensor): QNN intermediate output - cpu_output (torch.Tensor): CPU intermediate output - - Returns: - Tuple[Any, bool]: Return 2 elements: - 1) Score or anything that you would like to be printed under metrics category for svg graph or csv file. - 2) A boolean that indicates whether the evaluation result is acceptable or not. - """ - ... - - -class AtolEvaluator(MetricEvaluatorBase): - def __init__(self, threshold=1e-1): - self.threshold = threshold - - def metric_name(self) -> str: - return "Atol Similarity" - - def evaluate( - self, qnn_output: torch.Tensor, cpu_output: torch.Tensor - ) -> Tuple[Any, bool]: - avg_atol = torch.mean(torch.abs(qnn_output - cpu_output)) - valid = avg_atol < self.threshold - formatted_score = f"{avg_atol:.3f}" - return formatted_score, valid - - -class CosineSimilarityEvaluator(MetricEvaluatorBase): - def __init__(self, threshold=0.9): - self.threshold = threshold - - def metric_name(self) -> str: - return "Cosine Similarity" - - def evaluate( - self, qnn_output: torch.Tensor, cpu_output: torch.Tensor - ) -> Tuple[Any, bool]: - score = torch.nn.functional.cosine_similarity( - qnn_output.flatten(), cpu_output.flatten(), dim=0 - ).item() - valid = score > self.threshold - formatted_score = f"{score:.3f}" - return formatted_score, valid - - -class MeanSquaredErrorEvaluator(MetricEvaluatorBase): - def __init__(self, threshold=0.01): - self.threshold = threshold - - def metric_name(self) -> str: - return "Mean Squared Error" - - def evaluate( - self, qnn_output: torch.Tensor, cpu_output: torch.Tensor - ) -> Tuple[Any, bool]: - mse = torch.mean((qnn_output - cpu_output) ** 2) - valid = mse < self.threshold - return mse, valid diff --git a/backends/qualcomm/debugger/qcom_numerical_comparator_base.py b/backends/qualcomm/debugger/qcom_numerical_comparator_base.py new file mode 100644 index 00000000000..0068d28bfac --- /dev/null +++ b/backends/qualcomm/debugger/qcom_numerical_comparator_base.py @@ -0,0 +1,194 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import operator +from abc import abstractmethod +from dataclasses import dataclass +from typing import Dict, final, Optional, Tuple + +import executorch.exir as exir +import torch +from executorch.backends.qualcomm._passes.layout_transform import LayoutTransform +from executorch.backends.qualcomm.utils.constants import ( + QCOM_AXIS_ORDER, + QCOM_QUANT_ATTRS, + QCOM_SCALE, + QCOM_ZERO_POINT, +) +from executorch.devtools.inspector.numerical_comparator import ( + IntermediateOutputMapping, + NumericalComparatorBase, +) +from executorch.exir.sym_util import eval_shape + +logger = logging.getLogger(__name__) + + +@dataclass +class NodeMetaInfo: + node_name: str + scale: Optional[float] = None + zero_point: Optional[int] = None + axis_order: Optional[Tuple[int, ...]] = None + + +class QcomNumericalComparatorBase(NumericalComparatorBase): + """Base class for Qualcomm numerical comparators. + + This class locks down the `preprocessing` method to handle QNN-specific + tensor transformations (dequantization, layout conversion) internally. + Community users subclassing this base only need to implement `element_compare`. + + Attempting to override `preprocessing` in a subclass will raise TypeError + at class definition time. + """ + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + if "preprocessing" in cls.__dict__: + raise TypeError( + f"{cls.__name__} cannot override 'preprocessing'. " + "Qualcomm handles preprocessing (dequantization, layout conversion) internally." + ) + + def __init__(self, edge_ep: exir.ExportedProgram) -> None: + super().__init__() + self.edge_ep = edge_ep + + @abstractmethod + def metric_name(self) -> str: + """ + A name for this metric evaluation. + + Returns: + str: name of the metric evaluation. + """ + ... + + @abstractmethod + def is_valid_score(self, score: float) -> bool: + """ + Determine whether a comparison score is within an acceptable range. + + Args: + score: the comparison score to validate. + + Returns: + bool: True if the score is acceptable, False otherwise. + """ + ... + + @final + def preprocessing( # noqa: C901 + self, mapping: IntermediateOutputMapping + ) -> IntermediateOutputMapping: + + def _preprocess_tensor( + qnn_tensor: torch.Tensor, meta: NodeMetaInfo, golden_tensor: torch.Tensor + ) -> torch.Tensor: + if meta.scale is not None: + # Dequantize + qnn_tensor = ( + qnn_tensor.to(torch.float32) + .sub(meta.zero_point) + .mul(meta.scale) + .contiguous() + ) + if meta.axis_order: + # QNN to Pytorch layout + axis_order = LayoutTransform.get_axis_order( + eval_shape(qnn_tensor.shape), reverse=True + ) + qnn_tensor = qnn_tensor.permute(axis_order) + + assert ( + golden_tensor.shape == qnn_tensor.shape + ), f"{meta.node_name}'s golden and QNN tensor has different shape. Golden Tensor Shape: {golden_tensor.shape}. QNN Tensor Shape: {qnn_tensor.shape}." + + return qnn_tensor + + def _build_debug_handle_to_meta() -> ( + Dict[Tuple[int, ...], Dict[int, NodeMetaInfo]] + ): + debug_handle_to_meta: Dict[Tuple[int, ...], Dict[int, NodeMetaInfo]] = {} + for node in self.edge_ep.graph_module.graph.nodes: + if node.op != "call_function": + continue + + if (debug_handle := node.meta.get("debug_handle")) is None: + continue + else: + debug_handle = (debug_handle,) + + quant_attrs = node.meta.get(QCOM_QUANT_ATTRS, {}) + node_meta_info = NodeMetaInfo( + node_name=node.name, + scale=quant_attrs.get(QCOM_SCALE, None), + zero_point=quant_attrs.get(QCOM_ZERO_POINT, None), + axis_order=node.meta.get(QCOM_AXIS_ORDER, None), + ) + + if any(user.target == operator.getitem for user in node.users): + # Assume if a node user is getitem, all users are getitem + assert all( + user.target == operator.getitem for user in node.users + ), "[QNN Delegate Debugger]: Expect all users to be getitem node" + continue + + # Multi-output op's getitem node shares the same debug handle. + if node.target == operator.getitem: + output_idx = node.args[1] + debug_handle_to_meta.setdefault(debug_handle, {})[ + output_idx + ] = node_meta_info + else: + assert ( + debug_handle not in debug_handle_to_meta + ), f"[QNN Delegate Debugger]: Duplicate handle_id {debug_handle} found when visiting {node.name}." + debug_handle_to_meta[debug_handle] = {0: node_meta_info} + + return debug_handle_to_meta + + debug_handle_to_meta = _build_debug_handle_to_meta() + processed_mapping: IntermediateOutputMapping = {} + for (golden_handle, golden_output), (qnn_handle, qnn_output) in mapping.items(): + assert ( + golden_handle == qnn_handle + ), f"Expecting the handle to match, aot handle: {golden_handle}, qnn_handle: {qnn_handle}." + if node_meta_dict := debug_handle_to_meta.get(qnn_handle, None): + if isinstance(qnn_output, tuple): + assert len(qnn_output) <= len( + node_meta_dict + ), f"node_meta has {len(node_meta_dict)} entries but qnn_output has {len(qnn_output)} elements." + if len(node_meta_dict) != len(qnn_output): + logging.warning( + f"Number of QNN output {len(qnn_output)} mismatched with number of output for edge module {len(node_meta_dict)}. This is possibly due to multi-outputs and QNN does not use all outputs. Please verify the following meta from edge module and ensure this is desired: {node_meta_dict}." + ) + + processed = [] + for idx, q_tensor in enumerate(qnn_output): + processed.append( + _preprocess_tensor( + qnn_tensor=q_tensor, + meta=node_meta_dict[idx], + golden_tensor=golden_output[idx], + ) + ) + + qnn_output = tuple(processed) + else: + assert ( + len(node_meta_dict) == 1 and 0 in node_meta_dict + ), f"Single output expected node_meta_dict with key 0, got keys {list(node_meta_dict.keys())}" + qnn_output = _preprocess_tensor( + qnn_tensor=qnn_output, + meta=node_meta_dict[0], + golden_tensor=golden_output, + ) + + processed_mapping[(golden_handle, golden_output)] = (qnn_handle, qnn_output) + return processed_mapping diff --git a/backends/qualcomm/debugger/qcom_numerical_comparator_sample.py b/backends/qualcomm/debugger/qcom_numerical_comparator_sample.py new file mode 100644 index 00000000000..43783a64420 --- /dev/null +++ b/backends/qualcomm/debugger/qcom_numerical_comparator_sample.py @@ -0,0 +1,57 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Any + +import executorch.exir as exir +import torch +from executorch.backends.qualcomm.debugger.qcom_numerical_comparator_base import ( + QcomNumericalComparatorBase, +) + + +""" +This file provides some examples on how to implement a QcomNumericalComparator +""" + + +class QcomMSEComparator(QcomNumericalComparatorBase): + """Mean Squared Error comparator for Qualcomm intermediate outputs.""" + + def __init__(self, edge_ep: exir.ExportedProgram, threshold: float = 1e-3) -> None: + super().__init__(edge_ep) + self.threshold = threshold + + def metric_name(self) -> str: + return "mse" + + def is_valid_score(self, score: float) -> bool: + return score <= self.threshold + + def element_compare(self, a: Any, b: Any) -> float: + if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor): + return torch.mean(torch.square(a.float() - b.float())).item() + return float((a - b) ** 2) + + +class QcomCosineSimilarityComparator(QcomNumericalComparatorBase): + """Cosine Similarity comparator for Qualcomm intermediate outputs.""" + + def __init__(self, edge_ep: exir.ExportedProgram, threshold: float = 0.95) -> None: + super().__init__(edge_ep) + self.threshold = threshold + + def metric_name(self) -> str: + return "cosine_similarity" + + def is_valid_score(self, score: float) -> bool: + return score >= self.threshold + + def element_compare(self, a: Any, b: Any) -> float: + score = torch.nn.functional.cosine_similarity( + a.to(torch.float32).flatten(), b.to(torch.float32).flatten(), dim=0 + ).item() + return score diff --git a/backends/qualcomm/debugger/qnn_intermediate_debugger.py b/backends/qualcomm/debugger/qnn_intermediate_debugger.py index 904dd4f6ccb..7c7609cd7f5 100644 --- a/backends/qualcomm/debugger/qnn_intermediate_debugger.py +++ b/backends/qualcomm/debugger/qnn_intermediate_debugger.py @@ -4,136 +4,94 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import copy -import operator import os import warnings from enum import IntEnum -import torch - -from executorch.backends.qualcomm._passes.layout_transform import LayoutTransform -from executorch.backends.qualcomm.utils.constants import ( - QCOM_AXIS_ORDER, - QCOM_QUANT_ATTRS, - QCOM_SCALE, - QCOM_TENSOR_NAME, - QCOM_ZERO_POINT, +import executorch.exir as exir +from executorch.backends.qualcomm.debugger.qcom_numerical_comparator_base import ( + QcomNumericalComparatorBase, ) from executorch.devtools import Inspector -from executorch.exir.sym_util import eval_shape -from .format_outputs import export_csv, export_raw, export_svg -from .metrics_evaluator import MetricEvaluatorBase +from .format_outputs import export_csv, export_svg class OutputFormat(IntEnum): - SVG_GRAPHS = 0 - CSV_FILES = 1 - DUMP_RAW = 2 - - -class IntermediateModule(torch.nn.Module): - """ - This class serves as an intermediate point and is inserted right after the call_function node. - It also saves some metadata such as scale, offset, etc. - Since we just want to check the intermediate output, we will directly return the value during the forward call. - """ - - def __init__( - self, - module_name: str, - qnn_tensor_name: str, - node_name: str, - scale: float, - zero_point: int, - revert_order: bool = None, - ): - super().__init__() - self.module_name = module_name - self.qnn_tensor_name = qnn_tensor_name - self.node_name = node_name - self.scale = scale - self.zero_point = zero_point - self.revert_order = revert_order - - def forward(self, x): - return x + SVG_GRAPH = 0 + CSV_FILE = 1 class QNNIntermediateDebugger: """This is a debugger tool capable of retrieving intermediate results for CPU edge EP. - We can further compare these with QNN's intermediate output to identify any QNN accuracy issues. + We can further compare these with QNN's intermediate output to identify any accuracy issues. """ - def __init__(self): - self.intermediate_outputs = {} - - def set_edge_module(self, edge_module: torch.fx.graph_module.GraphModule): - self.orig_edge = copy.deepcopy(edge_module) - self.intermediate_output_module = self._insert_intermediate_module( - copy.deepcopy(edge_module) + def __init__(self, sample_input): + self.sample_input = sample_input + self.edge_ep = None + self.etrecord_file_path = None + self.inspector = None + # Support single to edge after transform forward graph for now. + self.reference_graph_name = "edge_after_transform/forward" + + def set_edge_ep(self, edge_ep: exir.ExirExportedProgram): + self.edge_ep = edge_ep + + def set_etrecord_file_path(self, etrecord_file_path: str): + self.etrecord_file_path = etrecord_file_path + + def setup_inspector(self, etdump_path: str, debug_buffer_path: str): + self.inspector = Inspector( + etdump_path=etdump_path, + debug_buffer_path=debug_buffer_path, + etrecord=self.etrecord_file_path, + reference_graph_name=self.reference_graph_name, ) + def create_comparator( + self, comparator_cls: type[QcomNumericalComparatorBase], **kwargs + ) -> QcomNumericalComparatorBase: + # No need to pass edge_ep — the factory injects it automatically. + # Just pass the comparator class and any comparator-specific args: + # comparator = debugger.create_comparator(QcomMSEComparator, threshold=1e-4) + assert ( + self.edge_ep is not None + ), "edge_ep must be set before creating a comparator." + return comparator_cls(edge_ep=self.edge_ep, **kwargs) + def generate_results( self, title: str, path: str, output_format: OutputFormat, - inspector: Inspector, - evaluator: MetricEvaluatorBase = None, - keep_qnn_layout: bool = False, + comparator: QcomNumericalComparatorBase, ): assert isinstance( output_format, OutputFormat ), "output_format passed in is not an instance of OutputFormat" os.makedirs(path, exist_ok=True) - if keep_qnn_layout: - warnings.warn( - "[QNN Delegate Debugger]: keep_qnn_layout is not recommended for general use case. " - "QNN and CPU has different dtype(FP V.S. Quantized) and data formats(NCHW V.S. NHWC) in a lot of cases.", - stacklevel=1, - ) - # Due to users can switch between keep_qnn_layout between generate_results, rematch this every time. - # Make this a class variable if repeat matching is taking too long and handle keep_qnn_layout. - node_tensor_map = self._match_tensors( - inspector=inspector, - keep_qnn_layout=keep_qnn_layout, + numeric_results = self.inspector.calculate_numeric_gap( + distance=comparator, reference_graph=self.reference_graph_name ) + numeric_results = numeric_results.set_index("runtime_debug_handle") - if output_format == OutputFormat.SVG_GRAPHS: - assert evaluator is not None, "Please provide an evaluator." + if output_format == OutputFormat.SVG_GRAPH: export_svg( title=title, path=path, - evaluator=evaluator, - edge_module=self.orig_edge, - node_tensor_map=node_tensor_map, + edge_ep=self.edge_ep, + numeric_results=numeric_results, + comparator=comparator, ) - elif output_format == OutputFormat.CSV_FILES: - assert evaluator is not None, "Please provide an evaluator." + elif output_format == OutputFormat.CSV_FILE: export_csv( title=title, path=path, - evaluator=evaluator, - edge_module=self.orig_edge, - node_tensor_map=node_tensor_map, - ) - elif output_format == OutputFormat.DUMP_RAW: - warnings.warn( - f"[QNN Delegate Debugger]: Param 'title' will be ignored, all raw files will be stored under: {path}", - stacklevel=1, - ) - if evaluator: - warnings.warn( - "[QNN Delegate Debugger]: Param 'evaluator' will be ignored as DUMP_RAW will only dump tensors to raw files but won't perform comparison.", - stacklevel=1, - ) - export_raw( - path=path, - edge_module=self.intermediate_output_module, - node_tensor_map=node_tensor_map, + edge_ep=self.edge_ep, + numeric_results=numeric_results, + comparator=comparator, ) else: warnings.warn( @@ -141,188 +99,3 @@ def generate_results( stacklevel=1, ) return - - def _insert_intermediate_module( # noqa: C901 - self, edge_module: torch.fx.graph_module.GraphModule - ): - """ - This feature is for intermediate tensor dump on the host CPU. - After we get an edge GraphModule, we insert submodule between each call_function node, - and we register forward hooks to store the intermediate results. - We have to use the edge GraphModule because this is the graph closest to what QNN is executing - while still being a valid graph to ExecuTorch. - - Args: - edge_module (exir.ExirExportedProgram): A deep copy of edge ir graph module. - We need to deep copy so we don't mess up the original edge_ep. - Returns: - exir.ExirExportedProgram: A deep copy of edge graph_module with intermediate modules inserted. - """ - - def hook_fn(module, input, output): - meta = {} - meta[QCOM_TENSOR_NAME] = module.qnn_tensor_name - meta["node_name"] = module.node_name - meta[QCOM_SCALE] = module.scale - meta[QCOM_ZERO_POINT] = module.zero_point - meta["revert_order"] = module.revert_order - meta["output"] = output # CPU output - - assert ( - module.qnn_tensor_name not in self.intermediate_outputs - ), f"{module.qnn_tensor_name} checked already, check if this is a potential error" - self.intermediate_outputs[module.qnn_tensor_name] = meta - - graph = edge_module.graph - module_count = 0 - for node in graph.nodes: - if node.op == "call_function": - module_name = f"intermediate_module_{module_count}" - module_count += 1 - with graph.inserting_after(node): - scale = None - zero_point = None - if QCOM_QUANT_ATTRS in node.meta: - scale = node.meta[QCOM_QUANT_ATTRS][QCOM_SCALE] - zero_point = node.meta[QCOM_QUANT_ATTRS][QCOM_ZERO_POINT] - - revert_order = QCOM_AXIS_ORDER in node.meta - - if node.target == operator.getitem: - index = node.args[1] - # Ex: topk -> intermediate_module -> get_item - src_node = node.args[0].args[0] - qnn_tensor_name = src_node.meta[QCOM_TENSOR_NAME][index] - elif any(user.target == operator.getitem for user in node.users): - # For cases like topK, qnn_tensor_name is stored in get_item instead of source_node itself. - assert all( - user.target == operator.getitem for user in node.users - ), "Expect all users to be get_item node" - qnn_tensor_name = node.name - elif QCOM_TENSOR_NAME in node.meta: - assert ( - len(node.meta[QCOM_TENSOR_NAME]) == 1 - ), "Expecting a single qnn_tensor name but get more than 1." - qnn_tensor_name = node.meta[QCOM_TENSOR_NAME][0] - else: - # Unused - qnn_tensor_name = node.name - - obs = IntermediateModule( - module_name=module_name, - qnn_tensor_name=qnn_tensor_name, - node_name=node.name, - scale=scale, - zero_point=zero_point, - revert_order=revert_order, - ) - setattr( - edge_module, - module_name, - obs, - ) - new_obs = graph.create_node("call_module", module_name, (node,), {}) - orig_users = list(node.users.keys()) - for user_node in orig_users: - if user_node is new_obs: - continue - user_node.replace_input_with(node, new_obs) - - # Register hooks for all intermediate layers - for ( - _, - layer, - ) in edge_module.named_modules(): - if isinstance(layer, IntermediateModule): - layer.register_forward_hook(hook_fn) - - graph.eliminate_dead_code() - edge_module.recompile() - - return edge_module - - def _process_qnn_output(self, qnn_output: torch.tensor, meta: dict) -> torch.tensor: - """ - QNN intermediate results are all quantized. - We need to dequantize them to match CPU float values. - Additionally, we need to revert the layout format for layout-sensitive nodes. - - Args: - qnn_output (torch.tensor): QNN intermediate output from inspector event - meta (dict): The meta for this tensor/node that is stored during insert_intermediate_module(). - - Returns: - torch.tensor: Processed tensor that should have same dtype and shape as CPU tensors. - """ - qnn_output = qnn_output.to(torch.float32) - if meta[QCOM_SCALE] is not None: - scale = meta[QCOM_SCALE] - zero_point = meta[QCOM_ZERO_POINT] - qnn_output = ( - qnn_output.sub(zero_point).mul(scale).to(torch.float32).contiguous() - ) - if meta["revert_order"]: - axis_order = LayoutTransform.get_axis_order( - eval_shape(qnn_output.shape), reverse=True - ) - qnn_output = qnn_output.permute(axis_order) - return qnn_output - - def _match_tensors(self, inspector: Inspector, keep_qnn_layout: bool = False): - """ - Map QNN tensors back to CPU tensors. - Create a map using the node name as the key and (preprocessed/postprocessed QNN tensor, CPU tensor, meta) as the value. - We need meta because it holds values such as scale, offset, layout sensitivity, etc. - - Args: - inspector (Inspector): Inspector that parse QNN runtime intermediate outputs - keep_qnn_layout (bool): If true, store QNN outputs in NHWC format. Not recommended for general users. - - Returns: - A dict storing {node_name : tuple(qnn_output, cpu_output, meta_info)} - Meta_info is the info stored during forward hook_fn. - """ - - # node_tensor_map {key: tuple(qnn_output, cpu_output, meta_info)} - node_tensor_map = {} - # OPs that only exists in QNN but not CPU Golden - unmatched_qnn_tensors = [] - # E.g.: DELEGATE_CALL (This is the model input data), 'Method::execute' - ignored_events = [] - # Collected with forward hook - intermediate_outputs = self.intermediate_outputs - for event_block in inspector.event_blocks: - if event_block.name == "Execute": - for event in event_block.events: - # If user enables profiling and dump intermediate outputs the same time, we need to skip the profiling event - if event.perf_data is not None and event.is_delegated_op: - continue - if meta := intermediate_outputs.get(event.name): - node_name = meta["node_name"] - cpu_output = meta["output"] - qnn_output = ( - event.debug_data[0] - if keep_qnn_layout - else self._process_qnn_output(event.debug_data[0], meta) - ) - node_tensor_map[node_name] = ( - qnn_output, - cpu_output, - meta, - ) - - else: - ( - unmatched_qnn_tensors.append(event.name) - if event.is_delegated_op - else ignored_events.append(event.name) - ) - - warnings.warn( - f"The following events are ignored: {ignored_events}", stacklevel=1 - ) - warnings.warn( - f"The following QNN OPs are missing CPU reference. OPs added during qnn_preprocess will not have CPU reference. Please ensure the operations below are created during qnn_preprocess. {unmatched_qnn_tensors}", - stacklevel=1, - ) - return node_tensor_map diff --git a/backends/qualcomm/export_utils.py b/backends/qualcomm/export_utils.py index 313573e523a..1bca168ad3f 100644 --- a/backends/qualcomm/export_utils.py +++ b/backends/qualcomm/export_utils.py @@ -53,7 +53,6 @@ is_qnn_sdk_version_less_than, to_edge_transform_and_lower_to_qnn, ) -from executorch.exir.backend.utils import get_delegates from executorch.exir.capture._config import ExecutorchBackendConfig from executorch.exir.passes.memory_planning_pass import MemoryPlanningPass from torchao.quantization.pt2e import MovingAverageMinMaxObserver @@ -119,7 +118,7 @@ class QnnConfig: ci: Optional[bool] = False seed: Optional[int] = None htp_performance_mode: QnnExecuTorchHtpPerformanceMode = ( - QnnExecuTorchHtpPerformanceMode.kHtpBurst, + QnnExecuTorchHtpPerformanceMode.kHtpBurst ) def __post_init__(self): @@ -494,6 +493,11 @@ def pull_debug_output(self, etdump_path, debug_ouput_path, callback=None): if callback: callback() + def pull_heap_output(self, src_file_path, dst_folder, callback=None): + self._adb(["pull", src_file_path, dst_folder]) + if callback: + callback() + def build_executorch_binary( model: torch.nn.Module, # noqa: B006 @@ -600,6 +604,7 @@ def build_executorch_binary( dep_table=passes_dependency, skip_node_id_set=qnn_config.skip_delegate_node_ids, skip_node_op_set=qnn_config.skip_delegate_node_ops, + generate_etrecord=qnn_intermediate_debugger is not None, ) else: edge_prog_mgr = to_edge_transform_and_lower_to_qnn( @@ -610,21 +615,9 @@ def build_executorch_binary( passes_job=passes_job, skip_node_id_set=qnn_config.skip_delegate_node_ids, skip_node_op_set=qnn_config.skip_delegate_node_ops, + generate_etrecord=qnn_intermediate_debugger is not None, ) - if qnn_intermediate_debugger: - lowered_module_nodes = get_delegates(edge_prog_mgr.exported_program().graph) - assert ( - len(lowered_module_nodes) == 1 - ), "Graph with partitions are currently unsupported." - - lowered_module_node = lowered_module_nodes[0] - lower_module = getattr( - edge_prog_mgr.exported_program().graph_module, lowered_module_node.name - ) - edge_module = lower_module.original_module.module() - qnn_intermediate_debugger.set_edge_module(edge_module=edge_module) - allocate_io = not (qnn_config.shared_buffer or qnn_config.direct_build_folder) executorch_config = ExecutorchBackendConfig( # For shared buffer, user must pass the memory address @@ -642,6 +635,16 @@ def build_executorch_binary( with open(pte_name, "wb") as file: exec_prog_mgr.write_to_file(file) + if qnn_intermediate_debugger: + etrecord = exec_prog_mgr.get_etrecord() + etrecord.update_representative_inputs(qnn_intermediate_debugger.sample_input) + edge_ep = etrecord.graph_map[qnn_intermediate_debugger.reference_graph_name] + # Use this edge_ep since edge_ep after etrecord serialize/deserialize will lose quant_attrs info. + qnn_intermediate_debugger.set_edge_ep(edge_ep=edge_ep) + etrecord_file_path = f"{os.path.dirname(pte_name)}/debug.etrecord" + qnn_intermediate_debugger.set_etrecord_file_path(etrecord_file_path) + etrecord.save(etrecord_file_path) + if qnn_config.compile_only: sys.exit(0) diff --git a/backends/qualcomm/partition/utils.py b/backends/qualcomm/partition/utils.py index a83444a56b2..93f00d4e994 100644 --- a/backends/qualcomm/partition/utils.py +++ b/backends/qualcomm/partition/utils.py @@ -68,6 +68,7 @@ def get_skip_decomp_table() -> List[torch._ops.OperatorBase]: torch.ops.aten.reflection_pad2d.default, torch.ops.aten.rms_norm.default, torch.ops.aten._safe_softmax.default, + torch.ops.aten.scatter.src, torch.ops.aten.stack.default, torch.ops.aten.upsample_bicubic2d.vec, # This request is ignored because it is in a blocklist. Refer to exir/program/_program.py diff --git a/backends/qualcomm/qnn_preprocess.py b/backends/qualcomm/qnn_preprocess.py index f423288640c..7ff9a336467 100644 --- a/backends/qualcomm/qnn_preprocess.py +++ b/backends/qualcomm/qnn_preprocess.py @@ -19,7 +19,10 @@ from executorch.backends.qualcomm.serialization.qc_schema_serialize import ( flatbuffer_to_option, ) -from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER +from executorch.backends.qualcomm.utils.constants import ( + QCOM_AXIS_ORDER, + QCOM_TENSOR_NAME, +) from executorch.backends.qualcomm.utils.qnn_manager_lifecycle import ( get_current_qnn_manager, ) @@ -181,10 +184,16 @@ def preprocess_multimethod( # noqa: C901 ) if qnn_manager.IsTensorDump(): for node in programs[i].graph.nodes: - if handle_id := node.meta.get(DEBUG_HANDLE_KEY): + # Skip multi-output nodes: devtools only supports + # single-output intermediate capture (len == 1). + if ( + (handle_id := node.meta.get(DEBUG_HANDLE_KEY)) + and QCOM_TENSOR_NAME in node.meta + and len(node.meta[QCOM_TENSOR_NAME]) == 1 + ): debug_handle_builder.insert_delegate_mapping_entry( handles=handle_id, - identifier=node.name, + identifier=node.meta[QCOM_TENSOR_NAME][0], ) if isinstance(py_op_wrappers, bytes): ctx_binary_list.append(py_op_wrappers) @@ -195,7 +204,6 @@ def preprocess_multimethod( # noqa: C901 for py_op_wrapper in py_op_wrappers ] ) - if len(py_op_wrapper_list) == len(edge_programs.values()): qnn_context_binary = qnn_manager.Compile( graph_names, py_op_wrapper_list diff --git a/backends/qualcomm/quantizer/annotators/htp_rules.py b/backends/qualcomm/quantizer/annotators/htp_rules.py index 342db1cb633..819c9f64136 100644 --- a/backends/qualcomm/quantizer/annotators/htp_rules.py +++ b/backends/qualcomm/quantizer/annotators/htp_rules.py @@ -345,6 +345,8 @@ class ColIm(GeneralOpDef): torch.ops.aten.zeros_like.default, torch.ops.aten.ones.default, torch.ops.aten.ones_like.default, + torch.ops.aten.rand.default, + torch.ops.aten.randn.default, ], qnn_op=None, ) @@ -1389,6 +1391,44 @@ class ScaledDotProductAttention(GeneralOpDef): pass +@register_annotator( + [torch.ops.aten.scatter.src], + qnn_op=None, +) +class ScatterElements(GeneralOpDef): + @staticmethod + def annotate(node: Node, quantization_config: QuantizationConfig) -> None: + if _is_annotated([node]): + return + + input_act = node.args[0] + if not isinstance(input_act, Node) or not _is_float_tensor(input_act): + return + + input_qspec_map = {} + input_qspec_map[input_act] = quantization_config.input_activation + + if ( + len(node.args) > 3 + and isinstance(node.args[3], Node) + and _is_float_tensor(node.args[3]) + ): + input_qspec_map[node.args[3]] = SharedQuantizationSpec((input_act, node)) + + output_act_qspec = ( + SharedQuantizationSpec((input_act, node)) + if _is_float_tensor(node) + else None + ) + + if len(input_qspec_map) > 0 or output_act_qspec is not None: + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=output_act_qspec, + _annotated=True, + ) + + @register_annotator( [torch.ops.aten.sigmoid, torch.ops.aten.sigmoid.default], QnnConstants.OpSigmoid.op_name, diff --git a/backends/qualcomm/quantizer/annotators/lpai_rules.py b/backends/qualcomm/quantizer/annotators/lpai_rules.py index 30a3cb1dc9d..6870882c2db 100644 --- a/backends/qualcomm/quantizer/annotators/lpai_rules.py +++ b/backends/qualcomm/quantizer/annotators/lpai_rules.py @@ -272,6 +272,8 @@ class ColIm(GeneralOpDef): torch.ops.aten.zeros_like.default, torch.ops.aten.ones.default, torch.ops.aten.ones_like.default, + torch.ops.aten.rand.default, + torch.ops.aten.randn.default, ], qnn_op=None, ) @@ -867,6 +869,44 @@ class ScaledDotProductAttention(GeneralOpDef): pass +@register_annotator( + [torch.ops.aten.scatter.src], + qnn_op=None, +) +class ScatterElements(GeneralOpDef): + @staticmethod + def annotate(node: Node, quantization_config: QuantizationConfig) -> None: + if _is_annotated([node]): + return + + input_act = node.args[0] + if not isinstance(input_act, Node) or not _is_float_tensor(input_act): + return + + input_qspec_map = {} + input_qspec_map[input_act] = quantization_config.input_activation + + if ( + len(node.args) > 3 + and isinstance(node.args[3], Node) + and _is_float_tensor(node.args[3]) + ): + input_qspec_map[node.args[3]] = SharedQuantizationSpec((input_act, node)) + + output_act_qspec = ( + SharedQuantizationSpec((input_act, node)) + if _is_float_tensor(node) + else None + ) + + if len(input_qspec_map) > 0 or output_act_qspec is not None: + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=output_act_qspec, + _annotated=True, + ) + + @register_annotator( [torch.ops.aten.sigmoid, torch.ops.aten.sigmoid.default], QnnConstants.OpSigmoid.op_name, diff --git a/backends/qualcomm/runtime/QnnBackendOptions.cpp b/backends/qualcomm/runtime/QnnBackendOptions.cpp index 0eb678b45e2..2117932bddc 100644 --- a/backends/qualcomm/runtime/QnnBackendOptions.cpp +++ b/backends/qualcomm/runtime/QnnBackendOptions.cpp @@ -52,6 +52,14 @@ template QnnExecuTorchProfileLevel get_option( QnnExecuTorchProfileLevel, const char*); +executorch::runtime::Error get_runtime_option( + const char* key, + executorch::runtime::BackendOption& backend_option) { + std::strncpy(backend_option.key, key, runtime::kMaxOptionKeyLength); + backend_option.key[runtime::kMaxOptionKeyLength - 1] = '\0'; + return get_option(QNN_BACKEND, backend_option); +} + } // namespace qnn } // namespace backends } // namespace executorch diff --git a/backends/qualcomm/runtime/QnnBackendOptions.h b/backends/qualcomm/runtime/QnnBackendOptions.h index c366755edd0..93e0de1fb61 100644 --- a/backends/qualcomm/runtime/QnnBackendOptions.h +++ b/backends/qualcomm/runtime/QnnBackendOptions.h @@ -37,6 +37,19 @@ struct RuntimeOption { template T get_option(T aot_option, const char* aot_key); +/** + * @brief + * Get the backend option. + * This method checks runtime option only. + * + * @param key The key of runtime option. + * @param backend_option The backend_option to be restored in runtime. + */ + +executorch::runtime::Error get_runtime_option( + const char* key, + executorch::runtime::BackendOption& backend_option); + } // namespace qnn } // namespace backends } // namespace executorch diff --git a/backends/qualcomm/runtime/QnnExecuTorch.h b/backends/qualcomm/runtime/QnnExecuTorch.h index 8a0ee3fed4b..9699e5b4735 100644 --- a/backends/qualcomm/runtime/QnnExecuTorch.h +++ b/backends/qualcomm/runtime/QnnExecuTorch.h @@ -25,6 +25,7 @@ #define QNN_RUNTIME_LPAI_CLIENT_PERF_TYPE "qnn_runtime_lpai_client_perf_type" #define QNN_RUNTIME_LPAI_AFFINITY "qnn_runtime_lpai_affinity" #define QNN_RUNTIME_LPAI_CORE_SELECTION "qnn_runtime_lpai_core_selection" +#define QNN_RUNTIME_HEAP_PROFILING_PATH "qnn_runtime_heap_profiling_path" #ifdef __cplusplus extern "C" { diff --git a/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp b/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp index 33cca5350d9..fdd70c0a8db 100644 --- a/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp +++ b/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp @@ -245,6 +245,13 @@ executorch::runtime::Error QnnExecuTorchBackend::set_option( qnn_runtime_lpai_core_selection_.value = *val; qnn_runtime_lpai_core_selection_.is_set = true; } + } else if (strcmp(option.key, QNN_RUNTIME_HEAP_PROFILING_PATH) == 0) { + if (auto* val = + std::get_if>( + &option.value)) { + qnn_runtime_heap_profiling_path_.value = *val; + qnn_runtime_heap_profiling_path_.is_set = true; + } } else { ET_LOG( Error, @@ -268,6 +275,7 @@ executorch::runtime::Error QnnExecuTorchBackend::get_option( executorch::runtime::BackendOptionContext& context, executorch::runtime::Span& backend_options) { + std::lock_guard guard(runtime_option_mutex_); size_t matches = backend_options.size(); for (size_t i = 0; i < backend_options.size(); ++i) { // Set the value to what was stored by set_option @@ -303,6 +311,10 @@ executorch::runtime::Error QnnExecuTorchBackend::get_option( strcmp(backend_options[i].key, QNN_RUNTIME_LPAI_CORE_SELECTION) == 0 && qnn_runtime_lpai_core_selection_.is_set) { backend_options[i].value = qnn_runtime_lpai_core_selection_.value; + } else if ( + strcmp(backend_options[i].key, QNN_RUNTIME_HEAP_PROFILING_PATH) == 0 && + qnn_runtime_heap_profiling_path_.is_set) { + backend_options[i].value = qnn_runtime_heap_profiling_path_.value; } else { // either runtime never called set_option or key does not exist matches--; diff --git a/backends/qualcomm/runtime/QnnExecuTorchBackend.h b/backends/qualcomm/runtime/QnnExecuTorchBackend.h index 942e61e2267..e3548c8752b 100644 --- a/backends/qualcomm/runtime/QnnExecuTorchBackend.h +++ b/backends/qualcomm/runtime/QnnExecuTorchBackend.h @@ -71,6 +71,7 @@ class QnnExecuTorchBackend final RuntimeOption qnn_runtime_lpai_client_perf_type_{false, 0}; RuntimeOption qnn_runtime_lpai_affinity_{false, 0}; RuntimeOption qnn_runtime_lpai_core_selection_{false, 0}; + RuntimeOption qnn_runtime_heap_profiling_path_{false, {}}; }; } // namespace qnn diff --git a/backends/qualcomm/runtime/backends/QnnBackendFactory.cpp b/backends/qualcomm/runtime/backends/QnnBackendFactory.cpp index fa2008befd5..4e819a43121 100644 --- a/backends/qualcomm/runtime/backends/QnnBackendFactory.cpp +++ b/backends/qualcomm/runtime/backends/QnnBackendFactory.cpp @@ -71,7 +71,8 @@ std::unique_ptr QnnBackendFactory::Create( qnn_device_ptr, backend_params->qnn_backend_cache_ptr_.get(), htp_options, - qnn_dlc_manager); + qnn_dlc_manager, + get_option(options->profile_level(), QNN_RUNTIME_PROFILE_LEVEL)); backend_params->qnn_graph_ptr_ = std::make_unique( implementation_ptr, diff --git a/backends/qualcomm/runtime/backends/QnnContextCommon.cpp b/backends/qualcomm/runtime/backends/QnnContextCommon.cpp index e16a173db6c..e81f92a8003 100644 --- a/backends/qualcomm/runtime/backends/QnnContextCommon.cpp +++ b/backends/qualcomm/runtime/backends/QnnContextCommon.cpp @@ -6,6 +6,7 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include @@ -13,12 +14,46 @@ namespace executorch { namespace backends { namespace qnn { +std::mutex QnnContext::htp_context_mutex_; +int QnnContext::htp_context_count_{0}; + +void QnnContext::WriteHeapProfile() { + executorch::runtime::BackendOption backend_option; + std::string heap_profiling_path; + if (get_runtime_option(QNN_RUNTIME_HEAP_PROFILING_PATH, backend_option) == + Error::Ok) { + auto* arr = std::get_if>( + &backend_option.value); + if (arr) { + heap_profiling_path = arr->data(); + } + } + Qnn_ErrorHandle_t error_profile = + qnn_profiler_->ProfileDataToFile(heap_profiling_path); + if (error_profile != QNN_SUCCESS) { + QNN_EXECUTORCH_LOG_ERROR( + "Failed to profile. Cannot get profile from handle. Error %d", + QNN_GET_ERROR_CODE(error_profile)); + } +} + QnnContext::~QnnContext() { const QnnInterface& qnn_interface = implementation_->GetQnnInterface(); Qnn_ErrorHandle_t error = QNN_SUCCESS; + if (handle_ != nullptr) { QNN_EXECUTORCH_LOG_INFO("Destroy Qnn context"); - error = qnn_interface.qnn_context_free(handle_, /*profile=*/nullptr); + + bool do_heap_profile = false; + { + std::lock_guard lock(htp_context_mutex_); + if (is_htp_backend_ && htp_context_count_ > 0 && need_to_profile_) { + --htp_context_count_; + do_heap_profile = (htp_context_count_ == 0); + } + } + error = qnn_interface.qnn_context_free( + handle_, do_heap_profile ? qnn_profiler_->GetHandle() : nullptr); if (error != QNN_SUCCESS) { QNN_EXECUTORCH_LOG_ERROR( "Failed to free QNN " @@ -26,6 +61,8 @@ QnnContext::~QnnContext() { "ID %u, error %d", qnn_interface.GetBackendId(), QNN_GET_ERROR_CODE(error)); + } else if (do_heap_profile) { + WriteHeapProfile(); } handle_ = nullptr; } @@ -45,21 +82,51 @@ Error QnnContext::Configure() { if (cache_->GetCacheState() == QnnBackendCache::DESERIALIZE) { const QnnExecuTorchContextBinary& qnn_context_blob = cache_->GetQnnContextBlob(); + /* + Total DSP heap usage can be measured in two conditions, first context + creation and last context free. By the QNN documentation, we need to insert + profileHandle in qnn_context_create_from_binary when creating first context + and closing last context. + + Limitations are two: + 1.Only supported on Android and QNX platforms. + 2.By enabling this feature initialization and cleanup time might be + impacted. + */ + + bool do_heap_profile = false; + { + std::lock_guard lock(htp_context_mutex_); + do_heap_profile = + is_htp_backend_ && (htp_context_count_ == 0) && need_to_profile_; + if (is_htp_backend_) { + ++htp_context_count_; + } + } error = qnn_interface.qnn_context_create_from_binary( backend_->GetHandle(), device_->GetHandle(), - temp_context_config.empty() ? nullptr : temp_context_config.data(), + (temp_context_config.empty() ? nullptr : temp_context_config.data()), static_cast(qnn_context_blob.buffer), qnn_context_blob.nbytes, &handle_, - /*profile=*/nullptr); + do_heap_profile ? qnn_profiler_->GetHandle() : nullptr); if (error != QNN_SUCCESS) { QNN_EXECUTORCH_LOG_ERROR( "Can't create context from " "binary. Error %d.", QNN_GET_ERROR_CODE(error)); + // Rollback the count since context creation failed + { + std::lock_guard lock(htp_context_mutex_); + if (is_htp_backend_ && htp_context_count_ > 0) { + --htp_context_count_; + } + } return Error::Internal; + } else if (do_heap_profile) { + WriteHeapProfile(); } } else if ( cache_->GetCacheState() == QnnBackendCache::SERIALIZE || diff --git a/backends/qualcomm/runtime/backends/QnnContextCommon.h b/backends/qualcomm/runtime/backends/QnnContextCommon.h index 7d507a4a50c..c0351b857b7 100644 --- a/backends/qualcomm/runtime/backends/QnnContextCommon.h +++ b/backends/qualcomm/runtime/backends/QnnContextCommon.h @@ -13,7 +13,10 @@ #include #include +#include + #include +#include namespace executorch { namespace backends { @@ -28,13 +31,22 @@ class QnnContext { QnnBackend* backend, QnnDevice* device, QnnBackendCache* cache, - QnnDlcManager* qnn_dlc_manager) + QnnDlcManager* qnn_dlc_manager, + const QnnExecuTorchProfileLevel& profile_level) : handle_(nullptr), implementation_(implementation), backend_(backend), device_(device), cache_(cache), - qnn_dlc_manager_(qnn_dlc_manager) {} + qnn_dlc_manager_(qnn_dlc_manager), + is_htp_backend_( + implementation->GetQnnInterface().GetBackendId() == + QNN_BACKEND_ID_HTP), + need_to_profile_( + profile_level != QnnExecuTorchProfileLevel::kProfileOff) { + qnn_profiler_ = + std::make_unique(implementation_, backend_, profile_level); + } virtual ~QnnContext(); @@ -73,6 +85,7 @@ class QnnContext { }; private: + void WriteHeapProfile(); Qnn_ContextHandle_t handle_; QnnImplementation* implementation_; QnnBackend* backend_; @@ -80,6 +93,12 @@ class QnnContext { QnnBackendCache* cache_; QnnContextCustomProtocol qnn_context_custom_protocol_; QnnDlcManager* qnn_dlc_manager_; + + std::unique_ptr qnn_profiler_; + bool is_htp_backend_; + bool need_to_profile_; + static std::mutex htp_context_mutex_; + static int htp_context_count_; }; } // namespace qnn } // namespace backends diff --git a/backends/qualcomm/runtime/backends/QnnProfiler.cpp b/backends/qualcomm/runtime/backends/QnnProfiler.cpp index b4650b30796..195c967a674 100644 --- a/backends/qualcomm/runtime/backends/QnnProfiler.cpp +++ b/backends/qualcomm/runtime/backends/QnnProfiler.cpp @@ -8,10 +8,34 @@ #include +#include +#include + namespace executorch { namespace backends { namespace qnn { +#define DEFINE_HEAP_BEFORE_CREATION "DSP:before_context_created" +#define DEFINE_HEAP_AFTER_FREED "DSP:after_context_freed" + +namespace { +const char* get_event_unit(QnnProfile_EventUnit_t unit) { + switch (unit) { + case QNN_PROFILE_EVENTUNIT_MICROSEC: + return " (us)"; + case QNN_PROFILE_EVENTUNIT_BYTES: + return " (bytes)"; + case QNN_PROFILE_EVENTUNIT_COUNT: + return " (count)"; + case QNN_PROFILE_EVENTUNIT_BACKEND: + // cycle unit is default appeared + case QNN_PROFILE_EVENTUNIT_CYCLES: + default: + return ""; + } +} +} // namespace + QnnProfile::QnnProfile( QnnImplementation* implementation, QnnBackend* backend, @@ -71,36 +95,36 @@ QnnProfile::QnnProfile( } } +Qnn_ErrorHandle_t QnnProfile::FetchEvents( + const QnnProfile_EventId_t** events_ptr, + std::uint32_t* num_events) { + if (handle_ == nullptr) { + QNN_EXECUTORCH_LOG_WARN("Profile handle is null, skipping FetchEvents"); + *num_events = 0; + return QNN_SUCCESS; + } + const QnnInterface& qnn_interface = implementation_->GetQnnInterface(); + Qnn_ErrorHandle_t error = + qnn_interface.qnn_profile_get_events(handle_, events_ptr, num_events); + if (error != QNN_SUCCESS) { + QNN_EXECUTORCH_LOG_ERROR( + "Failed to get profile events: %d", QNN_GET_ERROR_CODE(error)); + } + return error; +} + Qnn_ErrorHandle_t QnnProfile::ProfileData( executorch::runtime::EventTracer* event_tracer) { - const QnnInterface& qnn_interface = implementation_->GetQnnInterface(); const QnnProfile_EventId_t* events_ptr = nullptr; - const QnnProfile_EventId_t* sub_events_ptr = nullptr; std::uint32_t num_events = 0; - std::uint32_t num_sub_events = 0; - Qnn_ErrorHandle_t error = - qnn_interface.qnn_profile_get_events(handle_, &events_ptr, &num_events); + Qnn_ErrorHandle_t error = FetchEvents(&events_ptr, &num_events); if (error != QNN_SUCCESS) { QNN_EXECUTORCH_LOG_ERROR( - "ProfileData failed to get events: %d", QNN_GET_ERROR_CODE(error)); + "Failed to profile data in function FetchEvents: %d", + QNN_GET_ERROR_CODE(error)); return error; } - - auto get_unit = [](QnnProfile_EventUnit_t unit) { - switch (unit) { - case QNN_PROFILE_EVENTUNIT_MICROSEC: - return " (us)"; - case QNN_PROFILE_EVENTUNIT_BYTES: - return " (bytes)"; - case QNN_PROFILE_EVENTUNIT_COUNT: - return " (count)"; - case QNN_PROFILE_EVENTUNIT_BACKEND: - // cycle unit is default appeared - case QNN_PROFILE_EVENTUNIT_CYCLES: - default: - return ""; - } - }; + const QnnInterface& qnn_interface = implementation_->GetQnnInterface(); QnnProfile_EventData_t event_data; for (std::uint32_t i = 0; i < num_events; ++i) { error = @@ -115,7 +139,7 @@ Qnn_ErrorHandle_t QnnProfile::ProfileData( } // add events for other important metrics, e.g. RPC execution time std::string identifier = - std::string(event_data.identifier) + get_unit(event_data.unit); + std::string(event_data.identifier) + get_event_unit(event_data.unit); executorch::runtime::event_tracer_log_profiling_delegate( event_tracer, identifier.c_str(), @@ -125,48 +149,114 @@ Qnn_ErrorHandle_t QnnProfile::ProfileData( event_data.value); // Check an event's sub events only if it relates to graph execution time // (and its sub events are the individual op executions): - if (backend_->IsProfileEventTypeParentOfNodeTime(event_data.type)) { - error = qnn_interface.qnn_profile_get_sub_events( - events_ptr[i], &sub_events_ptr, &num_sub_events); + if (!backend_->IsProfileEventTypeParentOfNodeTime(event_data.type)) { + continue; + } + const QnnProfile_EventId_t* sub_events_ptr = nullptr; + std::uint32_t num_sub_events = 0; + error = qnn_interface.qnn_profile_get_sub_events( + events_ptr[i], &sub_events_ptr, &num_sub_events); + if (error != QNN_SUCCESS) { + QNN_EXECUTORCH_LOG_ERROR( + "ProfileData failed to get sub events " + "for event %d: %d", + i, + QNN_GET_ERROR_CODE(error)); + return error; + } + + QnnProfile_EventData_t sub_event_data; + for (std::uint32_t j = 0; j < num_sub_events; ++j) { + error = qnn_interface.qnn_profile_get_event_data( + sub_events_ptr[j], &sub_event_data); if (error != QNN_SUCCESS) { QNN_EXECUTORCH_LOG_ERROR( - "ProfileData failed to get sub events " - "for event %d: %d", + "ProfileData failed to get sub " + "event data for sub event %d of event %d: %d", + j, i, QNN_GET_ERROR_CODE(error)); return error; } - - QnnProfile_EventData_t sub_event_data; - for (std::uint32_t j = 0; j < num_sub_events; ++j) { - error = qnn_interface.qnn_profile_get_event_data( - sub_events_ptr[j], &sub_event_data); - if (error != QNN_SUCCESS) { - QNN_EXECUTORCH_LOG_ERROR( - "ProfileData failed to get sub " - "event data for sub event %d of event %d: %d", - j, - i, - QNN_GET_ERROR_CODE(error)); - return error; - } - if (sub_event_data.type == QNN_PROFILE_EVENTTYPE_NODE && - (sub_event_data.unit == QNN_PROFILE_EVENTUNIT_MICROSEC || - sub_event_data.unit == QNN_PROFILE_EVENTUNIT_CYCLES)) { - executorch::runtime::event_tracer_log_profiling_delegate( - event_tracer, - sub_event_data.identifier, - /*delegate_debug_id=*/ - static_cast(-1), - 0, - sub_event_data.value); - } + if (sub_event_data.type == QNN_PROFILE_EVENTTYPE_NODE && + (sub_event_data.unit == QNN_PROFILE_EVENTUNIT_MICROSEC || + sub_event_data.unit == QNN_PROFILE_EVENTUNIT_CYCLES)) { + executorch::runtime::event_tracer_log_profiling_delegate( + event_tracer, + sub_event_data.identifier, + /*delegate_debug_id=*/ + static_cast(-1), + 0, + sub_event_data.value); } } } return error; } +Qnn_ErrorHandle_t QnnProfile::ProfileDataToFile( + const std::string& profile_filename) { + if (handle_ == nullptr) { + QNN_EXECUTORCH_LOG_WARN( + "Profile handle is null, skipping ProfileDataToFile"); + return QNN_SUCCESS; + } + if (profile_filename.empty()) { + QNN_EXECUTORCH_LOG_WARN( + "Heap profiling path is empty. Please provide profiling filename from runtime option."); + return QNN_SUCCESS; + } + const QnnProfile_EventId_t* events_ptr = nullptr; + std::uint32_t num_events = 0; + Qnn_ErrorHandle_t error = FetchEvents(&events_ptr, &num_events); + if (error != QNN_SUCCESS) { + QNN_EXECUTORCH_LOG_ERROR( + "Failed to profile data in function FetchEvents: %d", + QNN_GET_ERROR_CODE(error)); + return error; + } + const QnnInterface& qnn_interface = implementation_->GetQnnInterface(); + QnnProfile_EventData_t event_data; + std::uint32_t count_num_events = 0; + for (std::uint32_t i = 0; i < num_events; ++i) { + error = + qnn_interface.qnn_profile_get_event_data(events_ptr[i], &event_data); + if (error != QNN_SUCCESS) { + QNN_EXECUTORCH_LOG_ERROR( + "ProfileData failed to get event data " + "for event %d: %d", + i, + QNN_GET_ERROR_CODE(error)); + return error; + } + + std::ios_base::openmode open_mode = std::ios::app; + if (strcmp(event_data.identifier, DEFINE_HEAP_BEFORE_CREATION) == 0) { + open_mode = std::ios::trunc; + } else if (strcmp(event_data.identifier, DEFINE_HEAP_AFTER_FREED) == 0) { + open_mode = std::ios::app; + } else { + count_num_events++; + continue; + } + std::string identifier = + std::string(event_data.identifier) + get_event_unit(event_data.unit); + std::ofstream ofs(profile_filename, open_mode); + if (!ofs) { + QNN_EXECUTORCH_LOG_ERROR( + "Error when opening profile file: %s", profile_filename.c_str()); + return QNN_COMMON_ERROR_GENERAL; + } + ofs << identifier << ", " << event_data.value << "\n"; + } + if (count_num_events == num_events) { + QNN_EXECUTORCH_LOG_WARN( + "Not HTP backend but enable htp profiling. Please check setting."); + return QNN_SUCCESS; + } + return QNN_SUCCESS; +} + QnnProfile::~QnnProfile() { const QnnInterface& qnn_interface = implementation_->GetQnnInterface(); if (handle_ != nullptr) { diff --git a/backends/qualcomm/runtime/backends/QnnProfiler.h b/backends/qualcomm/runtime/backends/QnnProfiler.h index de8fbd1d9d5..e8f2a3c0502 100644 --- a/backends/qualcomm/runtime/backends/QnnProfiler.h +++ b/backends/qualcomm/runtime/backends/QnnProfiler.h @@ -12,6 +12,9 @@ #include #include #include "QnnProfile.h" + +#include + namespace executorch { namespace backends { namespace qnn { @@ -24,6 +27,7 @@ class QnnProfile { const QnnExecuTorchProfileLevel& profile_level); ~QnnProfile(); Qnn_ErrorHandle_t ProfileData(executorch::runtime::EventTracer* event_tracer); + Qnn_ErrorHandle_t ProfileDataToFile(const std::string& profile_filename); Qnn_ProfileHandle_t GetHandle() { return handle_; @@ -33,6 +37,10 @@ class QnnProfile { Qnn_ProfileHandle_t handle_; QnnImplementation* implementation_; QnnBackend* backend_; + + Qnn_ErrorHandle_t FetchEvents( + const QnnProfile_EventId_t** events_ptr, + std::uint32_t* num_events); }; } // namespace qnn } // namespace backends diff --git a/backends/qualcomm/runtime/backends/gpu/GpuContext.cpp b/backends/qualcomm/runtime/backends/gpu/GpuContext.cpp index 07952e77eef..c6c6ace2bdf 100644 --- a/backends/qualcomm/runtime/backends/gpu/GpuContext.cpp +++ b/backends/qualcomm/runtime/backends/gpu/GpuContext.cpp @@ -21,7 +21,13 @@ GpuContext::GpuContext( QnnBackendCache* cache, QnnDlcManager* qnn_dlc_manager, const QnnExecuTorchGpuBackendOptions* gpu_options) - : QnnContext(implementation, backend, device, cache, qnn_dlc_manager) { + : QnnContext( + implementation, + backend, + device, + cache, + qnn_dlc_manager, + QnnExecuTorchProfileLevel::kProfileOff) { gpu_context_custom_config_ = std::make_unique(gpu_options); } diff --git a/backends/qualcomm/runtime/backends/htp/HtpContext.h b/backends/qualcomm/runtime/backends/htp/HtpContext.h index a0389ea5983..f00b709f607 100644 --- a/backends/qualcomm/runtime/backends/htp/HtpContext.h +++ b/backends/qualcomm/runtime/backends/htp/HtpContext.h @@ -25,10 +25,17 @@ class HtpContext : public QnnContext { QnnDevice* device, QnnBackendCache* cache, const QnnExecuTorchHtpBackendOptions* htp_options, - QnnDlcManager* qnn_dlc_manager) - : QnnContext(implementation, backend, device, cache, qnn_dlc_manager) { - htp_context_custom_config_ = - std::make_unique(this, htp_options); + QnnDlcManager* qnn_dlc_manager, + const QnnExecuTorchProfileLevel& profile_level) + : QnnContext( + implementation, + backend, + device, + cache, + qnn_dlc_manager, + profile_level) { + htp_context_custom_config_ = std::make_unique( + this, htp_options, profile_level); } ~HtpContext() {} diff --git a/backends/qualcomm/runtime/backends/htp/HtpContextCustomConfig.h b/backends/qualcomm/runtime/backends/htp/HtpContextCustomConfig.h index f0d4873b0d2..61a395fcb5b 100644 --- a/backends/qualcomm/runtime/backends/htp/HtpContextCustomConfig.h +++ b/backends/qualcomm/runtime/backends/htp/HtpContextCustomConfig.h @@ -26,12 +26,19 @@ class HtpContextCustomConfig { public: explicit HtpContextCustomConfig( const QnnContext* context, - const QnnExecuTorchHtpBackendOptions* htp_options) - : context_(context), htp_options_(htp_options) {} + const QnnExecuTorchHtpBackendOptions* htp_options, + const QnnExecuTorchProfileLevel& profile_level) + : profile_level_(profile_level), + context_(context), + htp_options_(htp_options) {} std::vector CreateContextCustomConfig(); private: + // profile_level_ is consumed only by the target build; the host build never + // reads it. Marked [[maybe_unused]] so the host build doesn't warn while the + // field stays available for the target side. + [[maybe_unused]] QnnExecuTorchProfileLevel profile_level_; QnnHtpContext_CustomConfig_t* AllocContextCustomConfig() { htp_context_config_.emplace_back( std::make_unique()); diff --git a/backends/qualcomm/runtime/backends/htp/host/HtpContextCustomConfig.cpp b/backends/qualcomm/runtime/backends/htp/host/HtpContextCustomConfig.cpp index 4850afa14a2..037998132a8 100644 --- a/backends/qualcomm/runtime/backends/htp/host/HtpContextCustomConfig.cpp +++ b/backends/qualcomm/runtime/backends/htp/host/HtpContextCustomConfig.cpp @@ -6,6 +6,7 @@ * LICENSE file in the root directory of this source tree. */ +#include #include namespace executorch { diff --git a/backends/qualcomm/runtime/backends/htp/target/HtpContextCustomConfig.cpp b/backends/qualcomm/runtime/backends/htp/target/HtpContextCustomConfig.cpp index 676795797f8..8488bf21e79 100644 --- a/backends/qualcomm/runtime/backends/htp/target/HtpContextCustomConfig.cpp +++ b/backends/qualcomm/runtime/backends/htp/target/HtpContextCustomConfig.cpp @@ -19,6 +19,17 @@ HtpContextCustomConfig::CreateContextCustomConfig() { QnnHtpContext_CustomConfig_t* p_custom_config = nullptr; const HtpContext* htp_ctx = static_cast(context_); + // TODO: Verify heap profile works with kProfileBasic once enabled. + if (profile_level_ != QnnExecuTorchProfileLevel::kProfileOff) { + QnnHtpContext_CustomConfig_t* p_custom_config_profile = nullptr; + p_custom_config_profile = AllocContextCustomConfig(); + p_custom_config_profile->option = + QNN_HTP_CONTEXT_CONFIG_OPTION_DSP_MEMORY_PROFILING_ENABLED; + p_custom_config_profile->dspMemoryProfilingEnabled = true; + ret.push_back( + static_cast(p_custom_config_profile)); + } + if (htp_options_->use_multi_contexts() && htp_options_->max_sf_buf_size() != 0) { p_custom_config = AllocContextCustomConfig(); diff --git a/backends/qualcomm/runtime/backends/ir/host/QnnDlcManager.cpp b/backends/qualcomm/runtime/backends/ir/host/QnnDlcManager.cpp index 47d583b5c15..62d01c78706 100644 --- a/backends/qualcomm/runtime/backends/ir/host/QnnDlcManager.cpp +++ b/backends/qualcomm/runtime/backends/ir/host/QnnDlcManager.cpp @@ -47,7 +47,8 @@ Error QnnDlcManager::Create() { backend_bundle_ptr_->qnn_backend_ptr.get(), backend_bundle_ptr_->qnn_device_ptr.get(), backend_params_ptr_->qnn_backend_cache_ptr_.get(), - nullptr); + nullptr, + QnnExecuTorchProfileLevel::kProfileOff); backend_params_ptr_->qnn_graph_ptr_ = std::make_unique( backend_bundle_ptr_->implementation.get(), diff --git a/backends/qualcomm/runtime/backends/lpai/LpaiContext.cpp b/backends/qualcomm/runtime/backends/lpai/LpaiContext.cpp index d5203898f6b..e0c9d3ed3d8 100644 --- a/backends/qualcomm/runtime/backends/lpai/LpaiContext.cpp +++ b/backends/qualcomm/runtime/backends/lpai/LpaiContext.cpp @@ -20,7 +20,13 @@ LpaiContext::LpaiContext( QnnDevice* device, QnnBackendCache* cache, QnnDlcManager* qnn_dlc_manager) - : QnnContext(implementation, backend, device, cache, qnn_dlc_manager) { + : QnnContext( + implementation, + backend, + device, + cache, + qnn_dlc_manager, + QnnExecuTorchProfileLevel::kProfileOff) { lpai_context_custom_config_ = std::make_unique(); } diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index 12d5e0902db..e190828f06a 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -875,6 +875,31 @@ def forward(self, x): return self.second(self.first(x)) +class ConvFull(torch.nn.Module): + def __init__(self, fill, full_shape): + super().__init__() + self.conv = torch.nn.Conv2d(8, 16, 3, padding=1) + self.fill = fill + self.full_shape = full_shape + + def forward(self, x): + y = self.conv(x) + c = torch.full(self.full_shape, self.fill, dtype=y.dtype) + return torch.cat([y, c], dim=1) + + +class ConvFullLike(torch.nn.Module): + def __init__(self, fill): + super().__init__() + self.conv = torch.nn.Conv2d(8, 16, 3, padding=1) + self.fill = fill + + def forward(self, x): + y = self.conv(x) + c = torch.full_like(y, self.fill) + return torch.cat([y, c], dim=1) + + class ConvTranspose1dSingle(torch.nn.Module): def __init__(self, bias=True, dilation=1): super().__init__() @@ -1929,6 +1954,14 @@ def forward(self, x): return torch.rand_like(x) + x +class Randn(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.randn_like(x) + x + + class Reciprocal(torch.nn.Module): def __init__(self): super().__init__() @@ -2168,6 +2201,15 @@ def forward(self, query_layer, key_layer, value_layer, attn_mask): return attn_output +class ScatterSrc(torch.nn.Module): + def __init__(self, dim=1): + super().__init__() + self.dim = dim + + def forward(self, data, index, src): + return torch.scatter(data, self.dim, index, src) + + class SelectCopy(torch.nn.Module): def __init__(self): super().__init__() @@ -2425,6 +2467,14 @@ def forward(self, x): return torch.swapaxes(x, axis0=self.axis0, axis1=self.axis1) +class Tan(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.tan(x) + + class Tanh(torch.nn.Module): def __init__(self): super().__init__() diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index 688dddf5c2a..940c54c2f8d 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -120,7 +120,7 @@ def setUp(self): debug=False, saver=False, online_prepare=TestQNN.online_prepare, - dump_intermediate_outputs=TestQNN.dump_intermediate_outputs, + dump_intermediate_outputs=False, profile_level=TestQNN.profile_level, shared_buffer=TestQNN.shared_buffer, ) @@ -1809,24 +1809,6 @@ def test_qnn_backend_prelu(self): index += 1 self.lower_module_and_test_output(module, sample_input) - def test_qnn_backend_rand(self): - sample_inputs = [ - (torch.randn(3, 4, 5),), - (torch.randn(2, 8),), - ( - torch.randn( - 10, - ), - ), - (torch.randn(1, 3, 32, 32),), - ] - for i, sample_input in enumerate(sample_inputs): - with self.subTest(i=i): - module = Rand() # noqa: F405 - self.lower_module_and_test_output( - module, sample_input, assert_output_equal=False - ) - def test_qnn_backend_reciprocal(self): module = Reciprocal() # noqa: F405 sample_input = (torch.randn([2, 2, 2, 2]),) @@ -1948,6 +1930,52 @@ def test_qnn_backend_round(self): sample_input = (torch.randn([3, 4]),) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_scatter_src(self): + test_comb = [ + { + QCOM_MODULE: [ScatterSrc(dim=1)], # noqa: F405 + QCOM_SAMPLE_INPUTS: [ + ( + torch.zeros(3, 5), + torch.tensor( + [[0, 1, 2, 3, 4], [4, 3, 2, 1, 0], [1, 0, 3, 4, 2]], + dtype=torch.int64, + ), + torch.rand(3, 5), + ), + ( + torch.zeros(3, 5, dtype=torch.float16), + torch.tensor( + [[0, 1, 2, 3, 4], [4, 3, 2, 1, 0], [1, 0, 3, 4, 2]], + dtype=torch.int64, + ), + torch.rand(3, 5, dtype=torch.float16), + ), + ], + }, + { + QCOM_MODULE: [ScatterSrc(dim=0)], # noqa: F405 + QCOM_SAMPLE_INPUTS: [ + ( + torch.zeros(3, 5), + torch.tensor( + [[2, 1, 0, 1, 2], [0, 2, 1, 2, 0], [1, 0, 2, 0, 1]], + dtype=torch.int64, + ), + torch.rand(3, 5), + ), + ], + }, + ] + + index = 0 + for comb in test_comb: + for module in comb[QCOM_MODULE]: + for sample_input in comb[QCOM_SAMPLE_INPUTS]: + with self.subTest(i=index): + index += 1 + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_rsqrt(self): module = Rsqrt() # noqa: F405 sample_input = (torch.abs(torch.randn([3, 4])),) @@ -2070,6 +2098,11 @@ def test_qnn_backend_swapaxes(self): sample_input = (torch.randn([1, 2, 3, 4]),) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_tan(self): + module = Tan() # noqa: F405 + sample_input = (torch.rand(2, 5, 1, 3) * 2 - 1,) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_tanh(self): module = Tanh() # noqa: F405 sample_input = (torch.randn(2, 5, 1, 3),) @@ -2219,7 +2252,7 @@ def setUp(self): debug=False, saver=False, online_prepare=TestQNN.online_prepare, - dump_intermediate_outputs=TestQNN.dump_intermediate_outputs, + dump_intermediate_outputs=False, profile_level=TestQNN.profile_level, shared_buffer=TestQNN.shared_buffer, ) @@ -2351,6 +2384,17 @@ def test_qnn_backend_einsum_outer_product_relu(self): ) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_full_layout_transformed(self): + full_shape = (1, 16, 4, 6) + module = ConvFull(0.5, full_shape) # noqa: F405 + sample_input = (torch.randn(1, 8, 4, 6),) + self.lower_module_and_test_output(module, sample_input) + + def test_qnn_backend_full_like_layout_transformed(self): + module = ConvFullLike(0.5) # noqa: F405 + sample_input = (torch.randn(1, 8, 4, 6),) + self.lower_module_and_test_output(module, sample_input) + # TODO: Create a new UT class for passes specific checks def test_qnn_backend_lift_add_tensor(self): module = LiftAddTensor() # noqa: F405 @@ -2463,7 +2507,7 @@ def setUp(self): debug=False, saver=False, online_prepare=TestQNN.online_prepare, - dump_intermediate_outputs=TestQNN.dump_intermediate_outputs, + dump_intermediate_outputs=False, profile_level=TestQNN.profile_level, shared_buffer=TestQNN.shared_buffer, ) @@ -4555,6 +4599,7 @@ def test_qnn_backend_prelu(self): self.lower_module_and_test_output(module, sample_input) def test_qnn_backend_rand(self): + module = Rand() # noqa: F405 sample_inputs = [ (torch.randn(3, 4, 5),), (torch.randn(2, 8),), @@ -4567,10 +4612,28 @@ def test_qnn_backend_rand(self): ] for i, sample_input in enumerate(sample_inputs): with self.subTest(i=i): - module = Rand() # noqa: F405 - module = self.get_qdq_module(module, sample_input) + qdq_module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output( - module, sample_input, assert_output_equal=False + qdq_module, sample_input, assert_output_equal=False + ) + + def test_qnn_backend_randn(self): + module = Randn() # noqa: F405 + sample_inputs = [ + (torch.randn(3, 4, 5),), + (torch.randn(2, 8),), + ( + torch.randn( + 10, + ), + ), + (torch.randn(1, 3, 32, 32),), + ] + for i, sample_input in enumerate(sample_inputs): + with self.subTest(i=i): + qdq_module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output( + qdq_module, sample_input, assert_output_equal=False ) def test_qnn_backend_reciprocal(self): @@ -4705,6 +4768,45 @@ def test_qnn_backend_rsqrt(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_scatter_src(self): + test_comb = [ + { + QCOM_MODULE: [ScatterSrc(dim=1)], # noqa: F405 + QCOM_SAMPLE_INPUTS: [ + ( + torch.zeros(3, 5), + torch.tensor( + [[0, 1, 2, 3, 4], [4, 3, 2, 1, 0], [1, 0, 3, 4, 2]], + dtype=torch.int64, + ), + torch.rand(3, 5), + ), + ], + }, + { + QCOM_MODULE: [ScatterSrc(dim=0)], # noqa: F405 + QCOM_SAMPLE_INPUTS: [ + ( + torch.zeros(3, 5), + torch.tensor( + [[2, 1, 0, 1, 2], [0, 2, 1, 2, 0], [1, 0, 2, 0, 1]], + dtype=torch.int64, + ), + torch.rand(3, 5), + ), + ], + }, + ] + + index = 0 + for comb in test_comb: + for module in comb[QCOM_MODULE]: + for sample_input in comb[QCOM_SAMPLE_INPUTS]: + with self.subTest(i=index): + index += 1 + qdq_module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(qdq_module, sample_input) + def test_qnn_backend_sdpa(self): modules = [ ScaledDotProductAttention(), # noqa: F405 @@ -4842,6 +4944,12 @@ def test_qnn_backend_swapaxes(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_tan(self): + module = Tan() # noqa: F405 + sample_input = (torch.rand(2, 5, 1, 3) * 2 - 1,) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_tanh(self): module = Tanh() # noqa: F405 sample_input = (torch.randn(2, 5, 1, 3),) @@ -5010,7 +5118,7 @@ def setUp(self): debug=False, saver=False, online_prepare=TestQNN.online_prepare, - dump_intermediate_outputs=TestQNN.dump_intermediate_outputs, + dump_intermediate_outputs=False, profile_level=TestQNN.profile_level, shared_buffer=TestQNN.shared_buffer, ) @@ -5270,6 +5378,19 @@ def test_qnn_backend_einsum_outer_product_relu(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_full_layout_transformed(self): + full_shape = (1, 16, 4, 6) + module = ConvFull(0.5, full_shape) # noqa: F405 + sample_input = (torch.randn(1, 8, 4, 6),) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + + def test_qnn_backend_full_like_layout_transformed(self): + module = ConvFullLike(0.5) # noqa: F405 + sample_input = (torch.randn(1, 8, 4, 6),) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + @unittest.skipIf(is_qnn_sdk_version_less_than("2.35"), "UT pass after QNN 2.35") def test_qnn_backend_masked_softmax(self): if self.enable_x86_64: @@ -5483,7 +5604,6 @@ class TestQNNFloatingPointUtils(TestQNN): def setUp(self): TestQNN.atol = 1e-1 TestQNN.rtol = 1e-1 - TestQNN.dump_intermediate_outputs = False TestQNN.enable_profile = False TestQNN.shared_buffer = False backend_options = generate_htp_compiler_spec(use_fp16=True) @@ -5522,7 +5642,6 @@ def output_callback(log_msg): ) def test_qnn_backend_dump_intermediate_outputs_topk(self): - TestQNN.dump_intermediate_outputs = True backend_options = generate_htp_compiler_spec(use_fp16=True) TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( soc_model=self.chipset_table[TestQNN.soc_model], @@ -5535,12 +5654,10 @@ def test_qnn_backend_dump_intermediate_outputs_topk(self): module, sample_input, expected_partitions=1, - expected_intermediate_events=7, - expected_compared_events=5, + expected_compared_events=3, ) def test_qnn_backend_dump_intermediate_outputs_simple_model(self): - TestQNN.dump_intermediate_outputs = True backend_options = generate_htp_compiler_spec(use_fp16=True) TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( soc_model=self.chipset_table[TestQNN.soc_model], @@ -5553,8 +5670,7 @@ def test_qnn_backend_dump_intermediate_outputs_simple_model(self): module, sample_input, expected_partitions=1, - expected_intermediate_events=20, - expected_compared_events=16, + expected_compared_events=14, ) def test_qnn_backend_skip_node_id(self): @@ -5703,6 +5819,42 @@ def test_qnn_backend_profile_op(self): ) TestQNN.profile_level = 0 + def test_qnn_backend_runtime_option_heap_profile(self): + module = SimpleModel() # noqa: F405 + sample_input = (torch.ones(1, 32, 28, 28), torch.ones(1, 32, 28, 28)) + + backend_options = generate_htp_compiler_spec( + use_fp16=True, + use_multi_contexts=True, + ) + + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=self.chipset_table[TestQNN.soc_model], + backend_options=backend_options, + profile_level=2, # if 0 for closing heap profiling + ) + + pass_jobs = get_capture_program_passes() + split_graph_pass, setting = self.split_graph(4) + pass_jobs[split_graph_pass] = setting + dep_table = get_passes_dependency_for_capture_program() + dep_table[split_graph_pass] = [FoldQDQ] + + edge_prog_mgr = to_edge_transform_and_lower_to_qnn( + module=module, + inputs=sample_input, + compiler_specs=compiler_specs, + dep_table=dep_table, + passes_job=pass_jobs, + ) + exec_prog = edge_prog_mgr.to_executorch() + self.verify_output( + module, + sample_input, + exec_prog, + save_heap_result=True, + ) + def test_qnn_backend_runtime_option_htp_performance(self): backend_options = generate_htp_compiler_spec(use_fp16=True) TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( @@ -6112,7 +6264,6 @@ def setUp(self): raise ValueError("Backend is not implemented yet") TestQNN.atol = 1e-1 TestQNN.rtol = 1 - TestQNN.dump_intermediate_outputs = False TestQNN.enable_profile = False TestQNN.shared_buffer = False backend_options = generate_htp_compiler_spec(use_fp16=False) @@ -6152,7 +6303,6 @@ def output_callback(log_msg): ) def test_qnn_backend_dump_intermediate_outputs_simple_model(self): - TestQNN.dump_intermediate_outputs = True backend_options = generate_htp_compiler_spec(use_fp16=False) TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( soc_model=self.chipset_table[TestQNN.soc_model], @@ -6166,13 +6316,11 @@ def test_qnn_backend_dump_intermediate_outputs_simple_model(self): module, sample_input, expected_partitions=1, - expected_intermediate_events=21, expected_compared_events=14, ) def test_qnn_backend_dump_intermediate_outputs_topk(self): torch.manual_seed(8) - TestQNN.dump_intermediate_outputs = True backend_options = generate_htp_compiler_spec(use_fp16=False) TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( soc_model=self.chipset_table[TestQNN.soc_model], @@ -6186,8 +6334,7 @@ def test_qnn_backend_dump_intermediate_outputs_topk(self): module, sample_input, expected_partitions=1, - expected_intermediate_events=9, - expected_compared_events=5, + expected_compared_events=3, ) def test_qnn_backend_dynamic_shape(self): @@ -6607,6 +6754,43 @@ def test_qnn_backend_profile_op(self): ) TestQNN.profile_level = 0 + def test_qnn_backend_runtime_option_heap_profile(self): + module = SimpleModel() # noqa: F405 + sample_input = (torch.ones(1, 32, 28, 28), torch.ones(1, 32, 28, 28)) + module1 = self.get_qdq_module(module, sample_input) + + backend_options = generate_htp_compiler_spec( + use_fp16=False, + use_multi_contexts=True, + ) + + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=self.chipset_table[TestQNN.soc_model], + backend_options=backend_options, + profile_level=2, # if 0 for closing heap profiling + ) + + pass_jobs = get_capture_program_passes() + split_graph_pass, setting = self.split_graph(4) + pass_jobs[split_graph_pass] = setting + dep_table = get_passes_dependency_for_capture_program() + dep_table[split_graph_pass] = [FoldQDQ] + + edge_prog_mgr = to_edge_transform_and_lower_to_qnn( + module=module1, + inputs=sample_input, + compiler_specs=compiler_specs, + dep_table=dep_table, + passes_job=pass_jobs, + ) + exec_prog = edge_prog_mgr.to_executorch() + self.verify_output( + module1, + sample_input, + exec_prog, + save_heap_result=True, + ) + def test_qnn_backend_runtime_option_htp_performance(self): backend_options = generate_htp_compiler_spec(use_fp16=False) TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( @@ -9665,9 +9849,7 @@ def test_intermediate_debugger(self): else: svg_path = msg["svg_path"] csv_path = msg["csv_path"] - min_accepted = 235 - max_accepted = 241 - # Having a +- 3 tolerance, expecting 238 events + expected_accepted_events = 234 assert os.path.exists(svg_path), f"Unable to find SVG file: {svg_path}" assert os.path.exists(csv_path), f"Unable to find CSV file: {csv_path}" @@ -9679,19 +9861,23 @@ def test_intermediate_debugger(self): for row in reader: if len(row) > index and row[index].strip().upper() == "TRUE": csv_valid_count += 1 - # We assume csv_valid_count == compared_events, since all compared events meet metric's threshold - assert ( - min_accepted <= csv_valid_count <= max_accepted - ), f"Expected CSV events with valid score is outside of expected range, number of valid score events found: {csv_valid_count}" + # We assume csv_valid_count == expected_accepted_events, since all compared events meet metric's threshold + self.assertEqual( + expected_accepted_events, + csv_valid_count, + msg=f"Expected CSV events: {expected_accepted_events}, found: {csv_valid_count}.", + ) svg_valid_count = 0 with open(svg_path, "r", encoding="utf-8") as svg_file: for line in svg_file: svg_valid_count += line.count("is_valid_score=True") - # We assume svg_valid_count == compared_events, since all compared events meet metric's threshold - assert ( - min_accepted <= svg_valid_count <= max_accepted - ), f"Expected SVG events with valid score is outside of expected range, number of valid score events found: {svg_valid_count}" + # We assume svg_valid_count == expected_accepted_events, since all compared events meet metric's threshold + self.assertEqual( + expected_accepted_events, + svg_valid_count, + msg=f"Expected SVG events: {expected_accepted_events}, found: {svg_valid_count}.", + ) print( f"CSV valid count: {csv_valid_count}. SVG valid count: {svg_valid_count}" ) @@ -9849,7 +10035,6 @@ def setup_environment(): TestQNN.oss_repo = args.oss_repo TestQNN.shared_buffer = args.shared_buffer TestQNN.enable_x86_64 = args.enable_x86_64 - TestQNN.dump_intermediate_outputs = args.dump_intermediate_outputs TestQNN.compile_only = args.compile_only TestQNN.pre_gen_pte = args.pre_gen_pte TestQNN.llama_artifacts = args.llama_artifacts diff --git a/backends/qualcomm/tests/utils.py b/backends/qualcomm/tests/utils.py index c6f0c1b840f..d8802f74e68 100644 --- a/backends/qualcomm/tests/utils.py +++ b/backends/qualcomm/tests/utils.py @@ -8,6 +8,7 @@ import subprocess import tempfile import unittest +from pathlib import Path from typing import Callable, Dict, List, Optional, OrderedDict, Tuple import numpy as np @@ -15,6 +16,10 @@ import torchao from executorch import exir from executorch.backends.qualcomm.builders.node_visitor import dq_ops + +from executorch.backends.qualcomm.debugger.qcom_numerical_comparator_sample import ( + QcomCosineSimilarityComparator, +) from executorch.backends.qualcomm.debugger.qnn_intermediate_debugger import ( QNNIntermediateDebugger, ) @@ -48,7 +53,6 @@ from executorch.examples.qualcomm.utils import make_output_dir from executorch.exir.backend.compile_spec_schema import CompileSpec -from executorch.exir.backend.utils import get_delegates from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass from executorch.exir.passes.memory_planning_pass import MemoryPlanningPass @@ -122,8 +126,27 @@ def generate_context_binary( shell=True, executable="/bin/bash", capture_output=True, + text=True, + env=_qnn_subprocess_env(qnn_sdk, target), ) - assert os.path.isfile(f"{artifact_dir}/model_ctx.bin"), print(result.stderr) + assert os.path.isfile(f"{artifact_dir}/model_ctx.bin"), ( + f"Failed to generate context binary at {artifact_dir}/model_ctx.bin. " + f"returncode={result.returncode}\n" + f"stdout:\n{result.stdout}\n" + f"stderr:\n{result.stderr}" + ) + + +def _qnn_subprocess_env(qnn_sdk: str, target: str) -> Dict[str, str]: + """Return an env dict with LD_LIBRARY_PATH set so QNN SDK prebuilt + binaries (e.g. qnn-context-binary-utility) can resolve their bundled + libc++.so.1, which is shipped at $QNN_SDK_ROOT/lib//. + """ + env = os.environ.copy() + qnn_lib_dir = f"{qnn_sdk}/lib/{target}" + existing = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = f"{qnn_lib_dir}:{existing}" if existing else qnn_lib_dir + return env def validate_context_binary(ctx_bin: bytes): @@ -149,8 +172,15 @@ def validate_context_binary(ctx_bin: bytes): shell=True, executable="/bin/bash", capture_output=True, + text=True, + env=_qnn_subprocess_env(qnn_sdk, target), + ) + assert os.path.isfile(f"{tmp_dir}/ctx.json"), ( + f"qnn-context-binary-utility failed to produce ctx.json. " + f"returncode={result.returncode}\n" + f"stdout:\n{result.stdout}\n" + f"stderr:\n{result.stderr}" ) - assert os.path.isfile(f"{tmp_dir}/ctx.json"), print(result.stderr) class TestQNN(unittest.TestCase): @@ -186,11 +216,11 @@ class TestQNN(unittest.TestCase): compile_only: bool = False pre_gen_pte: str = "" llama_artifacts: str = "" - dump_intermediate_outputs: bool = False inference_speed: float = 0.0 inference_speed_output_path = "outputs/inference_speed.txt" static_llm_eval_method = "" direct_build_folder: str = "" + dsp_heap_profile_filename = "htp_heap_usage.txt" @classmethod def setUpClass(cls): @@ -321,7 +351,6 @@ def verify_output( # noqa: C901 executorch_prog: ExecutorchProgram | ExecutorchProgramManager, etrecord_path: str = "etrecord.bin", expected_profile_events: int = -1, - expected_intermediate_events: int = -1, method_index: int = 0, input_encodings: Tuple = (), output_encodings: Tuple = (), @@ -332,6 +361,7 @@ def verify_output( # noqa: C901 save_inference_speed: bool = False, expected_compared_events: int = -1, qnn_intermediate_debugger: QNNIntermediateDebugger = None, + save_heap_result: bool = False, ): with tempfile.TemporaryDirectory() as tmp_dir: ( @@ -385,30 +415,51 @@ def validate_profile(): len(inspector.to_dataframe().index) >= expected_profile_events ) - def validate_intermediate_tensor(): - inspector = Inspector( - etdump_path=etdump_path, debug_buffer_path=debug_output_path + def validate_heap_profile(): + file_path = f"{tmp_dir}/{self.dsp_heap_profile_filename}" + self.assertTrue( + Path(file_path).exists(), f"File not found: {file_path}" ) - node_tensor_map = qnn_intermediate_debugger._match_tensors( - inspector=inspector, keep_qnn_layout=False + with open(file_path, "r") as f: + values = [ + int(line.split(",")[1].strip()) for line in f if line.strip() + ] + self.assertEqual(len(values), 2, f"Expected 2 entries, got {values}") + before, after = values + difference = after - before + + print(f"before_context_created: {before} bytes") + print(f"after_context_freed: {after} bytes") + print(f"difference: {difference:.2f} bytes") + + self.assertGreaterEqual( + after, + before, + "after_context_freed should be >= before_context_created", ) - self.assertEqual( - len(node_tensor_map), - expected_compared_events, - msg=f"Unexpected number of compared events, expecting {expected_compared_events}, but has {len(node_tensor_map)} events.", + + def validate_intermediate_tensor(): + qnn_intermediate_debugger.setup_inspector( + etdump_path=etdump_path, + debug_buffer_path=debug_output_path, ) - # Compare accuracy for each layer - for _, value in node_tensor_map.items(): - self._assert_outputs_equal( - value[0].to(torch.float32), value[1].to(torch.float32) + cos_comparator = qnn_intermediate_debugger.create_comparator( + QcomCosineSimilarityComparator + ) + numeric_results = ( + qnn_intermediate_debugger.inspector.calculate_numeric_gap( + distance=cos_comparator, + reference_graph=qnn_intermediate_debugger.reference_graph_name, ) - for event_block in inspector.event_blocks: - if event_block.name == "Execute": - self.assertEqual( - len(event_block.events), - expected_intermediate_events, - msg=f"Unexpected number of intermediate events, expecting {expected_intermediate_events}, but has {len(event_block.events)} events.", - ) + ) + numeric_results = numeric_results.set_index("runtime_debug_handle") + assert ( + len(numeric_results) == expected_compared_events + ), f"Unexpected number of compared events, expecting {expected_compared_events}, but has {len(numeric_results)} events." + for _, row in numeric_results.iterrows(): + assert cos_comparator.is_valid_score( + row.gap[0] + ), f"Node {row.aot_ops} is failing {cos_comparator.metric_name()} test, {row.gap[0]} is lower than {cos_comparator.threshold}." processed_inputs = list(sample_inputs) for i, enc in enumerate(input_encodings): @@ -477,7 +528,7 @@ def validate_intermediate_tensor(): "--method_index", str(method_index), ] - if self.dump_intermediate_outputs: + if expected_compared_events != -1: cmd.append("--dump_intermediate_outputs") cmd += extra_cmds.split() @@ -528,7 +579,7 @@ def validate_intermediate_tensor(): if expected_profile_events != -1: validate_profile() - if expected_intermediate_events != -1: + if expected_compared_events != -1: validate_intermediate_tensor() if save_inference_speed: @@ -545,7 +596,7 @@ def validate_intermediate_tensor(): device=self.device, host=self.host, soc_model=self.soc_model, - dump_intermediate_outputs=expected_intermediate_events != -1, + dump_intermediate_outputs=expected_compared_events != -1, direct_build_folder=self.direct_build_folder, ) @@ -576,6 +627,11 @@ def validate_intermediate_tensor(): adb.extra_cmds += ( f" --performance_output_path {self.inference_speed_output_path}" ) + + if save_heap_result: + adb.extra_cmds += ( + f" --heap_profiling_path {self.dsp_heap_profile_filename}" + ) adb.execute(custom_runner_cmd=f"rm -rf {adb.output_folder}") adb.execute(method_index=method_index, output_callback=output_callback) adb.pull(host_output_path=tmp_dir, callback=post_process) @@ -584,7 +640,7 @@ def validate_intermediate_tensor(): if expected_profile_events != -1: adb.pull_etdump(etdump_path, callback=validate_profile) - if expected_intermediate_events != -1: + if expected_compared_events != -1: adb.pull_debug_output( etdump_path, debug_output_path, @@ -595,6 +651,12 @@ def validate_intermediate_tensor(): f"{tmp_dir}/{self.inference_speed_output_path}", "r" ) as f: self.inference_speed = float(f.read()) + if save_heap_result: + adb.pull_heap_output( + f"{adb.workspace}/{self.dsp_heap_profile_filename}", + f"{tmp_dir}/{self.dsp_heap_profile_filename}", + callback=validate_heap_profile, + ) def lower_module_and_test_output( self, @@ -602,7 +664,6 @@ def lower_module_and_test_output( sample_inputs: Tuple[torch.Tensor], expected_partitions: int = 1, expected_profile_events: int = -1, - expected_intermediate_events: int = -1, expected_compared_events: int = -1, assert_output_equal: bool = True, passes_job: Optional[OrderedDict] = None, @@ -623,28 +684,9 @@ def lower_module_and_test_output( skip_node_id_set=skip_node_id_set, skip_node_op_set=skip_node_op_set, skip_mutable_buffer=skip_mutable_buffer, - generate_etrecord=self.profile_level != 0 - or expected_intermediate_events != -1, + generate_etrecord=self.profile_level != 0 or expected_compared_events != -1, ) - qnn_intermediate_debugger = None - if expected_intermediate_events != -1: - lowered_module_nodes = get_delegates( - delegated_program.exported_program().graph - ) - assert len(lowered_module_nodes) == 1, "Length not correct" - - lowered_module_node = lowered_module_nodes[0] - lower_module = getattr( - delegated_program.exported_program().graph_module, - lowered_module_node.name, - ) - edge_module = lower_module.original_module.module() - - qnn_intermediate_debugger = QNNIntermediateDebugger() - qnn_intermediate_debugger.set_edge_module(edge_module=edge_module) - qnn_intermediate_debugger.intermediate_output_module(*sample_inputs) - # Don't allocate if shared_buffer enabled or using direct_mode allocate_io = not (self.shared_buffer or self.direct_build_folder) exec_prog = delegated_program.to_executorch( @@ -675,11 +717,24 @@ def lower_module_and_test_output( etrecord_path = "etrecord.bin" if self.profile_level: exec_prog.get_etrecord().save(etrecord_path) + + qnn_intermediate_debugger = None + if expected_compared_events != -1: + etrecord = exec_prog.get_etrecord() + qnn_intermediate_debugger = QNNIntermediateDebugger(sample_inputs) + qnn_intermediate_debugger.set_etrecord_file_path(etrecord_path) + edge_ep = etrecord.graph_map[qnn_intermediate_debugger.reference_graph_name] + qnn_intermediate_debugger.set_edge_ep(edge_ep=edge_ep) + etrecord.update_representative_inputs( + qnn_intermediate_debugger.sample_input + ) + etrecord.save(etrecord_path) + # Check numerics if ( assert_output_equal or expected_profile_events != -1 - or expected_intermediate_events != -1 + or expected_compared_events != -1 ): self.verify_output( module=module, @@ -687,7 +742,6 @@ def lower_module_and_test_output( executorch_prog=exec_prog, etrecord_path=etrecord_path, expected_profile_events=expected_profile_events, - expected_intermediate_events=expected_intermediate_events, extra_cmds=extra_cmds, output_callback=output_callback, save_inference_speed=save_inference_speed, diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index bf11230065c..84c6ded0741 100644 --- a/backends/qualcomm/utils/utils.py +++ b/backends/qualcomm/utils/utils.py @@ -400,6 +400,21 @@ def ensure_graph_specific_dict(value, graph_names): return value return {graph_name: value for graph_name in graph_names} + # Ensure if user is using intermediate debugger, user only lower 1 method. + # This restriction is caused by conflict handle_id among graphs. + # This could be resolved with generating random debug_id(e.g., uuid). + for compiler_spec in ( + compiler_specs.values() + if isinstance(compiler_specs, Dict) + else [compiler_specs] + ): + option = generate_qnn_executorch_option(compiler_spec) + obj_options = flatbuffer_to_option(option) + if obj_options.dump_intermediate_outputs and isinstance(module, Dict): + assert ( + len(module) == 1 + ), "Intermediate Tensor Dump does not support multi-methods." + if not isinstance(module, dict): module = {"forward": module} diff --git a/backends/transforms/decompose_sdpa.py b/backends/transforms/decompose_sdpa.py index 13acaa32f11..a0f3bc3ab65 100644 --- a/backends/transforms/decompose_sdpa.py +++ b/backends/transforms/decompose_sdpa.py @@ -14,6 +14,7 @@ from executorch.exir.pass_base import ExportPass, PassResult from torch._decomp import get_decompositions from torch.fx.experimental.proxy_tensor import make_fx +from torch.fx.node import Argument class DecomposeScaledDotProductAttention(ExportPass): @@ -22,6 +23,11 @@ class DecomposeScaledDotProductAttention(ExportPass): """ _passes_required_after: Set[Type[ExportPass]] = set() + _SDPA_OPTIONAL_ARGS = ( + ("attn_mask", None), + ("dropout_p", 0.0), + ("is_causal", False), + ) def __init__(self, allow_non_fake_inputs: bool = True) -> None: super().__init__() @@ -42,6 +48,105 @@ def call( graph_module.recompile() return PassResult(graph_module, True) + @staticmethod + def _extract_arg_value(arg: object) -> object: + if isinstance(arg, torch.fx.Node): + if "val" not in arg.meta: + raise RuntimeError(f"Missing meta['val'] for SDPA arg node: {arg.name}") + return arg.meta["val"] + return arg + + @classmethod + def _canonicalize_sdpa_call( + cls, node: torch.fx.Node + ) -> tuple[tuple[object, ...], object, object]: + input_args = list(node.args) + input_kwargs = dict(node.kwargs) + + canonical_args = list(input_args[:3]) + for arg_index, (arg_name, default) in enumerate( + cls._SDPA_OPTIONAL_ARGS, start=3 + ): + if len(input_args) > arg_index: + canonical_args.append(input_args[arg_index]) + else: + canonical_args.append(input_kwargs.pop(arg_name, default)) + + raw_scale = input_kwargs.pop("scale", None) + canonical_args.append(raw_scale) + scale = cls._extract_arg_value(raw_scale) + enable_gqa = cls._extract_arg_value(input_kwargs.pop("enable_gqa", False)) + if input_kwargs: + raise RuntimeError( + "Unsupported kwargs for scaled_dot_product_attention: " + f"{', '.join(sorted(input_kwargs.keys()))}" + ) + + return tuple(canonical_args), scale, enable_gqa + + @staticmethod + def _copy_decomposed_graph( + graph: torch.fx.Graph, + node: torch.fx.Node, + decomposed_module: torch.fx.GraphModule, + canonical_inputs: tuple[object, ...], + scale: object, + ) -> None: + decomposed_node_to_subgraph_node: dict[torch.fx.Node, Argument] = {} + last_decomposed_node: torch.fx.Node | None = None + placeholder_nodes = [ + decomposed_node + for decomposed_node in decomposed_module.graph.nodes + if decomposed_node.op == "placeholder" + ] + if len(placeholder_nodes) != len(canonical_inputs): + raise RuntimeError( + "Unexpected placeholder count when decomposing " + "scaled_dot_product_attention" + ) + for decomposed_node, arg in zip(placeholder_nodes, canonical_inputs): + decomposed_node_to_subgraph_node[decomposed_node] = arg + + for decomposed_node in decomposed_module.graph.nodes: + if decomposed_node.op == "output": + output_arg = decomposed_node.args[0] + if not isinstance(output_arg, torch.fx.Node): + raise RuntimeError( + "Unexpected non-node output when decomposing " + "scaled_dot_product_attention" + ) + last_decomposed_node = output_arg + + for decomposed_node in decomposed_module.graph.nodes: + decomposed_node.meta["nn_module_stack"] = node.meta.get("nn_module_stack") + if decomposed_node.op == "placeholder": + continue + + if decomposed_node.op == "output" and last_decomposed_node is not None: + for user in node.users.copy(): + user.replace_input_with( + node, + decomposed_node_to_subgraph_node[last_decomposed_node], + ) + continue + + if scale is not None and decomposed_node.target in [ + torch.ops.aten.mul.Scalar + ]: + new_args = list(decomposed_node.args) + new_args[1] = math.sqrt(scale) + decomposed_node.args = tuple(new_args) + + subgraph_node = graph.node_copy( + decomposed_node, + arg_transform=lambda x: decomposed_node_to_subgraph_node[x], + ) + subgraph_node.meta["nn_module_stack"] = node.meta.get("nn_module_stack") + subgraph_node.meta["source_fn_stack"] = [ + (subgraph_node, subgraph_node.target) + ] + decomposed_node_to_subgraph_node[decomposed_node] = subgraph_node + def _decompose_sdpa_node( self, graph_module: torch.fx.GraphModule, @@ -49,12 +154,43 @@ def _decompose_sdpa_node( allow_non_fake_inputs: bool, ) -> None: graph = graph_module.graph - input_tensors = (input_node.meta["val"] for input_node in node.all_input_nodes) - scale = node.kwargs.get("scale", None) + + canonical_inputs, scale, enable_gqa = self._canonicalize_sdpa_call(node) + input_tensors = tuple(self._extract_arg_value(arg) for arg in canonical_inputs) + + def _sdpa_with_gqa( + q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None + ): + # Shapes: (B, H, T, D) + Hq = q.shape[1] + Hk = k.shape[1] + if Hq != Hk: + # LLaMA-style GQA: tile K and V heads to match Q + if not enable_gqa: + raise ValueError( + "SDPA head mismatch requires enable_gqa=True: " + f"Hq={Hq}, Hk={Hk}" + ) + if Hq % Hk != 0: + raise ValueError(f"GQA mismatch: Hq={Hq}, Hk={Hk}") + r = Hq // Hk + B, _, Tk, D = k.shape + k = k.unsqueeze(2).expand(B, Hk, r, Tk, D).reshape(B, Hq, Tk, D) + v = v.unsqueeze(2).expand(B, Hk, r, Tk, D).reshape(B, Hq, Tk, D) + return torch.ops.aten.scaled_dot_product_attention.default( + q, + k, + v, + attn_mask, + dropout_p, + is_causal, + scale=scale, + enable_gqa=enable_gqa, + ) # refer to pytorch/test/test_decomp.py decomposed_module = make_fx( - node.target, + _sdpa_with_gqa, decomposition_table=get_decompositions( # pyre-fixme[6] [ torch.ops.aten._scaled_dot_product_flash_attention_for_cpu.default, @@ -65,56 +201,8 @@ def _decompose_sdpa_node( )(*input_tensors) with graph.inserting_before(node): - name_to_input_tensor_map = {} - for i, arg in enumerate(node.args): - name_to_input_tensor_map[f"arg{i}_1"] = arg - - decomposed_node_to_subgraph_node: dict[torch.fx.Node, torch.fx.Node] = {} - last_decomposed_node = None - # Create a mapping from input nodes in decomposed module to original nodes. - # In decomposed module, there are only input tensors for placeholder op. - for decomposed_node in decomposed_module.graph.nodes: - if decomposed_node.op == "placeholder": - decomposed_node_to_subgraph_node[decomposed_node] = ( - name_to_input_tensor_map[decomposed_node.name] - ) - - if decomposed_node.op == "output": - last_decomposed_node = decomposed_node.args[0] - - # Copy node from decompose graph module - for decomposed_node in decomposed_module.graph.nodes: - node.meta["nn_module_stack"] = decomposed_node.meta.get( - "nn_module_stack" - ) - if decomposed_node.op == "placeholder": - continue - - if decomposed_node.op == "output" and last_decomposed_node is not None: - for user in node.users.copy(): - user.replace_input_with( - node, - decomposed_node_to_subgraph_node[last_decomposed_node], - ) - continue - - if scale is not None and decomposed_node.target in [ - torch.ops.aten.mul.Scalar - ]: - new_args = list(decomposed_node.args) - # Based on the implementation of _scaled_dot_product_attention_math, - # the scale is applied to q and k before matmul. - # refer to pytorch/aten/src/ATen/native/transformers/attention.cpp#L873 - new_args[1] = math.sqrt(scale) - decomposed_node.args = tuple(new_args) - - subgraph_node = graph.node_copy( - decomposed_node, - arg_transform=lambda x: decomposed_node_to_subgraph_node[x], - ) - subgraph_node.meta["source_fn_stack"] = [ - (subgraph_node, subgraph_node.target) - ] - decomposed_node_to_subgraph_node[decomposed_node] = subgraph_node + self._copy_decomposed_graph( + graph, node, decomposed_module, canonical_inputs, scale + ) graph.erase_node(node) diff --git a/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py b/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py index b8d6c75a174..f350120e7eb 100644 --- a/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py +++ b/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py @@ -20,7 +20,8 @@ class FuseCascadedTransposeOrPermuteOps(RemoveOrReplacePassInterface): """ Fuse a chain of transpose and permute ops into a single permute or a no-op. - Handles branches and chains permutes. + Handles branches and chains of permutes, including permute-view-permute + patterns where a squeeze/unsqueeze view sits between two permutes. """ transpose_or_permute_target = { @@ -28,20 +29,31 @@ class FuseCascadedTransposeOrPermuteOps(RemoveOrReplacePassInterface): exir_ops.edge.aten.permute_copy.default, } + _VIEW_OPS = { + exir_ops.edge.aten.view_copy.default, + exir_ops.edge.aten.view.default, + } + @property def targets(self) -> list[EdgeOpOverload]: return list(self.transpose_or_permute_target) def maybe_remove_or_replace(self, node: Node) -> bool: - # Fuse with the parent node if it's also a permute or a transpose. Since the - # pass interface traverses all ops in order the pass will properly fuse a chain - # of permutes. parent_node = get_arg(node, "input", Node) - if parent_node.target not in self.transpose_or_permute_target: - return False - input_of_parent = get_arg(parent_node, "input", Node) - # Compute combined effect of permutes. + # Case 1: Direct permute/transpose → permute/transpose + if parent_node.target in self.transpose_or_permute_target: + return self._fuse_direct(node, parent_node) + + # Case 2: permute → view_copy(squeeze/unsqueeze) → permute + if parent_node.target in self._VIEW_OPS: + return self._fuse_across_view(node, parent_node) + + return False + + def _fuse_direct(self, node: Node, parent_node: Node) -> bool: + """Fuse two adjacent permute/transpose ops.""" + input_of_parent = get_arg(parent_node, "input", Node) dims = list(range(node.meta["val"].ndim)) if parent_node.target == exir_ops.edge.aten.transpose_copy.int: @@ -54,7 +66,6 @@ def maybe_remove_or_replace(self, node: Node) -> bool: else: dims = get_permuted_dims(node, dims) - # If combined effect is identity replace the node with input. if dims == sorted(dims): node.replace_all_uses_with(input_of_parent) else: @@ -67,3 +78,104 @@ def maybe_remove_or_replace(self, node: Node) -> bool: node.replace_all_uses_with(new_permute) return True + + def _apply_view_to_dims( + self, dims: list[int], view_in_shape, view_out_shape + ) -> list[int] | None: + """Apply a squeeze or unsqueeze view to dimension mapping. + + Returns the updated dims, or None if the view cannot be mapped. + """ + if len(view_out_shape) == len(view_in_shape) + 1: + # unsqueeze: insert a new dim + index = self._find_extra_one(view_out_shape, view_in_shape) + if index == -1: + return None + dims = [x + 1 if x >= index else x for x in dims] + dims.insert(index, -1) # -1 marks the inserted dim + elif len(view_in_shape) == len(view_out_shape) + 1: + # squeeze: remove a dim + index = self._find_extra_one(view_in_shape, view_out_shape) + if index == -1: + return None + dims = list(dims) + del dims[index] + return dims + + def _fuse_across_view(self, node: Node, view_node: Node) -> bool: # noqa: C901 + """Fuse permute -> view(squeeze/unsqueeze) -> permute into a view_copy.""" + # view_node must have exactly one user (this permute node) + if len(view_node.users) != 1: + return False + # view_node's parent must be a permute/transpose + view_input = get_arg(view_node, "input", Node) + if view_input.target not in self.transpose_or_permute_target: + return False + # The view must be a squeeze or unsqueeze (rank differs by 1) + view_in_shape = view_input.meta["val"].shape + view_out_shape = view_node.meta["val"].shape + if abs(len(view_in_shape) - len(view_out_shape)) != 1: + return False + + # Get the input before the first permute + input_of_first_permute = get_arg(view_input, "input", Node) + + # Compute the combined effect on the original input dimensions + # Start with identity dims for the original input + original_ndim = input_of_first_permute.meta["val"].ndim + dims = list(range(original_ndim)) + + # Apply first permute + if view_input.target == exir_ops.edge.aten.transpose_copy.int: + dims = get_transposed_dims(view_input, dims) + else: + dims = get_permuted_dims(view_input, dims) + + # Apply the view (squeeze/unsqueeze) + dims = self._apply_view_to_dims(dims, view_in_shape, view_out_shape) + if dims is None: + return False + + # Apply second permute (node) + if node.target == exir_ops.edge.aten.transpose_copy.int: + node_dims = list(range(len(dims))) + node_dims = get_transposed_dims(node, node_dims) + dims = [dims[d] for d in node_dims] + elif node.target == exir_ops.edge.aten.permute_copy.default: + perm = get_arg(node, "dims") + dims = [dims[d] for d in perm] + else: + raise ValueError(f"Unexpected target: {node.target}") + + # Check if the combined effect (ignoring -1 inserted dims) is identity + real_dims = [d for d in dims if d != -1] + + if real_dims == sorted(real_dims): + # Combined permutations are identity — replace with view_copy + # (the only remaining effect is the squeeze/unsqueeze reshape) + output_shape = node.meta["val"].shape + if output_shape == input_of_first_permute.meta["val"].shape: + # Total no-op: replace with input + node.replace_all_uses_with(input_of_first_permute) + else: + with node.graph.inserting_before(node): + new_view = node.graph.call_function( + exir_ops.edge.aten.view_copy.default, + args=(input_of_first_permute, list(output_shape)), + ) + new_view.meta = node.meta + node.replace_all_uses_with(new_view) + return True + + return False + + @staticmethod + def _find_extra_one(longer: list[int], shorter: list[int]) -> int: + if len(longer) != len(shorter) + 1: + return -1 + for i in range(len(shorter)): + if longer[i] != shorter[i]: + if longer[i] == 1 and shorter[i:] == longer[i + 1 :]: + return i + return -1 + return len(shorter) if longer[-1] == 1 else -1 diff --git a/backends/transforms/quantize_fused_convbn_bias_pass.py b/backends/transforms/quantize_fused_convbn_bias_pass.py index a4f654f8f93..1f24b3447c4 100644 --- a/backends/transforms/quantize_fused_convbn_bias_pass.py +++ b/backends/transforms/quantize_fused_convbn_bias_pass.py @@ -238,6 +238,8 @@ def _quantize_fused_conv_bias( ) quant_min = -(2**31) + 1 if use_symmetric_quantization else -(2**31) + if len(weight_dequant.args) < 2: + continue if isinstance(weight_dequant.args[1], torch.fx.node.Node): weight_scale = get_weight_scale_tensor(weight_dequant.args[1]) bias_scale = input_dequant.args[1] * weight_scale diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 8a4170e207e..eec6bdc4e08 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -12,8 +12,8 @@ import torch import torch.fx +from executorch.backends.transforms.permute_pass_utils import get_arg, set_arg from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.pass_base import ExportPass, PassResult @@ -39,23 +39,143 @@ class Subgraph: constant_edges_in: set[tuple[torch.fx.Node, torch.fx.Node]] = field( default_factory=set ) + # Per-node expected end permutation (may differ from end_permute + # when the subgraph contains rank-changing views). + node_end_permute: dict[torch.fx.Node, list[int]] = field(default_factory=dict) + # Per-node expected start permutation for upstream traversal. + node_start_permute: dict[torch.fx.Node, list[int]] = field(default_factory=dict) - # Ops explicitly listed as permutable. This includes non-pointwise ops - # that need special dimension-argument handling (cat, mean, sum, slice) - # and quantize/dequantize ops not tagged as pointwise in ATen. - # In addition to this set, any op tagged with torch.Tag.pointwise is - # automatically considered permutable (see is_node_permutable). - permutable_ops: set[EdgeOpOverload] = { - exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, - # Ops that require special handling of dimension arguments. - exir_ops.edge.aten.cat.default, - exir_ops.edge.aten.mean.dim, - exir_ops.edge.aten.sum.dim_IntList, - exir_ops.edge.aten.slice_copy.Tensor, - } - - def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + def __init__(self, extra_permutable_ops: set | None = None) -> None: + super().__init__() + self._permutable_ops = { + exir_ops.edge.aten.add.Tensor, + exir_ops.edge.aten.mul.Tensor, + exir_ops.edge.aten.sub.Tensor, + exir_ops.edge.aten.hardtanh.default, + exir_ops.edge.aten.clamp.default, + exir_ops.edge.aten.cat.default, + exir_ops.edge.aten.mean.dim, + exir_ops.edge.aten.sum.dim_IntList, + exir_ops.edge.aten.slice_copy.Tensor, + } + try: + self._permutable_ops.add( + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default + ) + self._permutable_ops.add( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default + ) + except AttributeError: + pass + if extra_permutable_ops: + self._permutable_ops |= extra_permutable_ops + self._sq_unsq_cache: dict[torch.fx.Node, bool] = {} + + _VIEW_OPS = ( + exir_ops.edge.aten.view_copy.default, + exir_ops.edge.aten.view.default, + ) + + _UNSQUEEZE_OPS = (exir_ops.edge.aten.unsqueeze_copy.default,) + + _SQUEEZE_OPS = (exir_ops.edge.aten.squeeze_copy.dim,) + + @staticmethod + def _find_extra_one(longer: list[int], shorter: list[int]) -> int: + """If longer has exactly one more element of value 1, return its index. Else -1.""" + if len(longer) != len(shorter) + 1: + return -1 + for i in range(len(shorter)): + if longer[i] != shorter[i]: + if longer[i] == 1 and shorter[i:] == longer[i + 1 :]: + return i + return -1 + return len(shorter) if longer[-1] == 1 else -1 + + def _is_squeeze_unsqueeze_view(self, node: torch.fx.Node) -> bool: + """Check if a node is a squeeze, unsqueeze, or view_copy that only + adds or removes a single dim of size 1.""" + if node in self._sq_unsq_cache: + return self._sq_unsq_cache[node] + result = self._check_squeeze_unsqueeze_view(node) + self._sq_unsq_cache[node] = result + return result + + def _check_squeeze_unsqueeze_view(self, node: torch.fx.Node) -> bool: + if node.target in self._UNSQUEEZE_OPS or node.target in self._SQUEEZE_OPS: + return True + if node.target not in self._VIEW_OPS: + return False + inp = node.args[0] + assert isinstance(inp, torch.fx.Node) + in_shape = inp.meta["val"].shape + out_shape = node.meta["val"].shape + if len(out_shape) == len(in_shape) + 1: + return self._find_extra_one(out_shape, in_shape) != -1 + if len(in_shape) == len(out_shape) + 1: + return self._find_extra_one(in_shape, out_shape) != -1 + return False + + def _adapt_permute_across_view( + self, permute: list[int], node: torch.fx.Node + ) -> list[int] | None: + """Adjust a permutation across a squeeze/unsqueeze boundary. + + Adapts from input-rank to output-rank space (downstream direction). + Returns the adjusted permutation, or None if not possible. + """ + # Handle explicit unsqueeze_copy(dim) + if node.target in self._UNSQUEEZE_OPS: + dim = cast(int, node.args[1]) + rank = len(permute) + index = dim if dim >= 0 else dim + rank + 1 + new_perm = [x + 1 if x >= index else x for x in permute] + new_perm.insert(index, index) + return new_perm + + # Handle explicit squeeze_copy(dim) + if node.target in self._SQUEEZE_OPS: + dim = cast(int, node.args[1]) + rank = len(permute) + index = dim if dim >= 0 else dim + rank + # index is a POSITION in the tensor; the permutation VALUE at + # that position is the logical dim being removed. + squeezed_value = permute[index] + new_perm = [ + x - 1 if x > squeezed_value else x + for x in permute + if x != squeezed_value + ] + return new_perm + + # Handle view_copy (squeeze/unsqueeze-like reshape) + inp = node.args[0] + assert isinstance(inp, torch.fx.Node) + in_shape = inp.meta["val"].shape + out_shape = node.meta["val"].shape + + if len(out_shape) == len(in_shape) + 1: + # unsqueeze: insert identity mapping at the new dim + index = self._find_extra_one(out_shape, in_shape) + new_perm = [x + 1 if x >= index else x for x in permute] + new_perm.insert(index, index) + return new_perm + elif len(in_shape) == len(out_shape) + 1: + # squeeze via view_copy: find the squeezed dim and remove it + index = self._find_extra_one(in_shape, out_shape) + # index is a POSITION in in_shape; the permutation VALUE at + # that position is the logical dim being removed. + squeezed_value = permute[index] + new_perm = [ + x - 1 if x > squeezed_value else x + for x in permute + if x != squeezed_value + ] + return new_perm + return None + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 + self._sq_unsq_cache.clear() subgraphs_found: list[RemovePermutesAroundElementwiseOps.Subgraph] = [] processed_nodes: set[torch.fx.Node] = set() for node in graph_module.graph.find_nodes( @@ -67,18 +187,56 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # Expected end permutation for the subgraph. end_permute = [start_permute.index(i) for i in range(len(start_permute))] + # Try direct users first (same-rank matching) for user in node.users: - if user.target not in self.permutable_ops and not self._is_pointwise( - user.target - ): + if not self.is_node_permutable(user): continue - # Create a separate subgraph for each user since there may be cases - # where only a portion of the users are permutable. subgraph = self.Subgraph(start_permute, end_permute) if self.visit(user, subgraph, processed_nodes): subgraphs_found.append(subgraph) - for node in subgraph.nodes: - processed_nodes.add(node) + for n in subgraph.nodes: + processed_nodes.add(n) + + # Also try: permute → view(squeeze/unsqueeze) → chain → ... + # If the permute's sole user is a squeeze/unsqueeze view, + # adapt the permutation across the view and search for a + # matching end permute at the new rank. + users = list(node.users.keys()) + if ( + len(users) == 1 + and self._is_squeeze_unsqueeze_view(users[0]) + and node not in processed_nodes + ): + view_node = users[0] + adapted_start = self._adapt_permute_across_view( + start_permute, view_node + ) + if adapted_start is not None: + adapted_end = [ + adapted_start.index(i) for i in range(len(adapted_start)) + ] + for view_user in view_node.users: + if not self.is_node_permutable(view_user): + continue + subgraph = self.Subgraph(adapted_start, adapted_end) + # Include the view in the subgraph + subgraph.nodes.add(view_node) + subgraph.node_end_permute[view_node] = adapted_end + # Use the ORIGINAL start_permute for the view node + # so update_view_copy can remap its shape correctly + subgraph.node_start_permute[view_node] = start_permute + # The start permute feeds into the view + subgraph.edges_in.add((node, view_node)) + if self.visit( + view_user, + subgraph, + processed_nodes, + adapted_end, + adapted_start, + ): + subgraphs_found.append(subgraph) + for n in subgraph.nodes: + processed_nodes.add(n) modified = False for subgraph in subgraphs_found: @@ -97,40 +255,86 @@ def visit( # noqa: C901 node: torch.fx.Node, subgraph: Subgraph, processed_nodes: set[torch.fx.Node], + current_end_permute: list[int] | None = None, + current_start_permute: list[int] | None = None, ) -> bool: + if current_end_permute is None: + current_end_permute = subgraph.end_permute + if current_start_permute is None: + current_start_permute = subgraph.start_permute + if node in subgraph.nodes: return True if node in processed_nodes or not self.is_node_permutable(node): return False subgraph.nodes.add(node) + subgraph.node_end_permute[node] = current_end_permute + subgraph.node_start_permute[node] = current_start_permute + + # If this is a squeeze/unsqueeze view, adapt permutations for + # traversal across the rank change boundary. + downstream_end = current_end_permute + downstream_start = current_start_permute + if self._is_squeeze_unsqueeze_view(node): + # Adapt start permute for downstream (input-rank → output-rank) + adapted_start = self._adapt_permute_across_view(current_start_permute, node) + if adapted_start is None: + return False + downstream_start = adapted_start + + # Derive end permute as the inverse of adapted start to ensure + # consistency. Computing start and end independently via + # _adapt_permute_across_view can produce mismatched results for + # squeeze views because the formula differs for "forward" vs + # "inverse" permutations. + downstream_end = [adapted_start.index(i) for i in range(len(adapted_start))] # Traverse downstream: for user in node.users: - # Output should either go to a matching permute or another permutable op. if user.target == exir_ops.edge.aten.permute_copy.default: - if self.get_permutation(user) != subgraph.end_permute: + user_perm = self.get_permutation(user) + if user_perm == downstream_end: + subgraph.edges_out.add((node, user)) + else: + # Check if permute → view(squeeze/unsqueeze) forms an + # end boundary at a different rank. + user_users = list(user.users.keys()) + if len(user_users) == 1 and self._is_squeeze_unsqueeze_view( + user_users[0] + ): + view_after = user_users[0] + # Adapt the start permute across the view and derive + # the expected end permute as its inverse. + adapted_start_after = self._adapt_permute_across_view( + downstream_start, view_after + ) + if adapted_start_after is not None: + adapted = [ + adapted_start_after.index(i) + for i in range(len(adapted_start_after)) + ] + if user_perm == adapted: + # Include both the permute and the view as end edges + subgraph.edges_out.add((node, user)) + # Mark the view for inclusion so it gets preserved + continue return False - subgraph.edges_out.add((node, user)) elif user.op == "output": - # Graph output requires the data in its original layout. - # Removing permutes here would silently change the output - # format, so treat this as an invalid subgraph boundary. return False - elif not self.visit(user, subgraph, processed_nodes): + elif not self.visit( + user, subgraph, processed_nodes, downstream_end, downstream_start + ): return False # Traverse upstream: for inp in node.all_input_nodes: - # Input should either come from a matching permute or another permutable op. if inp.target == exir_ops.edge.aten.permute_copy.default: - if self.get_permutation(inp) != subgraph.start_permute: + if self.get_permutation(inp) != current_start_permute: return False subgraph.edges_in.add((inp, node)) elif self._is_constant(inp): - # Only accept the constant if we can insert a compensating - # permute or view. Otherwise reject the subgraph. const_rank = self._get_node_rank(inp) - permute_rank = len(subgraph.end_permute) + permute_rank = len(current_end_permute) if const_rank is None: return False if const_rank > permute_rank: @@ -138,7 +342,13 @@ def visit( # noqa: C901 if const_rank < permute_rank and inp.meta.get("val") is None: return False subgraph.constant_edges_in.add((inp, node)) - elif not self.visit(inp, subgraph, processed_nodes): + elif not self.visit( + inp, + subgraph, + processed_nodes, + current_end_permute, + current_start_permute, + ): return False return True @@ -146,21 +356,27 @@ def visit( # noqa: C901 def _is_constant(self, node: torch.fx.Node) -> bool: """Check if a node's value is available at compile time. Only considers direct constants (get_attr, parameter/buffer/constant - placeholders) — does not recurse into call_function chains to avoid - stack overflow on deep graphs.""" + placeholders, full ops producing scalar constants) — does not recurse + into call_function chains to avoid stack overflow on deep graphs.""" if node.op == "get_attr": return True if node.op == "placeholder": target = str(node.target) return target.startswith(("b_", "p_", "c_")) + # full.default creates scalar constants (e.g. epsilon in LayerNorm) + if ( + node.op == "call_function" + and node.target == exir_ops.edge.aten.full.default + ): + return True return False def _get_node_rank(self, node: torch.fx.Node) -> int | None: """Return the tensor rank of a node's output, or None if unknown.""" val = node.meta.get("val") - if val is not None and hasattr(val, "shape"): - return len(val.shape) - return None + if val is None: + return None + return len(val.shape) @staticmethod def _is_pointwise(target) -> bool: @@ -171,27 +387,51 @@ def _is_pointwise(target) -> bool: return False def is_node_permutable(self, node: torch.fx.Node) -> bool: - if node.target in self.permutable_ops: - # Special-case validation for dim-based ops. + if node.target in self._permutable_ops: if node.target in ( exir_ops.edge.aten.mean.dim, exir_ops.edge.aten.sum.dim_IntList, ): - # keepdim should be True. - if len(node.args) >= 3: - if not node.args[2]: - return False - elif "keepdim" in node.kwargs: - if not node.kwargs["keepdim"]: - return False - else: - # Default keepdim is False. + if not get_arg(node, "keepdim", bool): return False return True - # Accept any op tagged as pointwise in ATen (elementwise). + if self._is_squeeze_unsqueeze_view(node): + return True return self._is_pointwise(node.target) - def permute_subgraph(self, subgraph: Subgraph) -> None: + def permute_subgraph(self, subgraph: Subgraph) -> None: # noqa: C901 + # Handle dimension related node arguments FIRST, before + # bypassing permutes (which changes node inputs/metadata). + for node in subgraph.nodes: + node_start_perm = subgraph.node_start_permute.get( + node, subgraph.start_permute + ) + if node.target == exir_ops.edge.aten.cat.default: + self.update_cat(node, node_start_perm) + elif node.target in ( + exir_ops.edge.aten.mean.dim, + exir_ops.edge.aten.sum.dim_IntList, + ): + self.update_mean_dim(node, node_start_perm) + elif node.target == exir_ops.edge.aten.slice_copy.Tensor: + self.update_slice_copy(node, node_start_perm) + elif node.target in self._VIEW_OPS: + self.update_view_copy(node, node_start_perm) + elif node.target in self._UNSQUEEZE_OPS: + # unsqueeze dim is in output space (rank + 1) + dim = cast(int, node.args[1]) + rank = len(node_start_perm) + index = dim if dim >= 0 else dim + rank + 1 + if index < rank: + node.update_arg(1, node_start_perm[index]) + else: + # Inserting at or beyond existing dims — position unchanged + node.update_arg(1, index) + elif node.target in self._SQUEEZE_OPS: + # squeeze dim is in input space (rank) + dim = get_arg(node, "dim", int) + set_arg(node, "dim", node_start_perm[dim]) + # Skip incoming permutes. for inp, out in subgraph.edges_in: assert inp.target == exir_ops.edge.aten.permute_copy.default @@ -201,38 +441,30 @@ def permute_subgraph(self, subgraph: Subgraph) -> None: out.replace_input_with(inp, cast(torch.fx.Node, inp.kwargs["input"])) # Insert compensating permute on constant inputs. - # Since the subgraph's start permutes are being removed, the subgraph - # will operate in the un-permuted (original) layout. Constants that - # were in the permuted layout need end_permute (the inverse of - # start_permute) to convert back to the original layout. for const_node, user_node in subgraph.constant_edges_in: graph = const_node.graph const_rank = self._get_node_rank(const_node) - permute_rank = len(subgraph.end_permute) + # Use the node-specific end_permute for the correct rank + node_end_perm = subgraph.node_end_permute.get( + user_node, subgraph.end_permute + ) + permute_rank = len(node_end_perm) with graph.inserting_after(const_node): if const_rank is not None and const_rank == permute_rank: new_node = graph.create_node( "call_function", exir_ops.edge.aten.permute_copy.default, - args=(const_node, subgraph.end_permute), + args=(const_node, node_end_perm), ) elif ( const_rank is not None and const_rank < permute_rank and const_node.meta.get("val") is not None ): - # Rank mismatch (e.g. rank-1 bias with rank-4 permute). - # The constant is broadcastable and its shape is smaller - # than the permute rank, so we can't apply the permute - # directly. Instead, use view_copy to rearrange the - # shape according to the end_permute restricted to - # the trailing dimensions. original_shape = list(const_node.meta["val"].shape) - # Pad shape to match permute rank for reordering padded = [1] * (permute_rank - const_rank) + original_shape - target_shape = [padded[d] for d in subgraph.end_permute] - # Strip leading 1s back to original rank + target_shape = [padded[d] for d in node_end_perm] target_shape = target_shape[permute_rank - const_rank :] new_node = graph.create_node( "call_function", @@ -240,7 +472,6 @@ def permute_subgraph(self, subgraph: Subgraph) -> None: args=(const_node, target_shape), ) else: - # Cannot determine rank or handle this case; skip. continue user_node.replace_input_with(const_node, new_node) @@ -249,43 +480,52 @@ def permute_subgraph(self, subgraph: Subgraph) -> None: assert out.target == exir_ops.edge.aten.permute_copy.default out.replace_all_uses_with(inp) - # Handle dimension related node arguments. - for node in subgraph.nodes: - if node.target == exir_ops.edge.aten.cat.default: - self.update_cat(node, subgraph.start_permute) - elif node.target in ( - exir_ops.edge.aten.mean.dim, - exir_ops.edge.aten.sum.dim_IntList, - ): - self.update_mean_dim(node, subgraph.start_permute) - elif node.target == exir_ops.edge.aten.slice_copy.Tensor: - self.update_slice_copy(node, subgraph.start_permute) - def update_cat(self, node: torch.fx.Node, start_permute: list[int]) -> None: - if len(node.args) >= 2: - node.update_arg(1, start_permute[cast(int, node.args[1])]) - elif "dim" in node.kwargs: - node.update_kwarg("dim", start_permute[cast(int, node.kwargs["dim"])]) - else: - # Default cat dim is 0. - node.update_kwarg("dim", start_permute[0]) + dim = get_arg(node, "dim", int) + set_arg(node, "dim", start_permute[dim]) def update_mean_dim(self, node: torch.fx.Node, start_permute: list[int]) -> None: - if len(node.args) >= 2: - node.update_arg( - 1, [start_permute[dim] for dim in cast(list[int], node.args[1])] - ) - else: - node.update_kwarg( - "dim", - [start_permute[dim] for dim in cast(list[int], node.kwargs["dim"])], - ) + dims = get_arg(node, "dim") + set_arg(node, "dim", [start_permute[d] for d in cast(list[int], dims)]) def update_slice_copy(self, node: torch.fx.Node, start_permute: list[int]) -> None: - if len(node.args) >= 2: - node.update_arg(1, start_permute[cast(int, node.args[1])]) - else: - node.update_kwarg("dim", start_permute[cast(int, node.kwargs["dim"])]) + dim = get_arg(node, "dim", int) + set_arg(node, "dim", start_permute[dim]) + + def update_view_copy(self, node: torch.fx.Node, start_permute: list[int]) -> None: + """Adjust view_copy shape arg after permute removal. + + After removing the start permute, the view's input is in the original + (un-permuted) layout. Recompute the view's target shape accordingly. + """ + inp = node.args[0] + assert isinstance(inp, torch.fx.Node) + + in_shape = inp.meta["val"].shape + out_shape = node.meta["val"].shape + + # Compute un-permuted input shape + inverse_permute = [start_permute.index(i) for i in range(len(start_permute))] + unpermuted_in = [in_shape[inverse_permute[i]] for i in range(len(in_shape))] + + if len(out_shape) == len(in_shape) + 1: + # unsqueeze: find the inserted dim in the permuted output, + # then determine where it goes in the un-permuted layout + index = self._find_extra_one(out_shape, in_shape) + if index != -1: + new_shape = list(unpermuted_in) + new_shape.insert(index, 1) + node.update_arg(1, new_shape) + elif len(in_shape) == len(out_shape) + 1: + # squeeze: find the removed dim in the permuted input, + # map it to the un-permuted position, and remove it + index = self._find_extra_one(in_shape, out_shape) + if index != -1: + # Map the squeezed dim from permuted to un-permuted space + unpermuted_index = start_permute[index] + new_shape = list(unpermuted_in) + del new_shape[unpermuted_index] + node.update_arg(1, new_shape) def get_permutation(self, permute_node: torch.fx.Node) -> list[int] | None: assert permute_node.target == exir_ops.edge.aten.permute_copy.default diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index 5c3343469ce..8c3603e293d 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -330,6 +330,7 @@ def define_common_targets(): "//executorch/backends/...", ], deps = [ + ":permute_pass_utils", "//caffe2:torch", "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", @@ -376,6 +377,7 @@ def define_common_targets(): ":fuse_cascaded_transpose_or_permute_ops", ":fuse_cascaded_view_ops", ":postpone_permute_below_squeeze_view", + ":remove_permutes_around_elementwise_ops", ":replace_nop_transpose_or_permute_with_view", ], ) diff --git a/backends/transforms/test/test_decompose_sdpa.py b/backends/transforms/test/test_decompose_sdpa.py new file mode 100644 index 00000000000..9d79666919f --- /dev/null +++ b/backends/transforms/test/test_decompose_sdpa.py @@ -0,0 +1,96 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch +from executorch.backends.transforms.decompose_sdpa import ( + DecomposeScaledDotProductAttention, +) +from torch.export import export + + +class TestDecomposeScaledDotProductAttention(unittest.TestCase): + def test_decompose_sdpa_requires_enable_gqa_for_head_mismatch(self) -> None: + class Model(torch.nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, k, v, enable_gqa=True + ) + + q = torch.randn(1, 4, 3, 4) + k = torch.randn(1, 2, 3, 4) + v = torch.randn(1, 2, 3, 4) + + graph_module = export(Model().eval(), (q, k, v), strict=True).module() + for node in graph_module.graph.nodes: + if node.target == torch.ops.aten.scaled_dot_product_attention.default: + node.kwargs = {**node.kwargs, "enable_gqa": False} + break + + with self.assertRaisesRegex(ValueError, "enable_gqa=True"): + DecomposeScaledDotProductAttention()(graph_module) + + def test_decompose_sdpa_preserves_kwargs(self) -> None: + class Block(torch.nn.Module): + def forward(self, q, k, v, mask): + return torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=mask, + scale=0.25, + ) + + class Model(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.block = Block() + + def forward(self, q, k, v, mask): + return self.block(q, k, v, mask) + + q = torch.randn(1, 2, 3, 4) + k = torch.randn(1, 2, 3, 4) + v = torch.randn(1, 2, 3, 4) + mask = torch.tensor( + [[[[True, False, True], [True, True, False], [False, True, True]]]] + ) + + graph_module = export(Model().eval(), (q, k, v, mask), strict=True).module() + + before_output = graph_module(q, k, v, mask) + original_nn_module_stack = None + self.assertTrue( + any( + node.target == torch.ops.aten.scaled_dot_product_attention.default + for node in graph_module.graph.nodes + if node.op == "call_function" + ) + ) + for node in graph_module.graph.nodes: + if node.op == "call_function" and ( + node.target == torch.ops.aten.scaled_dot_product_attention.default + ): + original_nn_module_stack = node.meta.get("nn_module_stack") + break + + self.assertIsNotNone(original_nn_module_stack) + + DecomposeScaledDotProductAttention()(graph_module) + + self.assertFalse( + any( + node.target == torch.ops.aten.scaled_dot_product_attention.default + for node in graph_module.graph.nodes + if node.op == "call_function" + ) + ) + for node in graph_module.graph.nodes: + if node.op == "call_function": + self.assertEqual( + node.meta.get("nn_module_stack"), original_nn_module_stack + ) + torch.testing.assert_close(graph_module(q, k, v, mask), before_output) diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index bb326f125bc..808a599f81f 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -19,6 +19,9 @@ from executorch.backends.transforms.postpone_permute_below_squeeze_view import ( PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView, ) +from executorch.backends.transforms.remove_permutes_around_elementwise_ops import ( + RemovePermutesAroundElementwiseOps, +) from executorch.backends.transforms.replace_nop_transpose_or_permute_with_view import ( ReplaceNopTransposeOrPermuteWithViewPass, ) @@ -119,15 +122,12 @@ def test_cascaded_permutes_multiple_users(self) -> None: permute1 = builder.call_operator( op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 3, 1]) ) - # permute2 reverses permute1 => identity permute2 = builder.call_operator( op=exir_ops.edge.aten.permute_copy.default, args=(permute1, [0, 3, 1, 2]) ) - # permute3: different permutation permute3 = builder.call_operator( op=exir_ops.edge.aten.permute_copy.default, args=(permute1, [0, 2, 1, 3]) ) - # permute4 -> permute5: chained permute4 = builder.call_operator( op=exir_ops.edge.aten.permute_copy.default, args=(permute1, [3, 2, 0, 1]) ) @@ -148,6 +148,168 @@ def test_cascaded_permutes_multiple_users(self) -> None: "FuseCascadedTransposeOrPermuteOps", ) + def test_permute_view_permute_fuse(self) -> None: + """permute_3D([0,2,1]) → view(unsqueeze) → permute_4D([0,2,3,1]) should + be replaced with a single view_copy (permutations cancel out).""" + builder = GraphBuilder() + x_data = torch.randn(1, 40, 18) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + v = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(p1, [1, 18, 1, 40]) + ) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(v, [0, 2, 3, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = FuseCascadedTransposeOrPermuteOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + gm = result.graph_module + + self.assertEqual(count_node(gm, exir_ops.edge.aten.permute_copy.default), 0) + self.assertEqual(count_node(gm, exir_ops.edge.aten.view_copy.default), 1) + validate_numerics( + gm_before, + gm, + [x_data], + "FuseCascadedAcrossView", + ) + + def test_permute_view_squeeze_permute_fuse(self) -> None: + """permute_4D → view(squeeze) → permute_3D should fuse when + the combined permutation is identity.""" + builder = GraphBuilder() + x_data = torch.randn(1, 1, 40, 18) + x = builder.placeholder("x", x_data) + # NHWC-like permute + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2]) + ) + # Squeeze dim 2 (size 1) + v = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(p1, [1, 18, 40]) + ) + # Inverse 3D permute + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(v, [0, 2, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = FuseCascadedTransposeOrPermuteOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + gm = result.graph_module + + self.assertEqual(count_node(gm, exir_ops.edge.aten.permute_copy.default), 0) + validate_numerics( + gm_before, + gm, + [x_data], + "FuseCascadedSqueezeView", + ) + + def test_transpose_view_permute_fuse(self) -> None: + """transpose → view(unsqueeze) → permute should fuse when combined + permutations cancel out.""" + builder = GraphBuilder() + x_data = torch.randn(1, 40, 18) + x = builder.placeholder("x", x_data) + t1 = builder.call_operator( + op=exir_ops.edge.aten.transpose_copy.int, args=(x, 1, 2) + ) + v = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(t1, [1, 18, 1, 40]) + ) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(v, [0, 2, 3, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = FuseCascadedTransposeOrPermuteOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + gm = result.graph_module + + self.assertEqual(count_node(gm, exir_ops.edge.aten.permute_copy.default), 0) + self.assertEqual(count_node(gm, exir_ops.edge.aten.transpose_copy.int), 0) + validate_numerics( + gm_before, + gm, + [x_data], + "FuseTransposeViewPermute", + ) + + def test_no_fuse_non_squeeze_view(self) -> None: + """permute → view (not squeeze/unsqueeze, changes shape) → permute + should NOT fuse.""" + builder = GraphBuilder() + x_data = torch.randn(1, 6, 8) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + # This view reshapes 8x6 → 4x12, NOT a squeeze/unsqueeze + v = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(p1, [1, 4, 12]) + ) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(v, [0, 2, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + + p = FuseCascadedTransposeOrPermuteOps() + result = cast(PassResult, p(original)) + # The view is not a squeeze/unsqueeze so cross-view fusion should not fire + self.assertFalse(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 2 + ) + + def test_no_fuse_non_cancelling_across_view(self) -> None: + """permute → view(unsqueeze) → permute where combined permutations + are NOT identity should NOT be fused away.""" + builder = GraphBuilder() + x_data = torch.randn(1, 40, 18) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + v = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(p1, [1, 18, 1, 40]) + ) + # This permute does NOT cancel with p1 + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(v, [0, 1, 3, 2]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = FuseCascadedTransposeOrPermuteOps() + result = cast(PassResult, p(original)) + # Should NOT have removed both permutes + self.assertFalse(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 2 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "FuseNonCancellingAcrossView", + ) + # ────────────────────────────────────────────────────────────────────── # Tests for FuseCascadedViewOps @@ -247,7 +409,6 @@ def test_permute3_view4_chains(self) -> None: self.assertEqual(count_node(gm, exir_ops.edge.aten.view_copy.default), 2) self.assertEqual(count_node(gm, exir_ops.edge.aten.permute_copy.default), 2) - # Verify order: views before permutes targets = get_compute_nodes(gm) view_indices = [ i @@ -347,7 +508,6 @@ def test_negative_not_squeeze_like(self) -> None: count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 2, ) - # Order unchanged: view, permute, view, permute targets = get_compute_nodes(result.graph_module) self.assertEqual(targets[0], exir_ops.edge.aten.view_copy.default) self.assertEqual(targets[1], exir_ops.edge.aten.permute_copy.default) @@ -440,3 +600,401 @@ def test_replace_nop_permute_3d(self) -> None: validate_numerics( gm_before, gm_after, [x], "ReplaceNopTransposeOrPermuteWithViewPass" ) + + +# ────────────────────────────────────────────────────────────────────── +# Tests for RemovePermutesAroundElementwiseOps cross-view handling +# ────────────────────────────────────────────────────────────────────── + + +class RemovePermutesAcrossViewTest(unittest.TestCase): + def test_permute_view_squeeze_elementwise_view_unsqueeze_permute(self) -> None: + """permute(3D) → view(unsqueeze) → mul(4D) → view(squeeze) → permute(3D) + should have both permutes removed.""" + builder = GraphBuilder() + x_data = torch.randn(1, 128, 16) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + v1 = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(p1, [1, 16, 1, 128]) + ) + mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(v1, v1)) + v2 = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(mul, [1, 16, 128]) + ) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(v2, [0, 2, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "RemovePermutesAcrossView", + ) + + def test_4d_permute_squeeze_clamp_3d_permute(self) -> None: + """Cascade detector conv→LN boundary: permute_4D([0,3,1,2]) → + view(squeeze) → hardtanh → permute_3D([0,2,1]). + The two permutes should cancel across the squeeze+clamp.""" + builder = GraphBuilder() + x_data = torch.randn(1, 1, 16, 128) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2]) + ) + v1 = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(p1, [1, 128, 16]) + ) + clamp = builder.call_operator( + op=exir_ops.edge.aten.hardtanh.default, args=(v1,) + ) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(clamp, [0, 2, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "4D_permute_squeeze_clamp_3D_permute", + ) + + def test_permute_unsqueeze_cat_mul_squeeze_permute(self) -> None: + """Complex interaction: permute(3D) → view(unsqueeze to 4D) → + cat(two branches) → mul → view(squeeze to 3D) → permute(3D). + Tests cat + mul interacting with view/squeeze/unsqueeze boundaries.""" + builder = GraphBuilder() + x_data = torch.randn(1, 128, 16) + y_data = torch.randn(1, 128, 16) + x = builder.placeholder("x", x_data) + y = builder.placeholder("y", y_data) + # Permute both inputs + px = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + py = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(y, [0, 2, 1]) + ) + # Unsqueeze via view to 4D + vx = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(px, [1, 16, 1, 128]) + ) + vy = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(py, [1, 16, 1, 128]) + ) + # Cat along dim 2 (the unsqueezed dim) + cat = builder.call_operator( + op=exir_ops.edge.aten.cat.default, args=([vx, vy], 2) + ) + # Mul with itself + mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(cat, cat)) + # Squeeze back via view + v_sq = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(mul, [1, 16, 256]) + ) + # End permute + p_end = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(v_sq, [0, 2, 1]) + ) + builder.output([p_end]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + # The cat changes output shape so squeeze view won't match the + # original unsqueeze pattern; the pass should not fire here. + self.assertFalse(result.modified) + validate_numerics( + gm_before, + result.graph_module, + [x_data, y_data], + "permute_unsqueeze_cat_mul_squeeze_permute", + ) + + def test_permute_view_add_sub_mul_view_permute(self) -> None: + """Chain of multiple elementwise ops between view boundaries: + permute(3D) → view(unsqueeze) → add → sub → mul → view(squeeze) → permute(3D). + All three elementwise ops should be handled.""" + builder = GraphBuilder() + x_data = torch.randn(1, 128, 16) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + # Unsqueeze via view + v1 = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(p1, [1, 16, 1, 128]) + ) + # Chain of elementwise ops + add = builder.call_operator(op=exir_ops.edge.aten.add.Tensor, args=(v1, v1)) + sub = builder.call_operator(op=exir_ops.edge.aten.sub.Tensor, args=(add, v1)) + mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(sub, sub)) + # Squeeze via view + v2 = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(mul, [1, 16, 128]) + ) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(v2, [0, 2, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "permute_view_add_sub_mul_view_permute", + ) + + def test_permute_squeeze_clamp_add_permute(self) -> None: + """4D permute → squeeze(view) → hardtanh → add(with self) → 3D permute. + Tests clamp + add interacting across a squeeze boundary.""" + builder = GraphBuilder() + x_data = torch.randn(1, 1, 16, 128) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2]) + ) + # Squeeze dim 2 (size 1) via view + v1 = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(p1, [1, 128, 16]) + ) + clamp = builder.call_operator( + op=exir_ops.edge.aten.hardtanh.default, args=(v1,) + ) + add = builder.call_operator( + op=exir_ops.edge.aten.add.Tensor, args=(clamp, clamp) + ) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(add, [0, 2, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "permute_squeeze_clamp_add_permute", + ) + + def test_no_fire_non_squeeze_view(self) -> None: + """permute → view (not a squeeze/unsqueeze, changes shape) → mul → permute. + The pass should NOT remove permutes when the view is not a simple + squeeze/unsqueeze.""" + builder = GraphBuilder() + x_data = torch.randn(1, 6, 8) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + # This view reshapes 8x6 → 4x12, which is NOT a squeeze/unsqueeze + v1 = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(p1, [1, 4, 12]) + ) + mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(v1, v1)) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(mul, [0, 2, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + # Should NOT have removed permutes (view is not squeeze/unsqueeze-like) + self.assertFalse(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 2 + ) + + def test_permute_unsqueeze_copy_mul_squeeze_copy_permute(self) -> None: + """permute(3D) → unsqueeze_copy(dim=2) → mul(4D) → squeeze_copy(dim=2) → permute(3D). + Tests the explicit unsqueeze_copy/squeeze_copy code paths in + _adapt_permute_across_view (distinct from view_copy).""" + builder = GraphBuilder() + x_data = torch.randn(1, 128, 16) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + unsq = builder.call_operator( + op=exir_ops.edge.aten.unsqueeze_copy.default, args=(p1, 2) + ) + mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(unsq, unsq)) + sq = builder.call_operator( + op=exir_ops.edge.aten.squeeze_copy.dim, args=(mul, 2) + ) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(sq, [0, 2, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "permute_unsqueeze_copy_mul_squeeze_copy_permute", + ) + + def test_4d_permute_squeeze_copy_clamp_3d_permute(self) -> None: + """4D permute([0,3,1,2]) → squeeze_copy(dim=2) → hardtanh → 3D permute([0,2,1]). + Tests the squeeze_copy code path at the start boundary (entering the + subgraph via squeeze_copy rather than view_copy).""" + builder = GraphBuilder() + x_data = torch.randn(1, 1, 16, 128) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 3, 1, 2]) + ) + sq = builder.call_operator(op=exir_ops.edge.aten.squeeze_copy.dim, args=(p1, 2)) + clamp = builder.call_operator( + op=exir_ops.edge.aten.hardtanh.default, args=(sq,) + ) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(clamp, [0, 2, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "4D_permute_squeeze_copy_clamp_3D_permute", + ) + + def test_4d_permute_squeeze_view_slice_mul_3d_permute(self) -> None: + """4D permute([2,0,1,3]) → view(squeeze dim 0) → slice → mul → permute([1,0,2]). + Regression test for the Transformer pattern where the squeezed dim + position (0) differs from its permutation value (perm[0]=2). + Without the fix, _adapt_permute_across_view confuses the position + with the value, causing the pass to create an invalid subgraph that + leads to a shape mismatch at runtime.""" + builder = GraphBuilder() + # Distinct dim sizes to expose mismatched slicing + x_data = torch.randn(10, 32, 1, 64) + x = builder.placeholder("x", x_data) + # Permute puts the size-1 dim (input dim 2) at position 0 + # [10, 32, 1, 64] -> [1, 10, 32, 64] + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [2, 0, 1, 3]) + ) + # Squeeze dim 0 (size 1) via view_copy: [10, 32, 64] + v1 = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(p1, [10, 32, 64]) + ) + # Slice dim 0, taking 3 elements from size 10 + sl = builder.call_operator( + op=exir_ops.edge.aten.slice_copy.Tensor, args=(v1, 0, 0, 3) + ) + # Elementwise op + mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(sl, sl)) + # End permute [1, 0, 2]: swap dims 0 and 1 + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(mul, [1, 0, 2]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + # With the fix, the adapted permutation becomes identity [0,1,2], + # so no matching end permute is found and the graph is unchanged. + # Before the fix, the wrong adapted permutation [1,0,2] would match + # the end permute and create an invalid subgraph, causing a crash. + result = cast(PassResult, p(original)) + self.assertFalse(result.modified) + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "4D_permute_squeeze_view_slice_mul_3D_permute", + ) + + def test_permute_unsqueeze_copy_neg_dim_mul_squeeze_copy_permute(self) -> None: + """permute(3D) → unsqueeze_copy(dim=-1) → mul(4D) → squeeze_copy(dim=3) → permute(3D). + Tests unsqueeze with negative dim (output-space rank+1 normalization) + and dim=rank edge case that would IndexError with incorrect handling.""" + builder = GraphBuilder() + x_data = torch.randn(1, 128, 16) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + unsq = builder.call_operator( + op=exir_ops.edge.aten.unsqueeze_copy.default, args=(p1, -1) + ) + mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(unsq, unsq)) + sq = builder.call_operator( + op=exir_ops.edge.aten.squeeze_copy.dim, args=(mul, 3) + ) + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(sq, [0, 2, 1]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "permute_unsqueeze_copy_neg_dim_mul_squeeze_copy_permute", + ) diff --git a/backends/xnnpack/CMakeLists.txt b/backends/xnnpack/CMakeLists.txt index 1b46c993b17..cd0d945a84f 100644 --- a/backends/xnnpack/CMakeLists.txt +++ b/backends/xnnpack/CMakeLists.txt @@ -169,6 +169,12 @@ install( EXPORT ExecuTorchTargets DESTINATION ${CMAKE_INSTALL_LIBDIR} ) +if(DEFINED EXECUTORCH_XNNPACK_LOG_LEVEL) + target_compile_definitions( + xnnpack-logging PUBLIC XNN_LOG_LEVEL=${EXECUTORCH_XNNPACK_LOG_LEVEL} + ) +endif() + if(BUILD_TESTING) add_subdirectory(test) endif() diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 103bdeb6b82..e04aa78579f 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -179,9 +180,11 @@ Result getConstantDataPtr( uint32_t buffer_idx, GraphPtr flatbuffer_graph, const uint8_t* constant_data_ptr, + uint64_t constant_data_size, const NamedDataMap* named_data_map, std::vector& freeable_buffers, - XNNWeightsCache* weights_cache) { + XNNWeightsCache* weights_cache, + bool use_weight_cache) { if (buffer_idx) { if (!constant_data_ptr) { // TODO(T172265611): Remove constant_buffer in flatbuffer path after BC @@ -219,10 +222,20 @@ Result getConstantDataPtr( "Null constant_data entry at buffer_idx %u", buffer_idx); uint64_t offset = constant_data_offset->offset(); + uint64_t entry_size = constant_data_offset->size(); bool has_named_key = flatbuffers::IsFieldPresent( constant_data_offset, fb_xnnpack::ConstantDataOffset::VT_NAMED_KEY); // If there is no tensor name if (!has_named_key) { + ET_CHECK_OR_RETURN_ERROR( + offset <= constant_data_size && + entry_size <= constant_data_size - offset, + InvalidProgram, + "ConstantDataOffset {offset=%" PRIu64 ", size=%" PRIu64 + "} out of bounds for constant_data region of size %" PRIu64, + offset, + entry_size, + constant_data_size); return constant_data_ptr + offset; } else { ET_CHECK_OR_RETURN_ERROR( @@ -230,30 +243,30 @@ Result getConstantDataPtr( InvalidProgram, "Named key is null"); const std::string& data_name = constant_data_offset->named_key()->str(); -#ifdef ENABLE_XNNPACK_WEIGHTS_CACHE - Result data_ptr = - weights_cache->load_unpacked_data(data_name); - if (!data_ptr.ok()) { - ET_LOG(Error, "Failed to load weights from cache"); - return data_ptr.error(); - } - return data_ptr.get(); -#else - Result buffer = - named_data_map->get_data(data_name.c_str()); - if (!buffer.ok()) { - ET_LOG( - Error, - "Failed to get constant data for key %s from named_data_map. Error code: %u", - data_name.c_str(), - static_cast(buffer.error())); - return buffer.error(); + if (use_weight_cache) { + Result data_ptr = + weights_cache->load_unpacked_data(data_name); + if (!data_ptr.ok()) { + ET_LOG(Error, "Failed to load weights from cache"); + return data_ptr.error(); + } + return data_ptr.get(); + } else { + Result buffer = + named_data_map->get_data(data_name.c_str()); + if (!buffer.ok()) { + ET_LOG( + Error, + "Failed to get constant data for key %s from named_data_map. Error code: %u", + data_name.c_str(), + static_cast(buffer.error())); + return buffer.error(); + } + const uint8_t* data_ptr = + static_cast(buffer.get().data()); + freeable_buffers.push_back(std::move(buffer.get())); + return data_ptr; } - const uint8_t* data_ptr = - static_cast(buffer.get().data()); - freeable_buffers.push_back(std::move(buffer.get())); - return data_ptr; -#endif } } } @@ -265,16 +278,20 @@ Result getConstantDataPtr( const fb_xnnpack::XNNTensorValue* tensor_value, GraphPtr flatbuffer_graph, const uint8_t* constant_data_ptr, + uint64_t constant_data_size, const NamedDataMap* named_data_map, std::vector& freeable_buffers, - XNNWeightsCache* weights_cache) { + XNNWeightsCache* weights_cache, + bool use_weight_cache) { return getConstantDataPtr( tensor_value->constant_buffer_idx(), flatbuffer_graph, constant_data_ptr, + constant_data_size, named_data_map, freeable_buffers, - weights_cache); + weights_cache, + use_weight_cache); } /** @@ -288,12 +305,14 @@ Error defineTensor( ValuePtr value, GraphPtr flatbuffer_graph, const uint8_t* constant_data_ptr, + uint64_t constant_data_size, std::vector& input_ids, std::vector& output_ids, CompileAllocator& allocator, const NamedDataMap* named_data_map, std::vector& freeable_buffers, - XNNWeightsCache* weights_cache) { + XNNWeightsCache* weights_cache, + bool use_weight_cache) { const fb_xnnpack::XNNTensorValue* tensor_value = nullptr; const fb_xnnpack::XNNQuantizedTensorValue* qtensor_value = nullptr; @@ -345,9 +364,11 @@ Error defineTensor( tensor_value, flatbuffer_graph, constant_data_ptr, + constant_data_size, named_data_map, freeable_buffers, - weights_cache); + weights_cache, + use_weight_cache); if (!buffer_result.ok()) { return buffer_result.error(); } @@ -500,9 +521,11 @@ Error defineTensor( qparams->scale_buffer_idx(), flatbuffer_graph, constant_data_ptr, + constant_data_size, named_data_map, freeable_buffers, - weights_cache); + weights_cache, + use_weight_cache); if (!scale_result.ok()) { return scale_result.error(); } @@ -546,9 +569,11 @@ Error defineTensor( qparams->scale_buffer_idx(), flatbuffer_graph, constant_data_ptr, + constant_data_size, named_data_map, freeable_buffers, - weights_cache); + weights_cache, + use_weight_cache); if (!scale_data_result.ok()) { return scale_data_result.error(); } @@ -1976,10 +2001,12 @@ ET_NODISCARD Error XNNCompiler::compileModel( XNNExecutor* executor, XNNWeightsCache* weights_cache, xnn_workspace_t workspace, - const NamedDataMap* named_data_map) { + const NamedDataMap* named_data_map, + bool use_weight_cache) { Result header = XNNHeader::Parse(buffer_pointer, num_bytes); const uint8_t* flatbuffer_data = nullptr; const uint8_t* constant_data = nullptr; + uint64_t constant_data_size = 0; size_t flatbuffer_size = 0; CompileAllocator compile_allocator; @@ -1990,6 +2017,7 @@ ET_NODISCARD Error XNNCompiler::compileModel( flatbuffer_size = header->flatbuffer_size; constant_data = reinterpret_cast(buffer_pointer) + header->constant_data_offset; + constant_data_size = header->constant_data_size; } else if (header.error() == Error::NotFound) { flatbuffer_data = reinterpret_cast(buffer_pointer); flatbuffer_size = num_bytes; @@ -2081,12 +2109,14 @@ ET_NODISCARD Error XNNCompiler::compileModel( value, flatbuffer_graph, constant_data, + constant_data_size, input_ids, output_ids, compile_allocator, named_data_map, unpacked_buffers, - weights_cache); + weights_cache, + use_weight_cache); if (err != Error::Ok) { return err; @@ -2108,19 +2138,16 @@ ET_NODISCARD Error XNNCompiler::compileModel( xnn_runtime_t runtime_ptr = nullptr; - // XNNWeightsCache if weights cache is not enabled, then XNNWeightsCache - // just manages the unpacked weights until the runtime is created. -#ifdef ENABLE_XNNPACK_WEIGHTS_CACHE - ET_CHECK_OR_RETURN_ERROR( - unpacked_buffers.size() == 0, - Internal, - "Weight Cache is enabled, which means unpacked buffers should be owned by the cache"); - xnn_weights_cache_t weights_cache_ptr = - weights_cache->get_num_unpacked_data() > 0 ? weights_cache->get() - : nullptr; -#else xnn_weights_cache_t weights_cache_ptr = nullptr; -#endif + if (use_weight_cache) { + ET_CHECK_OR_RETURN_ERROR( + unpacked_buffers.size() == 0, + Internal, + "Weight Cache is enabled, which means unpacked buffers should be owned by the cache"); + weights_cache_ptr = weights_cache->get_num_unpacked_data() > 0 + ? weights_cache->get() + : nullptr; + } // NOLINTBEGIN(facebook-hte-NullableDereference) - weights cache is allowed to // be null @@ -2139,25 +2166,25 @@ ET_NODISCARD Error XNNCompiler::compileModel( "XNN Runtime creation failed with code: %s", xnn_status_to_string(status)); -#ifdef ENABLE_XNNPACK_WEIGHTS_CACHE - auto packed_weights_names = weights_cache->finalize_for_runtime(); - ET_CHECK_OR_RETURN_ERROR( - packed_weights_names.ok(), - Internal, - "Failed to finalize weights cache after creating the xnn runtime") -#else - for (auto& buffer : unpacked_buffers) { - buffer.Free(); + std::vector packed_weights_names; + if (use_weight_cache) { + auto packed_weights_names_result = weights_cache->finalize_for_runtime(); + ET_CHECK_OR_RETURN_ERROR( + packed_weights_names_result.ok(), + Internal, + "Failed to finalize weights cache after creating the xnn runtime"); + packed_weights_names = std::move(packed_weights_names_result.get()); + } else { + for (auto& buffer : unpacked_buffers) { + buffer.Free(); + } } - Result> packed_weights_names = - std::vector(); -#endif err = executor->initialize( // NOLINT: runtime_ptr is non-null runtime_ptr, std::move(input_ids), std::move(output_ids), - std::move(packed_weights_names.get())); + std::move(packed_weights_names)); return err; }; diff --git a/backends/xnnpack/runtime/XNNCompiler.h b/backends/xnnpack/runtime/XNNCompiler.h index bcc87351d7d..639df0438cb 100644 --- a/backends/xnnpack/runtime/XNNCompiler.h +++ b/backends/xnnpack/runtime/XNNCompiler.h @@ -29,7 +29,8 @@ class XNNCompiler { XNNExecutor* executor, XNNWeightsCache* weights_cache, xnn_workspace_t workspace, - const NamedDataMap* named_data_map); + const NamedDataMap* named_data_map, + bool use_weight_cache); }; } // namespace delegate diff --git a/backends/xnnpack/runtime/XNNPACKBackend.cpp b/backends/xnnpack/runtime/XNNPACKBackend.cpp index 23a3f4c4b1f..c20fa985f46 100644 --- a/backends/xnnpack/runtime/XNNPACKBackend.cpp +++ b/backends/xnnpack/runtime/XNNPACKBackend.cpp @@ -91,8 +91,13 @@ class XnnpackBackend final auto workspace = workspace_result.get(); bool use_weight_cache = options_.resolve_weight_cache(context); + // Hold the lock for the entire init-compile-finalize sequence to prevent + // concurrent inits from resetting is_finalized_ or overwriting + // named_data_map_ while compileModel is using the shared weights cache. + std::unique_lock lock_weights_cache( + weights_cache_mutex_, std::defer_lock); if (use_weight_cache) { - const std::lock_guard lock_weight_cache(weights_cache_mutex_); + lock_weights_cache.lock(); weights_cache_->initialize_for_runtime( context.get_runtime_allocator(), named_data_map); } @@ -110,7 +115,8 @@ class XnnpackBackend final executor, weights_cache_.get(), workspace_ptr, - named_data_map); + named_data_map, + use_weight_cache); // This backend does not need its processed data after compiling the model. processed->Free(); diff --git a/conftest.py b/conftest.py index 6c9df86a1ce..19d777a74e0 100644 --- a/conftest.py +++ b/conftest.py @@ -1,5 +1,7 @@ import sys +import torch + collect_ignore_glob: list[str] = [] # Skip Apple tests on Windows. Note that some Core ML tests can run on Linux, as the AOT flow @@ -10,3 +12,6 @@ collect_ignore_glob += [ "backends/apple/**", ] + +# Seed the run +torch.manual_seed(42) diff --git a/devtools/inspector/_inspector.py b/devtools/inspector/_inspector.py index e9fbc4778f5..ad914878347 100644 --- a/devtools/inspector/_inspector.py +++ b/devtools/inspector/_inspector.py @@ -22,6 +22,7 @@ Mapping, Optional, Sequence, + Set, Tuple, TypeAlias, TypedDict, @@ -1025,6 +1026,7 @@ def __init__( Callable[[Union[int, str], Union[int, float]], Union[int, float]] ] = None, enable_module_hierarchy: bool = False, + reference_graph_name: str = EDGE_DIALECT_GRAPH_KEY, ) -> None: r""" Initialize an `Inspector` instance with the underlying `EventBlock`\ s populated with data from the provided ETDump path or binary, @@ -1040,6 +1042,7 @@ def __init__( delegate_metadata_parser: Optional function to parse delegate metadata from an Profiling Event. Expected signature of the function is (delegate_metadata_list: List[bytes]) -> Union[List[str], Dict[str, Any]]. delegate_time_scale_converter: Optional function to convert the time scale of delegate profiling data. If not given, use the conversion ratio of target_time_scale/source_time_scale. enable_module_hierarchy: Enable submodules in the operator graph. Defaults to False. + reference_graph_name: The reference graph used to consume ETRecord Returns: None @@ -1104,9 +1107,9 @@ def __init__( # Key str is method name; value is list of ProgramOutputs because of list of test cases self._reference_outputs: Dict[str, List[ProgramOutput]] = {} self._enable_module_hierarchy = enable_module_hierarchy - self._consume_etrecord() + self._consume_etrecord(reference_graph_name) - def _consume_etrecord(self) -> None: + def _consume_etrecord(self, reference_graph_name) -> None: """ If an ETRecord is provided, connect it to the EventBlocks and populate the Event metadata. @@ -1147,7 +1150,7 @@ def _consume_etrecord(self) -> None: enable_module_hierarchy=self._enable_module_hierarchy, ) debug_handle_to_op_node_map = create_debug_handle_to_op_node_mapping( - self.op_graph_dict[EDGE_DIALECT_GRAPH_KEY], + self.op_graph_dict[reference_graph_name], ) for event_block in self.event_blocks: for event in event_block.events: @@ -1337,6 +1340,29 @@ def _get_runtime_intermediate_outputs_and_op_names( Retrieve the runtime intermediate outputs(debug handles and intermediate values mappings) from the event blocks, along with the corresponding debug handles and op names mapping. """ + # Collect debug handles already covered by fine-grained inner delegated + # events (i.e. delegated events with op_types and debug_data populated), + # grouped per backend. When such per-op intermediate outputs exist for a + # delegated subgraph, the wrapping DELEGATE_CALL event is redundant: its + # debug_handles span every internal handle. + inner_delegated_handles_by_backend: Dict[Optional[str], Set[int]] = defaultdict( + set + ) + for event_block in self.event_blocks: + for event in event_block.events: + if ( + event.is_delegated_op + and event.op_types + and event.debug_data + and event.debug_handles is not None + ): + handles = event.debug_handles + if isinstance(handles, int): + handles = (handles,) + inner_delegated_handles_by_backend[ + event.delegate_backend_name + ].update(handles) + debug_handle_to_output = {} debug_handle_to_op_names = {} for event_block in self.event_blocks: @@ -1353,6 +1379,17 @@ def _get_runtime_intermediate_outputs_and_op_names( debug_handle = (debug_handle,) else: debug_handle = tuple(debug_handle) + # Skip a DELEGATE_CALL whose handles are already covered by + # fine-grained inner delegated events from the same backend + # (see comment above). + if event.name == "DELEGATE_CALL": + backend_inner = inner_delegated_handles_by_backend.get( + event.delegate_backend_name + ) + if backend_inner and not set(debug_handle).isdisjoint( + backend_inner + ): + continue current_entry = debug_handle_to_output.get( debug_handle, (-1, None, event.num_outputs) ) diff --git a/devtools/inspector/numerical_comparator/numerical_comparator_base.py b/devtools/inspector/numerical_comparator/numerical_comparator_base.py index c4f8a90f78f..c9d3629f3fd 100644 --- a/devtools/inspector/numerical_comparator/numerical_comparator_base.py +++ b/devtools/inspector/numerical_comparator/numerical_comparator_base.py @@ -246,6 +246,7 @@ def compare( continue rows.append( { + "aot_debug_handle": aot_debug_handle, "aot_ops": find_op_names( aot_debug_handle, aot_debug_handle_to_op_names ), @@ -253,6 +254,7 @@ def compare( "runtime_ops": find_op_names( runtime_debug_handle, runtime_debug_handle_to_op_names ), + "runtime_debug_handle": runtime_debug_handle, "runtime_intermediate_output": runtime_intermediate_output, "gap": self._compare_intermediate_outputs( aot_intermediate_output, runtime_intermediate_output diff --git a/devtools/inspector/tests/inspector_test.py b/devtools/inspector/tests/inspector_test.py index 4ee8b71760e..b33c5b37164 100644 --- a/devtools/inspector/tests/inspector_test.py +++ b/devtools/inspector/tests/inspector_test.py @@ -740,9 +740,11 @@ def test_calculate_numeric_gap(self): self.assertEqual(len(df), 2) cols = set(df.columns) expected_cols = { + "aot_debug_handle", "aot_ops", "aot_intermediate_output", "runtime_ops", + "runtime_debug_handle", "runtime_intermediate_output", "gap", "stacktraces", @@ -831,9 +833,11 @@ def test_calculate_numeric_gap_with_stacktraces(self): self.assertEqual(len(df), 2) cols = set(df.columns) expected_cols = { + "aot_debug_handle", "aot_ops", "aot_intermediate_output", "runtime_ops", + "runtime_debug_handle", "runtime_intermediate_output", "gap", "stacktraces", @@ -918,9 +922,11 @@ def element_compare(self, a, b): self.assertEqual(len(df), 2) cols = set(df.columns) expected_cols = { + "aot_debug_handle", "aot_ops", "aot_intermediate_output", "runtime_ops", + "runtime_debug_handle", "runtime_intermediate_output", "gap", "stacktraces", @@ -1066,9 +1072,11 @@ def preprocessing( self.assertEqual(len(df_mse), 2) cols = set(df_mse.columns) expected_cols = { + "aot_debug_handle", "aot_ops", "aot_intermediate_output", "runtime_ops", + "runtime_debug_handle", "runtime_intermediate_output", "gap", "stacktraces", diff --git a/docs/README.md b/docs/README.md index 845267b32f6..da9e6a6a5df 100644 --- a/docs/README.md +++ b/docs/README.md @@ -52,7 +52,7 @@ To build the documentation locally: Or a Conda environment: ```bash - conda create -yn executorch python=3.10.0 && conda activate executorch + conda create -yn executorch python=3.10 && conda activate executorch ``` 1. Install dependencies: diff --git a/docs/source/android-examples.md b/docs/source/android-examples.md index b56ed660f63..24af5750d56 100644 --- a/docs/source/android-examples.md +++ b/docs/source/android-examples.md @@ -2,9 +2,9 @@ - [Working with LLMs - Android Examples](https://github.com/meta-pytorch/executorch-examples/blob/main/llm/android/LlamaDemo/README.md) - ExecuTorch Llama Android Demo App - [Demo Apps](https://github.com/meta-pytorch/executorch-examples/tree/main/dl3/android/DeepLabV3Demo#executorch-android-demo-app) - DeepLab v3 model for image segmentation -- {doc}`backends/arm-vgf/tutorials/arm-vgf-tutorials` — Export a simple PyTorch model for the ExecuTorch VGF backend +- {doc}`backends/arm-vgf/tutorials/vgf-getting-started` — Export a simple PyTorch model for the ExecuTorch VGF backend ```{toctree} :hidden: -backends/arm-vgf/tutorials/arm-vgf-tutorials +backends/arm-vgf/tutorials/vgf-getting-started ``` diff --git a/docs/source/android-samsung-exynos.md b/docs/source/android-samsung-exynos.md index 4c5a470edca..d2c020ba0fd 100644 --- a/docs/source/android-samsung-exynos.md +++ b/docs/source/android-samsung-exynos.md @@ -1 +1,2 @@ -```{include} backends-samsung-exynos.md +```{include} backends/samsung/samsung-overview.md +``` diff --git a/docs/source/android-section.md b/docs/source/android-section.md index a5774352bc1..dedb7df22d7 100644 --- a/docs/source/android-section.md +++ b/docs/source/android-section.md @@ -16,8 +16,16 @@ Deploy ExecuTorch on Android devices with hardware acceleration support. - {doc}`android-examples` — Explore Android Examples & Demos +## API Reference + +- [Java API Reference (Javadoc)](https://pytorch.org/executorch/main/javadoc/index.html) — Full Java class and method reference + ```{toctree} +:maxdepth: 1 :hidden: + using-executorch-android android-backends android-examples +Java API Reference (Javadoc) +``` diff --git a/docs/source/api-life-cycle.md b/docs/source/api-life-cycle.md index 3ccaa4eddb1..fcf82094d73 100644 --- a/docs/source/api-life-cycle.md +++ b/docs/source/api-life-cycle.md @@ -80,21 +80,72 @@ communicate state to developers. Language + API state + Code Documentation + + All + + Stable + + + +No annotation is required. APIs are considered stable if they are not marked as +experimental or deprecated. + + + + +No warning is required. + + + + + All + + Deleted + + + +Remove the API from code after the deprecation period has passed. + + + + +Remove references to the deleted API from documentation. + + + Python + Deprecated + Use the executorch.exir._warnings.deprecated decorator. -

+ + + +Use .. warning:: in the docstring. Clearly point to the +replacement API when one exists. + + + + + Python + + Experimental + + + Use the executorch.exir._warnings.experimental decorator. @@ -102,25 +153,22 @@ decorator. -Use .. warning:: in the docstrings of deprecated and experimental -APIs. See +Use .. warning:: in the docstring. State that the API is +experimental and may change or be removed without notice. See example usage. - C++ + Deprecated + Use the ET_DEPRECATED annotation macro. See example usage. -

-

-Use the ET_EXPERIMENTAL annotation macro. - @@ -128,23 +176,33 @@ Start Doxygen comments with DEPRECATED: See example usage. -

-

+ + + + C++ + + Experimental + + + +Use the ET_EXPERIMENTAL annotation macro. + + + + Start Doxygen comments with EXPERIMENTAL:. + Java + Deprecated + Use java.lang.Deprecated. -

-

- -Use androidx.annotation.RequiresOptIn. -

@@ -152,6 +210,19 @@ Use androidx.annotation.RequiresOptIn. + + + 

/**
 * Warning: This API is experimental.
@@ -161,12 +232,11 @@ Use 
-

-

-__attribute__((deprecated("This API is experimental and may change without notice.")));
    
    
 

@@ -175,6 +245,18 @@ Use 
+   
+   
 

 

 /**
@@ -186,17 +268,28 @@ Use 
+   
+   
 

+/// - Warning: Deprecated. Use `newMethod()` instead.
+   
+  
+  
+   Swift
+   
+   Experimental
+   
+   
 

 @available(*, message: "This API is experimental")
    
    
 

-/// - Warning: Deprecated. Use `newMethod()` instead.
-

 /// - Warning: This API is experimental.
    
   
diff --git a/docs/source/backends-overview.md b/docs/source/backends-overview.md
index 3d54a150eac..ad3d3b507c1 100644
--- a/docs/source/backends-overview.md
+++ b/docs/source/backends-overview.md
@@ -65,4 +65,5 @@ backends/arm-vgf/arm-vgf-overview
 build-run-openvino
 backends/nxp/nxp-overview
 backends-cadence
-backends/samsung/samsung-overview
+Samsung Exynos Backend 
+```
diff --git a/docs/source/backends/arm-ethos-u/arm-ethos-u-overview.md b/docs/source/backends/arm-ethos-u/arm-ethos-u-overview.md
index faffedece35..34b52ff5c35 100644
--- a/docs/source/backends/arm-ethos-u/arm-ethos-u-overview.md
+++ b/docs/source/backends/arm-ethos-u/arm-ethos-u-overview.md
@@ -4,7 +4,7 @@ The Arm® Ethos™-U backend targets Edge/IoT-type AI use-cases by enabli
 [Arm® Ethos™-U55 NPU](https://www.arm.com/products/silicon-ip-cpu/ethos/ethos-u55), [Arm® Ethos™-U65 NPU](https://www.arm.com/products/silicon-ip-cpu/ethos/ethos-u65), and
 [Arm® Ethos™-U85 NPU](https://www.arm.com/products/silicon-ip-cpu/ethos/ethos-u85), leveraging [TOSA](https://www.mlplatform.org/tosa/) and the
 [ethos-u-vela](https://pypi.org/project/ethos-u-vela/) graph compiler. This document is a technical reference for using the Ethos-U backend, for a top level view with code examples
-please refer to the [Arm Ethos-U Backend Tutorial](tutorials/ethos-u-getting-started.md).
+please refer to the [Arm Ethos-U Backend Tutorial](tutorials/ethos-u-getting-started.md). 
 
 ## Features
 
@@ -27,7 +27,7 @@ For the AOT flow, compilation of a model to `.pte` format using the Ethos-U back
 - [TOSA Serialization Library](https://www.mlplatform.org/tosa/software.html) for serializing the Exir IR graph into TOSA IR.
 - [Ethos-U Vela graph compiler](https://pypi.org/project/ethos-u-vela/) for compiling TOSA flatbuffers into an Ethos-U command stream.
 
-And for building and running the example application available in `examples/arm/executor_runner/`:
+And for building and running the example application available in `examples/arm/executor_runner/` through the standalone CMake entry point:
 - [Arm GNU Toolchain](https://developer.arm.com/Tools%20and%20Software/GNU%20Toolchain) for cross compilation.
 - [Arm® Corstone™ SSE-300 FVP](https://developer.arm.com/documentation/100966/1128/Arm--Corstone-SSE-300-FVP) for testing on a Arm® Cortex®-M55+Ethos-U55 reference design.
 - [Arm® Corstone™ SSE-320 FVP](https://developer.arm.com/documentation/109760/0000/SSE-320-FVP) for testing on a Arm® Cortex®-M85+Ethos-U85 reference design.
@@ -111,7 +111,7 @@ For more information on quantization, see [Quantization](arm-ethos-u-quantizatio
 
 ## Runtime Integration
 
-An example runtime application is available in [examples/arm/executor_runner](https://github.com/pytorch/executorch/blob/main/examples/arm/executor_runner/), and the steps requried for building and deploying it on a FVP it is explained in the previously mentioned [Arm Ethos-U Backend Tutorial](tutorials/ethos-u-getting-started.md).
+An example runtime application is available in [examples/arm/executor_runner](https://github.com/pytorch/executorch/blob/main/examples/arm/executor_runner/), with a standalone CMake entry point in `examples/arm/executor_runner/standalone`. The steps required for building and deploying it on an FVP are explained in the previously mentioned [Arm Ethos-U Backend Tutorial](tutorials/ethos-u-getting-started.md). 
 The example application is recommended to use for testing basic functionality of your lowered models, as well as a starting point for developing runtime integrations for your own targets.
 For an in-depth explanation of the architecture of the executor_runner and the steps required for doing such an integration, please refer to [Ethos-U porting guide](https://github.com/pytorch/executorch/blob/main/examples/arm/ethos-u-porting-guide.md).
 
@@ -209,7 +209,7 @@ ExecuTorch for the Ethos-U backend, you automatically install the compiler conta
 
 **→{doc}`/backends/arm-ethos-u/arm-ethos-u-troubleshooting` — Troubleshooting and common issues.**
 
-**→{doc}`/backends/arm-ethos-u/tutorials/arm-ethos-u-tutorials` — Tutorials.**
+**→{doc}`/backends/arm-ethos-u/tutorials/ethos-u-getting-started` — Getting started tutorial.**
 
 **→{doc}`/backends/arm-ethos-u/U55_op_support` — Ethos-U55 supported operators.**
 
@@ -224,7 +224,7 @@ ExecuTorch for the Ethos-U backend, you automatically install the compiler conta
 arm-ethos-u-partitioner
 arm-ethos-u-quantization
 arm-ethos-u-troubleshooting
-tutorials/arm-ethos-u-tutorials
+tutorials/ethos-u-getting-started
 U55_op_support
 U85_op_support
 ```
diff --git a/docs/source/backends/arm-ethos-u/tutorials/arm-ethos-u-tutorials.md b/docs/source/backends/arm-ethos-u/tutorials/arm-ethos-u-tutorials.md
deleted file mode 100644
index 4b540f2179d..00000000000
--- a/docs/source/backends/arm-ethos-u/tutorials/arm-ethos-u-tutorials.md
+++ /dev/null
@@ -1,10 +0,0 @@
-# Arm Ethos-U Backend Tutorials
-
-**→{doc}`ethos-u-getting-started`**
-
-```{toctree}
-:maxdepth: 2
-:hidden:
-:caption: Tutorials
-
-ethos-u-getting-started
diff --git a/docs/source/backends/arm-ethos-u/tutorials/ethos-u-getting-started.md b/docs/source/backends/arm-ethos-u/tutorials/ethos-u-getting-started.md
index 841827cff9b..5fdb3530023 100644
--- a/docs/source/backends/arm-ethos-u/tutorials/ethos-u-getting-started.md
+++ b/docs/source/backends/arm-ethos-u/tutorials/ethos-u-getting-started.md
@@ -20,7 +20,7 @@ In this tutorial you will learn how to export a simple PyTorch model for the Exe
 ```{tip}
 If you are already familiar with this delegate, you may want to jump directly to the examples:
 * [Examples in the ExecuTorch repository](https://github.com/pytorch/executorch/tree/main/examples/arm)
-* [A commandline compiler for example models](https://github.com/pytorch/executorch/blob/main/backends/arm/scripts/aot_arm_compiler.py)
+* [A commandline compiler for quick tests and example models](https://github.com/pytorch/executorch/blob/main/backends/arm/scripts/aot_arm_compiler.py)
 ```
 
 This tutorial serves as an introduction to using ExecuTorch to deploy PyTorch models on Arm® Ethos™-U targets. It is based on `ethos_u_minimal_example.ipynb`, provided in Arm’s examples folder.
@@ -142,42 +142,36 @@ save_pte_program(executorch_program_manager, "ethos_u_minimal_example.pte")
 
 
 ```{tip}
-For a quick start, you can use the script `backends/arm/scripts/aot_arm_compiler.py` to produce the pte.
+For a quick test, you can use the script `backends/arm/scripts/aot_arm_compiler.py` to produce the pte.
 To produce a pte file equivalent to the one above, run
-`python -m backends.arm.scripts.aot_arm_compiler --model_name=add --delegate --quantize --output=ethos_u_minimal_example.pte`
+`python -m backends.arm.scripts.aot_arm_compiler --model_name=add --delegate --quantize --output=ethos_u_minimal_example.pte`.
+For production use, you should instead use the stable Python API shown above.
 ```
 
 ### Runtime:
 
-After the AOT compilation flow is done, the runtime can be cross compiled and linked to the produced `.pte`-file using the Arm cross-compilation toolchain. This is done in two steps:
+After the AOT compilation flow is done, the runtime can be cross compiled and linked to the produced `.pte`-file using the Arm cross-compilation toolchain. Configure the standalone Arm executor runner CMake project to pull in the ExecuTorch build graph, link the Ethos-U delegate, and generate kernel bindings for any non-delegated ops. This produces the `arm_executor_runner` program that will run on target.
 
-First, build and install the ExecuTorch libraries and EthosUDelegate:
 ```
 # In ExecuTorch top-level, with sourced setup_path.sh
-cmake -DCMAKE_BUILD_TYPE=Release --preset arm-baremetal -B cmake-out-arm .
-cmake --build cmake-out-arm --target install -j$(nproc)
-```
-Second, build and link the `arm_executor_runner` and generate kernel bindings for any non delegated ops. This is the actual program that will run on target.
-
-```
-# In ExecuTorch top-level, with sourced setup_path.sh
-cmake -DCMAKE_TOOLCHAIN_FILE=`pwd`/examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake \
+cmake -S examples/arm/executor_runner/standalone \
+      -B ethos_u_minimal_example \
+      -DEXECUTORCH_ROOT=$(pwd) \
+      -DCMAKE_TOOLCHAIN_FILE=$(pwd)/examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake \
       -DCMAKE_BUILD_TYPE=Release \
       -DET_PTE_FILE_PATH=ethos_u_minimal_example.pte \
       -DTARGET_CPU=cortex-m55 \
       -DETHOSU_TARGET_NPU_CONFIG=ethos-u55-128 \
       -DMEMORY_MODE=Shared_Sram \
-      -DSYSTEM_CONFIG=Ethos_U55_High_End_Embedded \
-      -Bethos_u_minimal_example \
-      examples/arm/executor_runner
+      -DSYSTEM_CONFIG=Ethos_U55_High_End_Embedded
 cmake --build ethos_u_minimal_example -j$(nproc) -- arm_executor_runner
 ```
 
 ```{tip}
-For a quick start, you can use the script `backends/arm/scripts/build_executor_runner.sh` to build the runner.
+For a quick start, you can use the script `backends/arm/scripts/build_executor_runner.sh` to configure and build the standalone runner.
 To build a runner equivalent to the one above, run
 `./backends/arm/scripts/build_executor_runner.sh --pte=ethos_u_minimal_example.pte`
-````
+```
 
 The block diagram below shows, at the high level, how the various build artifacts are generated and are linked together to generate the final bare-metal executable.
 
@@ -196,7 +190,6 @@ The example application is by default built with an input of ones, so the expect
 ## Takeaways
 
 In this tutorial you have learned how to use ExecuTorch to export a PyTorch model to an executable that can run on an embedded target, and then run that executable on simulated hardware.
-
 To learn more, check out the [ExecuTorch on Arm Practical Labs](https://github.com/arm-education/executorch_on_arm_labs) series. This series provides a structured entry-point to developing with ExecuTorch on Arm, across both CPU and Ethos-U NPU.
 
 For quick learning paths showcasing short tutorials:
diff --git a/docs/source/backends/arm-vgf/arm-vgf-overview.md b/docs/source/backends/arm-vgf/arm-vgf-overview.md
index dd3b00eb356..2f4523a1eb9 100644
--- a/docs/source/backends/arm-vgf/arm-vgf-overview.md
+++ b/docs/source/backends/arm-vgf/arm-vgf-overview.md
@@ -129,7 +129,7 @@ described in the rest of this guide but with a concrete end-to-end sample.
 
 **→{doc}`/backends/arm-vgf/arm-vgf-troubleshooting` — Debug common issues.**
 
-**→{doc}`/backends/arm-vgf/tutorials/arm-vgf-tutorials` — Tutorials.**
+**→{doc}`/backends/arm-vgf/tutorials/vgf-getting-started` — Getting started tutorial.**
 
 **→{doc}`/backends/arm-vgf/VGF_op_support` — VGF supported operators.**
 
@@ -142,6 +142,6 @@ described in the rest of this guide but with a concrete end-to-end sample.
 arm-vgf-partitioner
 arm-vgf-quantization
 arm-vgf-troubleshooting
-tutorials/arm-vgf-tutorials
+tutorials/vgf-getting-started
 VGF_op_support
 ```
diff --git a/docs/source/backends/arm-vgf/tutorials/arm-vgf-tutorials.md b/docs/source/backends/arm-vgf/tutorials/arm-vgf-tutorials.md
deleted file mode 100644
index ceb4304a814..00000000000
--- a/docs/source/backends/arm-vgf/tutorials/arm-vgf-tutorials.md
+++ /dev/null
@@ -1,10 +0,0 @@
-# Arm VGF Backend Tutorials
-
-**→{doc}`vgf-getting-started`**
-
-```{toctree}
-:maxdepth: 2
-:hidden:
-:caption: Tutorials
-
-vgf-getting-started
diff --git a/docs/source/backends/arm-vgf/tutorials/vgf-getting-started.md b/docs/source/backends/arm-vgf/tutorials/vgf-getting-started.md
index 9600f3b9d54..fcb77452ac3 100644
--- a/docs/source/backends/arm-vgf/tutorials/vgf-getting-started.md
+++ b/docs/source/backends/arm-vgf/tutorials/vgf-getting-started.md
@@ -26,7 +26,7 @@ You may encounter some rough edges and features which may be documented or plann
 ```{tip}
 If you are already familiar with this delegate, you may want to jump directly to the examples:
 * [Examples in the ExecuTorch repository](https://github.com/pytorch/executorch/tree/main/examples/arm)
-* [A commandline compiler for example models](https://github.com/pytorch/executorch/blob/main/backends/arm/scripts/aot_arm_compiler.py)
+* [A commandline compiler for quick tests and example models](https://github.com/pytorch/executorch/blob/main/backends/arm/scripts/aot_arm_compiler.py)
 ```
 
 This tutorial serves as an introduction to using ExecuTorch to deploy PyTorch models on VGF targets. The tutorial is based on `vgf_minimal_example.ipyb`, provided in Arm's example folder.
@@ -163,14 +163,15 @@ assert os.path.exists(pte_path), "Build failed; no .pte-file found"
 
 
 ```{tip}
-For a quick start, you can use the script `backends/arm/scripts/aot_arm_compiler.py` to produce the pte.
+For a quick test, you can use the script `backends/arm/scripts/aot_arm_compiler.py` to produce the pte.
 To produce a pte file equivalent to the one above, run
-`python -m backends.arm.scripts.aot_arm_compiler --model_name=add --delegate --quantize --output=simple_example.pte --target=vgf`
+`python -m backends.arm.scripts.aot_arm_compiler --model_name=add --delegate --quantize --output=simple_example.pte --target=vgf`.
+For production use, you should instead use the stable Python API shown above.
 ```
 
-### Runtime:
+## Runtime
 
-## Build executor runtime
+### Build executor runtime
 
 After the AOT compilation flow is done, we can build the executor runner target. For this tutorial, the default runner can be used. Build it with the following configuration:
 
@@ -200,7 +201,7 @@ The block diagram below demonstrates, at the high level, how the various build a
 ![](arm-delegate-runtime-build.svg)
 
 
-## Deploying and running on device
+### Deploying and running on device
 
 Since we are using the Vulkan emulation layer, we can run the executor runner with the VGF delegate on the host machine:
 
diff --git a/docs/source/backends/nxp/nxp-overview.md b/docs/source/backends/nxp/nxp-overview.md
index 6070f86e458..00b173eed04 100644
--- a/docs/source/backends/nxp/nxp-overview.md
+++ b/docs/source/backends/nxp/nxp-overview.md
@@ -24,10 +24,10 @@ Among currently supported machine learning models are:
 
 - [MCUXpresso IDE](https://www.nxp.com/design/design-center/software/development-software/mcuxpresso-software-and-tools-/mcuxpresso-integrated-development-environment-ide:MCUXpresso-IDE) or [MCUXpresso Visual Studio Code extension](https://www.nxp.com/design/design-center/software/development-software/mcuxpresso-software-and-tools-/mcuxpresso-for-visual-studio-code:MCUXPRESSO-VSC)
 - [MCUXpresso SDK 25.12](https://mcuxpresso.nxp.com/mcuxsdk/25.12.00/html/index.html)
-- eIQ Neutron SDK version 3.1.0, what you can download from eIQ PyPI:
+- eIQ Neutron SDK version 3.1.1, what you can download from eIQ PyPI:
 
 ```commandline
-$ pip install --index-url https://eiq.nxp.com/repository eiq-neutron-sdk==3.1.0
+$ pip install --index-url https://eiq.nxp.com/repository eiq-neutron-sdk==3.1.1
 ```
 
 Instead of manually installing requirements, except MCUXpresso IDE and SDK, you can use the setup script: 
diff --git a/docs/source/backends/xnnpack/xnnpack-quantization.md b/docs/source/backends/xnnpack/xnnpack-quantization.md
index e0180393f9e..74d8eafba72 100644
--- a/docs/source/backends/xnnpack/xnnpack-quantization.md
+++ b/docs/source/backends/xnnpack/xnnpack-quantization.md
@@ -61,7 +61,7 @@ See [PyTorch 2 Export Post Training Quantization](https://docs.pytorch.org/ao/ma
 
 The XNNPACK backend also supports quantizing models with the [torchao](https://github.com/pytorch/ao) quantize_ API.  This is most commonly used for LLMs, requiring more advanced quantization.  Since quantize_ is not backend aware, it is important to use a config that is compatible with CPU/XNNPACK:
 
-* Quantize embeedings with `IntxWeightOnlyConfig` (with weight_dtype torch.int2, torch.int4, or torch.int8, using PerGroup or PerAxis granularity)
+* Quantize embeddings with `IntxWeightOnlyConfig` (with weight_dtype torch.int2, torch.int4, or torch.int8, using PerGroup or PerAxis granularity)
 * Quantize linear layers with 4 bit weight and 8bit dynamic activation, use `Int8DynamicActivationIntxWeightConfig` (with weight_dtype=torch.int4, using PerGroup or PerAxis granularity)
 
 Below is a simple example, but a more detailed tutorial including accuracy evaluation on popular LLM benchmarks can be found in the [torchao documentation](https://docs.pytorch.org/ao/main/serving.html#mobile-deployment-with-executorch).
@@ -79,13 +79,13 @@ embedding_config = IntxWeightOnlyConfig(
     weight_dtype=torch.int8,
     granularity=PerAxis(0),
 )
-qunatize_(
+quantize_(
     eager_model,
     lambda m, fqn: isinstance(m, torch.nn.Embedding),
 )
 
 
-# Quatize linear layers with 8-bit dynamic activations and 4-bit weights
+# Quantize linear layers with 8-bit dynamic activations and 4-bit weights
 linear_config = Int8DynamicActivationIntxWeightConfig(
     weight_dtype=torch.int4,
     weight_granularity=PerGroup(32),
diff --git a/docs/source/conf.py b/docs/source/conf.py
index d5799fa042b..4414281bd37 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -321,6 +321,8 @@
     "llm/llama-demo-android": "../using-executorch-android.html",
     "tutorial-arm-ethos-u": "backends/arm-ethos-u/tutorials/ethos-u-getting-started.html",
     "tutorial-arm-vgf": "backends/arm-vgf/tutorials/vgf-getting-started.html",
+    "backends/arm-ethos-u/tutorials/arm-ethos-u-tutorials": "ethos-u-getting-started.html",
+    "backends/arm-vgf/tutorials/arm-vgf-tutorials": "vgf-getting-started.html",
     "visualization": "visualize.html",
 }
 
diff --git a/docs/source/desktop-coreml.md b/docs/source/desktop-coreml.md
index ff6551aa0c2..00149b03e68 100644
--- a/docs/source/desktop-coreml.md
+++ b/docs/source/desktop-coreml.md
@@ -1 +1,2 @@
 ```{include} backends/coreml/coreml-overview.md
+```
\ No newline at end of file
diff --git a/docs/source/desktop-mps.md b/docs/source/desktop-mps.md
index 13717675ba5..84e5b8b7792 100644
--- a/docs/source/desktop-mps.md
+++ b/docs/source/desktop-mps.md
@@ -1 +1,2 @@
 ```{include} backends/mps/mps-overview.md
+```
\ No newline at end of file
diff --git a/docs/source/embedded-arm-ethos-u.md b/docs/source/embedded-arm-ethos-u.md
index 6eb9622925d..cd3b7e525eb 100644
--- a/docs/source/embedded-arm-ethos-u.md
+++ b/docs/source/embedded-arm-ethos-u.md
@@ -1 +1,2 @@
 ```{include} backends/arm-ethos-u/arm-ethos-u-overview.md
+```
\ No newline at end of file
diff --git a/docs/source/embedded-backends.md b/docs/source/embedded-backends.md
index 147f6cfc151..31b04f1839d 100644
--- a/docs/source/embedded-backends.md
+++ b/docs/source/embedded-backends.md
@@ -13,7 +13,7 @@ Available hardware acceleration backends for embedded systems.
 
 ## NPU Acceleration
 
-- {doc}`embedded-arm-ethos-u` — ARM Ethos-U NPU acceleration
+- {doc}`embedded-arm-ethos-u` — Arm Ethos-U NPU acceleration
 - {doc}`embedded-nxp` — NXP eIQ Neutron Backend
 
 
@@ -23,3 +23,4 @@ embedded-arm-cortex-m
 embedded-cadence
 embedded-arm-ethos-u
 embedded-nxp
+```
diff --git a/docs/source/ir-exir.md b/docs/source/ir-exir.md
index d3357d34a03..0c819cf558a 100644
--- a/docs/source/ir-exir.md
+++ b/docs/source/ir-exir.md
@@ -175,4 +175,19 @@ for all core ATen ops.
 
 ## Backend Dialect
 
-See this [doc](compiler-backend-dialect.md)
+Backend dialect is the optional, target-aware stage after Edge dialect. It is
+used when a backend needs to rewrite the graph with backend-specific operators,
+metadata, or delegated lowered modules before the program is converted to an
+ExecuTorch program. For example, a backend pass can replace an Edge dialect
+subgraph such as `addmm` followed by `relu` with a single backend operator that
+the target can execute more efficiently.
+
+Unlike custom operators that may appear in eager mode, ATen dialect, or Edge
+dialect, backend-specific operators are introduced only by passes that run after
+Edge dialect. They are useful for target-specific fusions, lowering patterns,
+or delegate integration where the resulting graph contains nodes meaningful only
+to the selected backend.
+
+For details on when to use backend-specific operators, how they differ from
+custom operators and delegates, and how to register replacement patterns with
+`bind_pattern_to_op`, see [Backend Dialect](compiler-backend-dialect.md).
diff --git a/docs/source/llm/export-llm-optimum.md b/docs/source/llm/export-llm-optimum.md
index e2c8ee14743..b7de8d99689 100644
--- a/docs/source/llm/export-llm-optimum.md
+++ b/docs/source/llm/export-llm-optimum.md
@@ -45,15 +45,11 @@ Optimum ExecuTorch supports a wide range of model architectures including decode
 
 For the complete list of supported models, see the [Optimum ExecuTorch documentation](https://github.com/huggingface/optimum-executorch#-supported-models).
 
-## Export Methods
+## CLI Export
 
-Optimum ExecuTorch offers two ways to export models:
+The `optimum-cli` command is the recommended way to export Hugging Face models. It provides a single invocation that downloads the model from the Hub, applies the configured optimizations, and writes the resulting `.pte` file.
 
-### Method 1: CLI Export
-
-The CLI is the simplest way to export models. It provides a single command to convert models from Hugging Face Hub to ExecuTorch format.
-
-#### Basic Export
+### Basic Export
 
 ```bash
 optimum-cli export executorch \
@@ -63,7 +59,7 @@ optimum-cli export executorch \
     --output_dir="./smollm2_exported"
 ```
 
-#### With Optimizations
+### With Optimizations
 
 Add custom SDPA, KV cache optimization, and quantization:
 
@@ -79,7 +75,7 @@ optimum-cli export executorch \
     --output_dir="./smollm2_exported"
 ```
 
-#### Available CLI Arguments
+### Available CLI Arguments
 
 Key arguments for LLM export include `--model`, `--task`, `--recipe` (backend), `--use_custom_sdpa`, `--use_custom_kv_cache`, `--qlinear` (linear quantization), `--qembedding` (embedding quantization), and `--max_seq_len`.
 
@@ -156,8 +152,8 @@ print(generated_text)
 After verifying your model works correctly, deploy it to device:
 
 - [Running with C++](run-with-c-plus-plus.md) - Run exported models using ExecuTorch's C++ runtime
-- [Running on Android](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/android) - Deploy to Android devices
-- [Running on iOS](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/apple) - Deploy to iOS devices
+- [Running on Android](run-on-android.md) - Java APIs for the `executorch-android` AAR (sample app: [LlamaDemo](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/android))
+- [Running on iOS](run-on-ios.md) - Objective-C and Swift APIs for the `ExecuTorchLLM` framework (sample app: [etLLM](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/apple))
 
 ## Performance
 
diff --git a/docs/source/llm/getting-started.md b/docs/source/llm/getting-started.md
index 95caae6ddd9..1985a610cae 100644
--- a/docs/source/llm/getting-started.md
+++ b/docs/source/llm/getting-started.md
@@ -25,6 +25,6 @@ Deploying LLMs to ExecuTorch can be boiled down to a two-step process: (1) expor
 
 ### Running
 - [Running with C++](run-with-c-plus-plus.md)
-- [Running on Android (XNNPack)](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/android)
+- [Running on Android](run-on-android.md)
 - [Running on Android (Qualcomm)](build-run-llama3-qualcomm-ai-engine-direct-backend.md)
-- [Running on iOS](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/apple)
+- [Running on iOS](run-on-ios.md)
diff --git a/docs/source/llm/run-on-android.md b/docs/source/llm/run-on-android.md
new file mode 100644
index 00000000000..81abd6a79d5
--- /dev/null
+++ b/docs/source/llm/run-on-android.md
@@ -0,0 +1,202 @@
+# Running LLMs on Android
+
+ExecuTorch's LLM-specific runtime components provide an experimental Java interface around the core C++ LLM runtime, available through the `executorch-android` AAR.
+
+## Prerequisites
+
+Make sure you have a model and tokenizer files ready, as described in the prerequisites section of the [Running LLMs with C++](run-with-c-plus-plus.md) guide.
+
+To add the `executorch-android` library to your app, see [Using ExecuTorch on Android](../using-executorch-android.md). The LLM runner classes are bundled inside the same AAR as the generic `Module` API.
+
+## Runtime API
+
+Once the `executorch-android` AAR is on your classpath, you can import the LLM runner classes from the `org.pytorch.executorch.extension.llm` package.
+
+### Importing
+
+```java
+import org.pytorch.executorch.extension.llm.LlmModule;
+import org.pytorch.executorch.extension.llm.LlmModuleConfig;
+import org.pytorch.executorch.extension.llm.LlmGenerationConfig;
+import org.pytorch.executorch.extension.llm.LlmCallback;
+```
+
+### LlmModule
+
+The `LlmModule` class provides a simple Java interface for loading a text-generation model, configuring its tokenizer, generating token streams, and stopping execution. It also supports multimodal models that accept image and audio inputs alongside a text prompt.
+
+This API is experimental and subject to change.
+
+#### Initialization
+
+Create an `LlmModule` by specifying paths to your serialized model (`.pte`) and tokenizer files. For text-only models, the simple constructor is enough:
+
+```java
+LlmModule module = new LlmModule(
+    "/data/local/tmp/llama-3.2-instruct.pte",
+    "/data/local/tmp/tokenizer.model",
+    0.8f);
+```
+
+For finer control (multimodal model type, BOS/EOS handling, supplementary data files, load mode), use `LlmModuleConfig` with the fluent builder:
+
+```java
+LlmModuleConfig config = LlmModuleConfig.create()
+    .modulePath("/data/local/tmp/llama-3.2-instruct.pte")
+    .tokenizerPath("/data/local/tmp/tokenizer.model")
+    .temperature(0.8f)
+    .modelType(LlmModuleConfig.MODEL_TYPE_TEXT)
+    .loadMode(LlmModuleConfig.LOAD_MODE_MMAP)
+    .build();
+
+LlmModule module = new LlmModule(config);
+```
+
+Available load modes are `LOAD_MODE_FILE`, `LOAD_MODE_MMAP` (default), `LOAD_MODE_MMAP_USE_MLOCK`, and `LOAD_MODE_MMAP_USE_MLOCK_IGNORE_ERRORS`. Available model types are `MODEL_TYPE_TEXT`, `MODEL_TYPE_TEXT_VISION`, and `MODEL_TYPE_MULTIMODAL`.
+
+Construction itself is lightweight and does not load the program data immediately.
+
+#### Loading
+
+Explicitly load the model before generation to avoid paying the load cost during your first `generate` call.
+
+```java
+int status = module.load();
+if (status != 0) {
+  // Handle load failure (status is an ExecuTorch runtime error code).
+}
+```
+
+If you skip this step, the model is loaded lazily on the first `generate` call.
+
+#### Generating
+
+Generate tokens from a text prompt by passing an `LlmCallback` that receives each token as it is produced. The same callback also receives a JSON-encoded statistics string when generation completes.
+
+```java
+LlmCallback callback = new LlmCallback() {
+  @Override
+  public void onResult(String token) {
+    // Called once per generated token. Append to your UI buffer here.
+    System.out.print(token);
+  }
+
+  @Override
+  public void onStats(String statsJson) {
+    // Called once when generation finishes. See extension/llm/runner/stats.h
+    // for the field definitions.
+    System.out.println("\n" + statsJson);
+  }
+
+  @Override
+  public void onError(int errorCode, String message) {
+    // Called if the runtime reports an error during generation.
+  }
+};
+
+module.generate("Once upon a time", callback);
+```
+
+For full control over generation parameters, use `LlmGenerationConfig`:
+
+```java
+LlmGenerationConfig genConfig = LlmGenerationConfig.create()
+    .seqLen(2048)
+    .temperature(0.8f)
+    .echo(false)
+    .build();
+
+module.generate("Once upon a time", genConfig, callback);
+```
+
+`LlmGenerationConfig` exposes `echo`, `maxNewTokens`, `seqLen`, `temperature`, `numBos`, `numEos`, and `warming`. Defaults match the C++ `GenerationConfig` documented in [Running LLMs with C++](run-with-c-plus-plus.md).
+
+#### Stopping Generation
+
+If you need to interrupt a long-running generation, call `stop()` from another thread (or from inside the `onResult` callback):
+
+```java
+module.stop();
+```
+
+Generation also runs synchronously on the calling thread, so make sure you invoke `generate()` off the main thread (for example, on a `HandlerThread` or via a `java.util.concurrent.Executor`).
+
+#### Resetting
+
+To clear the prefilled tokens from the KV cache and reset the start position to 0, call:
+
+```java
+module.resetContext();
+```
+
+This is the equivalent of `reset()` on the iOS runner and `reset()` on the C++ `IRunner`.
+
+### Multimodal Inputs
+
+For models declared as `MODEL_TYPE_TEXT_VISION` or `MODEL_TYPE_MULTIMODAL`, image and audio data are provided through dedicated prefill methods. After prefilling all modalities, call `generate()` with the text prompt to produce the response.
+
+#### Images
+
+Raw uint8 pixel data in CHW order can be supplied as an `int[]`, or as a direct `ByteBuffer` to avoid JNI array copies:
+
+```java
+// As int[]
+int[] pixels = ...;       // length == channels * height * width
+module.prefillImages(pixels, /*width=*/336, /*height=*/336, /*channels=*/3);
+
+// As direct ByteBuffer (preferred for large images)
+ByteBuffer buffer = ByteBuffer.allocateDirect(3 * 336 * 336);
+buffer.put(rawBytes).rewind();
+module.prefillImages(buffer, 336, 336, 3);
+```
+
+Pre-normalized float pixel data is also supported, both as a `float[]` and as a direct `ByteBuffer` in native byte order:
+
+```java
+float[] normalized = ...;  // length == channels * height * width
+module.prefillImages(normalized, 336, 336, 3);
+
+ByteBuffer floatBuffer = ByteBuffer
+    .allocateDirect(3 * 336 * 336 * Float.BYTES)
+    .order(ByteOrder.nativeOrder());
+// fill floatBuffer with normalized values, then:
+module.prefillNormalizedImage(floatBuffer, 336, 336, 3);
+```
+
+#### Audio
+
+Preprocessed audio features (for example mel spectrograms produced by a Whisper preprocessor) can be supplied as `byte[]` or `float[]`:
+
+```java
+module.prefillAudio(features, /*batchSize=*/1, /*nBins=*/128, /*nFrames=*/3000);
+```
+
+Raw audio samples can be supplied with `prefillRawAudio`:
+
+```java
+module.prefillRawAudio(samples, /*batchSize=*/1, /*nChannels=*/1, /*nSamples=*/16000);
+```
+
+#### Generating with Multimodal Prefill
+
+After prefilling each modality, run `generate()` with the text prompt as usual:
+
+```java
+module.prefillImages(pixels, 336, 336, 3);
+module.generate("What's in this image?", callback);
+```
+
+For text-vision models, a convenience overload accepts the image and prompt together:
+
+```java
+module.generate(
+    pixels, /*width=*/336, /*height=*/336, /*channels=*/3,
+    "What's in this image?",
+    /*seqLen=*/768,
+    callback,
+    /*echo=*/false);
+```
+
+## Demo
+
+See the [Llama Android demo app](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/android/LlamaDemo) in `executorch-examples` for an end-to-end project that wires `LlmModule`, `LlmCallback`, and a `HandlerThread` into a chat UI.
diff --git a/docs/source/llm/working-with-llms.md b/docs/source/llm/working-with-llms.md
index e4088efd12b..ce6daff6ce8 100644
--- a/docs/source/llm/working-with-llms.md
+++ b/docs/source/llm/working-with-llms.md
@@ -15,5 +15,6 @@ export-llm-optimum
 export-custom-llm
 run-with-c-plus-plus
 build-run-llama3-qualcomm-ai-engine-direct-backend
+run-on-android
 run-on-ios
 ```
diff --git a/docs/source/raspberry_pi_llama_tutorial.md b/docs/source/raspberry_pi_llama_tutorial.md
index 1e886db694a..46d6e3d4fb0 100644
--- a/docs/source/raspberry_pi_llama_tutorial.md
+++ b/docs/source/raspberry_pi_llama_tutorial.md
@@ -65,7 +65,7 @@ cd executorch
 
 ```bash
 # Create conda environment
-conda create -yn executorch python=3.10.0
+conda create -yn executorch python=3.10
 conda activate executorch
 
 # Upgrade pip
diff --git a/docs/source/tools-section.md b/docs/source/tools-section.md
index 6d8061dd33a..d2bb44086f8 100644
--- a/docs/source/tools-section.md
+++ b/docs/source/tools-section.md
@@ -13,7 +13,7 @@ In this section, explore ExecuTorch's comprehensive developer tools for profilin
 - {doc}`model-inspector` — Model Inspector
 - {doc}`memory-planning-inspection` — Memory Planning Inspection
 - {doc}`devtools-tutorial` — Development Utilities
-- {doc}`visualize` — Model Visualization
+- [Model Visualization](visualize.md) — Visualize exported models as computational graphs
 
 ```{toctree}
 :hidden:
@@ -29,4 +29,5 @@ model-debugging
 model-inspector
 memory-planning-inspection
 devtools-tutorial
-visualize
+Model Visualization 
+```
diff --git a/docs/source/tutorials_source/devtools-debugging-tutorial.py b/docs/source/tutorials_source/devtools-debugging-tutorial.py
index 54bcbb50542..6534f7edf1a 100644
--- a/docs/source/tutorials_source/devtools-debugging-tutorial.py
+++ b/docs/source/tutorials_source/devtools-debugging-tutorial.py
@@ -208,9 +208,14 @@
 #
 # The returned DataFrame contains columns for each operator including:
 #
+# - ``aot_debug_handle``: The debug handle tuple identifying the AOT operator(s).
+#   For ops grouped into a single delegated subgraph chunk, this is the tuple of
+#   every internal handle the chunk implements.
 # - ``aot_ops``: The operators in the eager model graph
 # - ``aot_intermediate_output``: Intermediate outputs from eager model
-# - ``runtime_ops``: The operators executed at runtime (may show DELEGATE_CALL for delegated ops)
+# - ``runtime_ops``: The kernel-level operators executed at runtime. ``DELEGATE_CALL``
+#   appears as a single entry for each delegated subgraph chunk.
+# - ``runtime_debug_handle``: The debug handle tuple from the runtime
 # - ``runtime_intermediate_output``: Intermediate outputs from runtime
 # - ``gap``: The numerical gap (MSE) between eager and runtime outputs
 # - ``stacktraces``: A dictionary mapping each operator name to its source code stack trace
@@ -219,14 +224,19 @@
 #
 # .. code-block:: text
 #
-#    |    | aot_ops                                                         | aot_intermediate_output                            | runtime_ops                                        | runtime_intermediate_output                        | gap                        | stacktraces                                        |
-#    |----|----------------------------------------------------------------|----------------------------------------------------|----------------------------------------------------|----------------------------------------------------| ---------------------------|----------------------------------------------------|
-#    | 0  | [conv2d]                                                        | [[[tensor([-0.0130,  0.0075, -0.0334, -0.0122,...  | [DELEGATE_CALL]                                    | [[[tensor([-0.0130,  0.0075, -0.0334, -0.0122,...  | [3.2530690555343034e-15]   | {'conv2d': 'File "model.py", line 10...'}         |
-#    | 1  | [permute, cat, add, dropout]                                    | [[[tensor(-0.0024), tensor(0.0054), tensor(0.0...  | [DELEGATE_CALL]                                    | [[[tensor(-0.0024), tensor(0.0054), tensor(0.0...  | [3.2488685838924244e-15]   | {'permute': 'File "model.py", line 15...', ...}   |
-#    ...
-#    | 4  | [transpose, linear, unflatten, unsqueeze, tran...]              | [[[tensor(0.0045), tensor(-0.0084), tensor(0.0...  | [DELEGATE_CALL, DELEGATE_CALL, DELEGATE_CALL, ...] | [[tensor(0.0045), tensor(-0.0084), tensor(0.00...  | [0.00010033142876115867]   | {'transpose': 'File "model.py", line 20...', ...} |
-#    ...
-#    | 59 | [transpose_66, linear_44, unflatten_11, unsque...]              | [[[tensor(-0.3346), tensor(0.1540), tensor(-0....  | [DELEGATE_CALL, DELEGATE_CALL, DELEGATE_CALL, ...] | [[tensor(-0.3346), tensor(0.1540), tensor(-0.0...  | [0.02629170972698486]      | {'transpose_66': 'File "model.py", line 50...'... |
+#    |    | aot_debug_handle               | aot_ops                                       | aot_intermediate_output                  | runtime_ops                                       | runtime_debug_handle           | runtime_intermediate_output              | gap                      | stacktraces                                  |
+#    |----|--------------------------------|-----------------------------------------------|------------------------------------------|---------------------------------------------------|--------------------------------|------------------------------------------|--------------------------|----------------------------------------------|
+#    | 0  | (1, 2, 3, 5, 6, 7)             | [conv2d, reshape, permute, cat, add, dropout] | [[[tensor(-0.0024), tensor(0.0054),...   | [DELEGATE_CALL]                                   | (1, 2, 3, 5, 6, 7)             | [[[tensor(-0.0024), tensor(0.0054),...   | [0.0]                    | {'conv2d': 'File "vit.py", line 10...'}      |
+#    | 1  | (4,)                           | [expand]                                      | [[[tensor(-0.0012), tensor(0.0027),...   | [native_call_expand_copy.out]                     | (4,)                           | [[[tensor(-0.0012), tensor(0.0027),...   | [0.0]                    | {'expand': 'File "vit.py", line 50...'}      |
+#    | 2  | (8,)                           | [layer_norm]                                  | [[[tensor(-0.0001), tensor(0.0009),...   | [native_call_native_layer_norm.out]               | (8,)                           | [[[tensor(-0.0001), tensor(0.0009),...   | [1.5257813119012582e-16] | {'layer_norm': 'File "vit.py", line 80...'}  |
+#    | 3  | (9, 10, 11, 12, 13)            | [transpose, linear, unflatten, unsqueeze,...] | [[[tensor(0.0132), tensor(-0.0005),...   | [DELEGATE_CALL, native_call_expand_copy.out, ...] | (9, 10, 11, 12, 13)            | [[[tensor(0.0132), tensor(-0.0005),...   | [6.72759426534678e-16]   | {'transpose': 'File "vit.py", line 90...'}   |
+#    | ...|                                |                                               |                                          |                                                   |                                |                                          |                          |                                              |
+#    | 52 | (166,)                         | [linear_48]                                   | [[tensor(-0.8981), tensor(0.6046),...    | [DELEGATE_CALL]                                   | (166,)                         | [[tensor(-0.8981), tensor(0.6046),...    | [8.760367306096254e-12]  | {'linear_48': 'File "vit.py", line 120...'}  |
+#
+# Each row is one runtime intermediate-output cluster: either a single
+# non-delegated op or one delegated subgraph chunk grouped under a
+# ``DELEGATE_CALL``. Multi-handle chunks list every internal AOT op together
+# so you can map the delegate boundary back to the eager graph.
 #
 # The ``stacktraces`` column is particularly useful for tracing operators back to the
 # original PyTorch source code. Each entry is a dictionary where keys are operator names
@@ -274,34 +284,24 @@
 # .. code-block:: text
 #
 #    Top 5 operators with largest numerical discrepancies:
-#                                                  aot_ops                            aot_intermediate_output                                        runtime_ops                        runtime_intermediate_output                     gap                                        stacktraces
-#    59  [transpose_66, linear_44, unflatten_11, unsque...  [[[tensor(-0.3346), tensor(0.1540), tensor(-0....  [DELEGATE_CALL, DELEGATE_CALL, DELEGATE_CALL, ...  [[tensor(-0.3346), tensor(0.1540), tensor(-0.0...   [0.02629170972698486]  {'transpose_66': 'File "vit.py", line 125...'}
-#    24  [transpose_24, linear_16, unflatten_4, unsquee...  [[[tensor(0.0344), tensor(-0.0583), tensor(-0....  [DELEGATE_CALL, DELEGATE_CALL, DELEGATE_CALL, ...  [[tensor(0.0344), tensor(-0.0583), tensor(-0.0...  [0.010045093258604096]  {'transpose_24': 'File "vit.py", line 125...'}
-#    29  [transpose_30, linear_20, unflatten_5, unsquee...  [[[tensor(0.0457), tensor(0.0266), tensor(-0.0...  [DELEGATE_CALL, DELEGATE_CALL, DELEGATE_CALL, ...  [[tensor(0.0457), tensor(0.0266), tensor(-0.05...  [0.008497326594593926]  {'transpose_30': 'File "vit.py", line 125...'}
-#    34  [transpose_36, linear_24, unflatten_6, unsquee...  [[[tensor(-0.1336), tensor(-0.0154), tensor(-0...  [DELEGATE_CALL, DELEGATE_CALL, DELEGATE_CALL, ...  [[tensor(-0.1336), tensor(-0.0154), tensor(-0....  [0.007672668965640913]  {'transpose_36': 'File "vit.py", line 125...'}
-#    19  [transpose_18, linear_12, unflatten_3, unsquee...  [[[tensor(-0.0801), tensor(0.0458), tensor(-0....  [DELEGATE_CALL, DELEGATE_CALL, DELEGATE_CALL, ...  [[tensor(-0.0801), tensor(0.0458), tensor(-0.0...  [0.007446783635888463]  {'transpose_18': 'File "vit.py", line 125...'}
-#
-#    --- Operator 59 ---
-#    Operators: ['transpose_66', 'linear_44', 'unflatten_11', 'unsqueeze_11', 'transpose_67']
-#    Gap: [0.02629170972698486]
-#    Stack traces:
-#      transpose_66:
-#        File "torchvision/models/vision_transformer.py", line 125, in forward
-#          x = self.self_attention(x)
-#        File "torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
+#                       aot_debug_handle                                            aot_ops                            aot_intermediate_output                                        runtime_ops            runtime_debug_handle                        runtime_intermediate_output                       gap                                        stacktraces
+#    49  (158, 159, 160, 161, 162, 163)  [linear_46, gelu_11, dropout_35, linear_47, dr...  [[[tensor(-0.6154), tensor(-0.0025), tensor(-0...                                    [DELEGATE_CALL]  (158, 159, 160, 161, 162, 163)  [[[tensor(-0.6154), tensor(-0.0025), tensor(-0...  [3.7287804660557997e-11]  {'linear_46': 'File "vit.py", line 100...'}
+#    45  (145, 146, 147, 148, 149, 150)  [linear_42, gelu_10, dropout_32, linear_43, dr...  [[[tensor(-0.2785), tensor(-0.2209), tensor(-0...                                    [DELEGATE_CALL]  (145, 146, 147, 148, 149, 150)  [[[tensor(-0.2785), tensor(-0.2208), tensor(-0...  [3.267051816109279e-11]  {'linear_42': 'File "vit.py", line 100...'}
+#    47       (152, 153, 154, 155, 156)  [transpose_66, linear_44, unflatten_11, unsque...  [[[tensor(-0.5681), tensor(-0.0251), tensor(-0...  [DELEGATE_CALL, native_call_expand_copy.out, D...       (152, 153, 154, 155, 156)  [[[tensor(-0.5681), tensor(-0.0251), tensor(-0...  [3.043042878567803e-11]  {'transpose_66': 'File "vit.py", line 90...'}
+#    50                          (164,)                                    [layer_norm_24]  [[[tensor(-0.8486), tensor(-0.0315), tensor(-0...                [native_call_native_layer_norm.out]                          (164,)  [[[tensor(-0.8486), tensor(-0.0315), tensor(-0...  [2.3309619445739474e-11]  {'layer_norm_24': 'File "vit.py", line 78...'}
+#    41  (132, 133, 134, 135, 136, 137)  [linear_38, gelu_9, dropout_29, linear_39, dro...  [[[tensor(-0.3233), tensor(-0.1968), tensor(-0...                                    [DELEGATE_CALL]  (132, 133, 134, 135, 136, 137)  [[[tensor(-0.3233), tensor(-0.1968), tensor(-0...  [1.8573805941767968e-11]  {'linear_38': 'File "vit.py", line 100...'}
 #
 #    Operators with MSE > 0.0001:
-#                                                  aot_ops                            aot_intermediate_output                                        runtime_ops                        runtime_intermediate_output                       gap                                        stacktraces
-#    4   [transpose, linear, unflatten, unsqueeze, tran...  [[[tensor(0.0045), tensor(-0.0084), tensor(0.0...  [DELEGATE_CALL, DELEGATE_CALL, DELEGATE_CALL, ...  [[tensor(0.0045), tensor(-0.0084), tensor(0.00...  [0.00010033142876115867]  {'transpose': 'File "vit.py", line 125...'}
-#    9   [transpose_6, linear_4, unflatten_1, unsqueeze...  [[[tensor(0.0113), tensor(-0.0737), tensor(-0....  [DELEGATE_CALL, DELEGATE_CALL, DELEGATE_CALL, ...  [[tensor(0.0113), tensor(-0.0737), tensor(-0.0...   [0.0005611182577030275]  {'transpose_6': 'File "vit.py", line 125...'}
-#    14  [transpose_12, linear_8, unflatten_2, unsqueez...  [[[tensor(-0.0476), tensor(-0.0941), tensor(-0...  [DELEGATE_CALL, DELEGATE_CALL, DELEGATE_CALL, ...  [[tensor(-0.0476), tensor(-0.0941), tensor(-0....    [0.004658652508649068]  {'transpose_12': 'File "vit.py", line 125...'}
-#    ...
-#
-# In this example, we can see that the attention layers (transpose + linear + unflatten patterns)
-# show the largest numerical discrepancies, which is expected behavior for delegated operators
-# using different precision. The ``stacktraces`` column shows that these operators originate from
-# ``self.self_attention(x)`` in the ViT model's forward method, helping you identify exactly
-# where in your model code the discrepancies arise.
+#    Empty DataFrame
+#    Columns: [aot_debug_handle, aot_ops, aot_intermediate_output, runtime_ops, runtime_debug_handle, runtime_intermediate_output, gap, stacktraces]
+#    Index: []
+#
+# The largest numerical gaps come from MLP blocks (``linear, gelu, dropout`` groups
+# inside a ``DELEGATE_CALL``) and attention QKV blocks, both at the order of
+# ``~3e-11`` — floating-point rounding accumulated across the chunk at float32
+# precision. ``layer_norm`` rows are non-delegated CPU kernels and show similar
+# magnitudes. No operators exceed the ``1e-4`` threshold, confirming that
+# XNNPACK float32 delegation is numerically accurate.
 
 ######################################################################
 # Pipeline 2: CMake Runtime
diff --git a/docs/source/using-executorch-android.md b/docs/source/using-executorch-android.md
index ef55ade68aa..0c034d99924 100644
--- a/docs/source/using-executorch-android.md
+++ b/docs/source/using-executorch-android.md
@@ -21,6 +21,7 @@ All ExecuTorch Android libraries are packaged into an Android library (AAR), exe
 The AAR artifact contains the Java library for users to integrate with their Java/Kotlin application code, as well as the corresponding JNI library (.so file), which is loaded by the Java code during initialization.
 
 - [Java library](https://github.com/pytorch/executorch/tree/main/extension/android/executorch_android/src/main/java/org/pytorch/executorch)
+- [Java API Reference (Javadoc)](https://pytorch.org/executorch/main/javadoc/index.html)
 - JNI contains the JNI binding for the corresponding Java code, and ExecuTorch native library, including
   - Core ExecuTorch runtime libraries
   - XNNPACK backend
@@ -240,4 +241,4 @@ using ExecuTorch AAR package.
 
 ## Java API reference
 
-Please see [Java API reference](https://pytorch.org/executorch/main/javadoc/).
+Please see [Java API reference](https://pytorch.org/executorch/main/javadoc/index.html).
diff --git a/docs/source/using-executorch-building-from-source.md b/docs/source/using-executorch-building-from-source.md
index 5fabcfab682..170c3f8758d 100644
--- a/docs/source/using-executorch-building-from-source.md
+++ b/docs/source/using-executorch-building-from-source.md
@@ -45,7 +45,7 @@ portability details.
    ```bash
    git clone -b viable/strict https://github.com/pytorch/executorch.git
    cd executorch
-   conda create -yn executorch python=3.10.0
+   conda create -yn executorch python=3.10
    conda activate executorch
    ```
 
diff --git a/docs/source/using-executorch-export.md b/docs/source/using-executorch-export.md
index d37dfae2ef7..30f2a22368e 100644
--- a/docs/source/using-executorch-export.md
+++ b/docs/source/using-executorch-export.md
@@ -45,6 +45,10 @@ Commonly used hardware backends are listed below. For mobile, consider using XNN
 
 The export process takes in a standard PyTorch model, typically a `torch.nn.Module`. This can be an custom model definition, or a model from an existing source, such as TorchVision or HuggingFace. See [Getting Started with ExecuTorch](getting-started.md) for an example of lowering a TorchVision model.
 
+:::{tip}
+Exporting a model from the [Hugging Face Hub](https://huggingface.co/models)? Use the [Optimum ExecuTorch](llm/export-llm-optimum.md) integration. It wraps the export and lowering steps below in a single CLI invocation and supports a wide range of decoder, encoder, multimodal, and seq2seq architectures out of the box.
+:::
+
 Model export is done from Python. This is commonly done through a Python script or from an interactive Python notebook, such as Jupyter or Colab. The example below shows instantiation and inputs for a simple PyTorch model. The inputs are prepared as a tuple of torch.Tensors, and the model can run with these inputs.
 
 ```python
diff --git a/examples/apple/coreml/llama/llama_transformer.py b/examples/apple/coreml/llama/llama_transformer.py
index ae98c327b45..0175d0deaf6 100644
--- a/examples/apple/coreml/llama/llama_transformer.py
+++ b/examples/apple/coreml/llama/llama_transformer.py
@@ -14,7 +14,7 @@
 
 import torch
 import torch.nn.functional as F
-from executorch.examples.models.llama.norm import RMSNorm
+from executorch.examples.models.llama.norm import RMSNorm, RMSNormCoreML  # noqa: F401
 
 from executorch.examples.models.llama.rope import (
     hf_apply_rotary_emb,
@@ -109,65 +109,6 @@ def __post_init__(self):
             self.head_dim = self.dim // self.n_heads
 
 
-class CoreMLRMSNorm(torch.nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-6):
-        """
-        Initialize the RMSNorm normalization layer.
-
-        Args:
-            dim (int): The dimension of the input tensor.
-            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
-
-        Attributes:
-            eps (float): A small value added to the denominator for numerical stability.
-            weight (nn.Parameter): Learnable scaling parameter.
-
-        """
-        super().__init__()
-        self.dim = dim
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-
-    def _norm(self, x):
-        """
-        Apply the RMSNorm normalization to the input tensor.
-
-        Args:
-            x (torch.Tensor): The input tensor.
-
-        Returns:
-            torch.Tensor: The normalized tensor.
-
-        """
-        # CoreML ignores casts to FP32, so existing implementation of RMSNorm was not stable
-        # We instead use (x * sqrt(n)) / norm(x, dim=-1)
-        # Using torch.norm and preserving this op in CoreML improves stability
-        # Note, we ignore eps, but could add it by using torch.norm(torch.concat(x, sqrt(n*eps))) in the denominator
-        # In future, we want to add CoreML support for the functional RMSNorm op
-        # We have yet to do large scale evaluations on the numeric stability of this solution, but note that
-        # it appears better than what exists currently (removing FP32 casts and using FP16)
-        rms_norm_eps0 = (
-            x
-            * torch.sqrt(torch.tensor(self.dim, dtype=x.dtype))
-            * torch.reciprocal(torch.linalg.vector_norm(x, dim=-1, keepdim=True))
-        )
-        return rms_norm_eps0
-
-    def forward(self, x):
-        """
-        Forward pass through the RMSNorm layer.
-
-        Args:
-            x (torch.Tensor): The input tensor.
-
-        Returns:
-            torch.Tensor: The output tensor after applying RMSNorm.
-
-        """
-        output = self._norm(x)
-        return output * self.weight
-
-
 class Rope(torch.nn.Module):
     def __init__(self, params: ModelArgs):
         super().__init__()
diff --git a/examples/arm/README.md b/examples/arm/README.md
index bcd8a1e1d0a..c5f5bb24862 100644
--- a/examples/arm/README.md
+++ b/examples/arm/README.md
@@ -1,3 +1,10 @@
+
+
 ## ExecuTorch for Arm backends Ethos-U, VGF and Cortex-M
 
 This project contains scripts to help you setup and run a PyTorch
@@ -11,22 +18,25 @@ The main scripts are `setup.sh`, `run.sh` and
 `setup.sh` will install the needed tools and with --root-dir  
 you can change the path to a scratch folder where it will download and generate build
 artifacts. If supplied, you must also supply the same folder to run.sh with
---scratch-dir= If not supplied both script will use examples/arm/arm-scratch
+--scratch-dir= If not supplied both scripts will use examples/arm/arm-scratch.
 
 `run.sh` can be used to build, run and test a model in an easy way and it will call cmake for you
 and in cases you want to run a simulator it will start it also. The script will call `aot_arm_compiler.py`
 to convert a model and include it in the build/run.
 
+For bare-metal Ethos-U builds `run.sh` configures the standalone
+`examples/arm/executor_runner/standalone` CMake entry point automatically. If
+`--build-dir` is omitted, the script creates and owns a build tree under
+`arm_test/_`. Supplying `--build-dir` reuses an existing tree
+(for example a VGF host build or out-of-tree configuration) and `run.sh`
+verifies it exposes the runner options it needs before compiling.
+
 Build and test artifacts are by default placed under the folder arm_test folder
 this can be changed with --et_build_root=
 
 `aot_arm_compiler.py` is used to convert a Python model or a saved .pt model to a PTE file and is used by `run.sh`
 and other test script but can also be used directly.
 
-If you prefer to use the ExecuTorch API, there is also the `ethos_u_minimal_example.ipynb` notebook example.
-This shows the workflow if you prefer to integrate a python torch.export and ExecuTorch flow directly into your
-model codebase. This is particularly useful if you want to perform more complex training, such as quantization
-aware training using the ArmQuantizer.
 
 ## Create a PTE file for Arm backends
 
@@ -64,6 +74,16 @@ $ python3 -m backends.arm.scripts.aot_arm_compiler --model_name=mv2 --target=eth
 
 `aot_arm_compiler.py` is called from the scripts below so you don't need to, but it can be useful to do by hand in some cases.
 
+## Host VGF example applications
+
+The Arm examples directory also contains host-side VGF reference flows for
+specific tasks:
+
+- `examples/arm/image_classification_example_vgf` for DEiT image
+  classification.
+- `examples/arm/super_resolution_example_vgf` for Swin2SR image
+  super-resolution.
+
 
 ## ExecuTorch on Arm Ethos-U55/U65 and U85
 
diff --git a/examples/arm/cortex_m_mv2_example.ipynb b/examples/arm/cortex_m_mv2_example.ipynb
index c2fe4342773..36844b4e5fd 100644
--- a/examples/arm/cortex_m_mv2_example.ipynb
+++ b/examples/arm/cortex_m_mv2_example.ipynb
@@ -136,7 +136,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "%%bash \n# Build example executor runner application to examples/arm/cortex_m_mv2_example\n# Note that this is the same runner as used in the Ethos-U example, creating some overlap in the config even though the Ethos-U is not used.\ncmake -DCMAKE_TOOLCHAIN_FILE=$(pwd)/ethos-u-setup/arm-none-eabi-gcc.cmake \\\n      -DCMAKE_BUILD_TYPE=Release \\\n      -DET_PTE_FILE_PATH=cortex_m_mv2_example.bpte \\\n      -DTARGET_CPU=cortex-m55 \\\n      -DETHOSU_TARGET_NPU_CONFIG=ethos-u55-128 \\\n      -DMEMORY_MODE=Shared_Sram \\\n      -DET_BUNDLE_IO=ON \\\n      -DSYSTEM_CONFIG=Ethos_U55_High_End_Embedded \\\n      -Bcortex_m_mv2_example \\\n      executor_runner\ncmake --build cortex_m_mv2_example -j$(nproc) -- arm_executor_runner"
+   "source": "%%bash \n# Build example executor runner application to examples/arm/cortex_m_mv2_example\n# Note that this is the same runner as used in the Ethos-U example, creating some overlap in the config even though the Ethos-U is not used.\ncmake -DCMAKE_TOOLCHAIN_FILE=$(pwd)/ethos-u-setup/arm-none-eabi-gcc.cmake \\\n      -DCMAKE_BUILD_TYPE=Release \\\n      -DET_PTE_FILE_PATH=cortex_m_mv2_example.bpte \\\n      -DTARGET_CPU=cortex-m55 \\\n      -DETHOSU_TARGET_NPU_CONFIG=ethos-u55-128 \\\n      -DMEMORY_MODE=Shared_Sram \\\n      -DET_BUNDLE_IO=ON \\\n      -DSYSTEM_CONFIG=Ethos_U55_High_End_Embedded \\\n      -Bcortex_m_mv2_example \\\n      -S executor_runner/standalone\ncmake --build cortex_m_mv2_example -j$(nproc) -- arm_executor_runner"
   },
   {
    "cell_type": "markdown",
@@ -179,4 +179,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
\ No newline at end of file
+}
diff --git a/examples/arm/ethos_u_minimal_example.ipynb b/examples/arm/ethos_u_minimal_example.ipynb
index fbb15cd0e57..11f24019d23 100644
--- a/examples/arm/ethos_u_minimal_example.ipynb
+++ b/examples/arm/ethos_u_minimal_example.ipynb
@@ -171,26 +171,8 @@
    "source": [
     "## Build executor runtime\n",
     "\n",
-    "After the AOT compilation flow is done, the runtime can be cross compiled and linked to the produced .pte-file using the Arm cross-compilation toolchain. This is done in two steps:\n",
-    "1. Build and install the executorch libraries and EthosUDelegate.\n",
-    "2. Build and link the `arm_executor_runner` and generate kernel bindings for any non delegated ops."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%%bash\n",
-    "# Ensure the arm-none-eabi-gcc toolchain and FVP:s are available on $PATH\n",
-    "source arm-scratch/setup_path.sh\n",
-    "\n",
-    "# Build executorch libraries cross-compiled for arm baremetal to executorch/cmake-out-arm\n",
-    "cmake --preset arm-baremetal \\\n",
-    "-DCMAKE_BUILD_TYPE=Release \\\n",
-    "-B../../cmake-out-arm ../..\n",
-    "cmake --build ../../cmake-out-arm --target install -j$(nproc) "
+    "After the AOT compilation flow finishes, cross-compile and link the runtime by configuring the standalone `examples/arm/executor_runner/standalone` CMake project with the Arm toolchain.\n",
+    "It automatically pulls the ExecuTorch checkout in as a dependency so the delegate, kernels, and runner util are rebuilt alongside the application, and it generates kernel bindings for any non-delegated ops found in the `.pte`.\n"
    ]
   },
   {
@@ -201,6 +183,8 @@
    "source": [
     "%%bash \n",
     "source arm-scratch/setup_path.sh\n",
+    "# Ensure CMake resolves the ExecuTorch checkout root regardless of caller env\n",
+    "export EXECUTORCH_ROOT=$(cd ../.. && pwd)\n",
     "\n",
     "# Build example executor runner application to examples/arm/ethos_u_minimal_example\n",
     "cmake -DCMAKE_TOOLCHAIN_FILE=$(pwd)/ethos-u-setup/arm-none-eabi-gcc.cmake \\\n",
@@ -211,7 +195,7 @@
     "      -DMEMORY_MODE=Shared_Sram \\\n",
     "      -DSYSTEM_CONFIG=Ethos_U55_High_End_Embedded \\\n",
     "      -Bethos_u_minimal_example \\\n",
-    "      executor_runner\n",
+    "      -S executor_runner/standalone\n",
     "cmake --build ethos_u_minimal_example -j$(nproc) -- arm_executor_runner"
    ]
   },
@@ -232,6 +216,8 @@
    "source": [
     "%%bash \n",
     "source arm-scratch/setup_path.sh\n",
+    "# Ensure CMake resolves the ExecuTorch checkout root regardless of caller env\n",
+    "export EXECUTORCH_ROOT=$(cd ../.. && pwd)\n",
     "\n",
     "# Run the example\n",
     "../../backends/arm/scripts/run_fvp.sh --elf=ethos_u_minimal_example/arm_executor_runner --target=ethos-u55-128"
diff --git a/examples/arm/executor_runner/CMakeLists.txt b/examples/arm/executor_runner/CMakeLists.txt
index c169f5d447a..d84947a75ad 100644
--- a/examples/arm/executor_runner/CMakeLists.txt
+++ b/examples/arm/executor_runner/CMakeLists.txt
@@ -3,8 +3,49 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-cmake_minimum_required(VERSION 3.20)
-project(arm_executor_runner)
+if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR)
+  cmake_minimum_required(VERSION 3.20)
+  project(arm_executor_runner_redirect LANGUAGES C CXX)
+  message(
+    FATAL_ERROR
+      "Configure standalone arm_executor_runner builds from ${CMAKE_CURRENT_LIST_DIR}/standalone instead of ${CMAKE_CURRENT_LIST_DIR}."
+  )
+endif()
+
+get_filename_component(
+  _default_executorch_root "${CMAKE_CURRENT_LIST_DIR}/../../.." ABSOLUTE
+)
+
+if(NOT DEFINED EXECUTORCH_ROOT)
+  set(EXECUTORCH_ROOT
+      "${_default_executorch_root}"
+      CACHE PATH "Path to an ExecuTorch checkout"
+  )
+endif()
+
+set(ET_DIR_PATH
+    "${EXECUTORCH_ROOT}"
+    CACHE PATH "Kept for backward compatibility; synonym for EXECUTORCH_ROOT"
+)
+if(NOT DEFINED ET_INCLUDE_PATH)
+  set(ET_INCLUDE_PATH
+      "${EXECUTORCH_ROOT}"
+      CACHE
+        PATH
+        "Kept for backward compatibility; include root for ExecuTorch headers"
+  )
+endif()
+
+if(NOT EXISTS "${EXECUTORCH_ROOT}/CMakeLists.txt")
+  message(
+    FATAL_ERROR
+      "EXECUTORCH_ROOT (${EXECUTORCH_ROOT}) does not contain an ExecuTorch CMake project."
+  )
+endif()
+
+if(NOT COMMAND executorch_target_link_options_shared_lib)
+  include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)
+endif()
 
 option(
   ET_MODEL_PTE_ADDR
@@ -21,6 +62,12 @@ option(ET_LOG_DUMP_INPUT "Dump input in log" OFF)
 option(ET_LOG_DUMP_OUTPUT "Dump output in log" ON)
 
 option(ET_BUNDLE_IO "Set to compile in BundleIO support" OFF)
+set(BUNDLED_PROGRAM_LIBRARY_DIR
+    ""
+    CACHE
+      PATH
+      "Optional directory that contains a prebuilt libbundled_program.a when ET_BUNDLE_IO is enabled without building devtools."
+)
 set(ET_ATOL
     "0.01"
     CACHE STRING "Set atol to use for BundleIO testing (Requires ET_BUNDLE_IO)"
@@ -55,13 +102,46 @@ option(
   OFF
 )
 
+if(NOT DEFINED PYTHON_EXECUTABLE)
+  find_package(
+    Python3
+    COMPONENTS Interpreter
+    REQUIRED
+  )
+  set(PYTHON_EXECUTABLE "${Python3_EXECUTABLE}")
+endif()
+
+include(${EXECUTORCH_ROOT}/backends/arm/scripts/corstone_utils.cmake)
+include(${EXECUTORCH_ROOT}/backends/arm/cmake/ArmEthosUSDK.cmake)
+include(${EXECUTORCH_ROOT}/backends/arm/cmake/ArmRunnerUtils.cmake)
+
+arm_runner_require_baremetal_targets()
+
+# Keep the default scratch location aligned with the scratch tree used by
+# setup.sh/run.sh so developers who just ran those scripts do not need extra
+# CMake flags.
+set(ETHOS_SDK_PATH
+    "${EXECUTORCH_ROOT}/examples/arm/arm-scratch/ethos-u"
+    CACHE PATH "Path to Ethos-U bare metal driver/env"
+)
+
+arm_ethos_u_default_fetch("${ETHOS_SDK_PATH}" _fetch_ethos_u_default)
 option(FETCH_ETHOS_U_CONTENT
-       "Fetch ethos_u dependencies instead of relying on pre-downloads" ON
+       "Fetch ethos_u dependencies instead of relying on pre-downloads"
+       ${_fetch_ethos_u_default}
+)
+arm_ensure_ethos_u_content(
+  "${ETHOS_SDK_PATH}" "${EXECUTORCH_ROOT}" ${FETCH_ETHOS_U_CONTENT}
 )
 
-if(NOT DEFINED ET_MODEL_PTE_ADDR
-   AND NOT DEFINED ET_PTE_FILE_PATH
-   AND NOT DEFINED SEMIHOSTING
+set(ET_PTE_FILE_PATH
+    ""
+    CACHE PATH "Path to ExecuTorch model pte"
+)
+
+if(NOT ET_MODEL_PTE_ADDR
+   AND "${ET_PTE_FILE_PATH}" STREQUAL ""
+   AND NOT SEMIHOSTING
 )
   message(
     FATAL_ERROR
@@ -72,39 +152,16 @@ if(NOT DEFINED ET_MODEL_PTE_ADDR
   )
 endif()
 
-# Example ExecuTorch demo for bare metal Cortex-M based systems
-set(ET_DIR_PATH
-    "${CMAKE_CURRENT_SOURCE_DIR}/../../.."
-    CACHE PATH "Path to ExecuTorch dir"
+if(NOT SEMIHOSTING
+   AND NOT ET_MODEL_PTE_ADDR
+   AND NOT "${ET_PTE_FILE_PATH}" STREQUAL ""
 )
-include(${ET_DIR_PATH}/tools/cmake/Utils.cmake)
-set(ET_BUILD_DIR_PATH
-    "${ET_DIR_PATH}/cmake-out-arm"
-    CACHE PATH "Path to ExecuTorch build/install dir"
-)
-set(ET_INCLUDE_PATH
-    "${ET_DIR_PATH}/.."
-    CACHE PATH "Path to ExecuTorch headers"
-)
-set(ET_PTE_FILE_PATH
-    ""
-    CACHE PATH "Path to ExecuTorch model pte"
-)
-set(ETHOS_SDK_PATH
-    "${ET_DIR_PATH}/examples/arm/arm-scratch/ethos-u"
-    CACHE PATH "Path to Ethos-U bare metal driver/env"
-)
-set(PYTHON_EXECUTABLE
-    "python"
-    CACHE PATH "Define to override python executable used"
-)
-
-# Include corstone help functions
-include(${ET_DIR_PATH}/backends/arm/scripts/corstone_utils.cmake)
-
-if(FETCH_ETHOS_U_CONTENT)
-  # Download ethos_u dependency if needed.
-  fetch_ethos_u_content(${ETHOS_SDK_PATH} ${ET_DIR_PATH})
+  if(NOT EXISTS "${ET_PTE_FILE_PATH}")
+    message(
+      FATAL_ERROR
+        "ET_PTE_FILE_PATH is set to ${ET_PTE_FILE_PATH}, but no file was found. Generate the model first or point ET_PTE_FILE_PATH at an existing .pte/.bpte."
+    )
+  endif()
 endif()
 
 # Selects timing adapter values matching system_config. Default is
@@ -154,23 +211,33 @@ message(
 add_corstone_subdirectory(${SYSTEM_CONFIG} ${ETHOS_SDK_PATH})
 configure_timing_adapters(${SYSTEM_CONFIG} ${MEMORY_MODE})
 
-# Dependencies from the ExecuTorch build
-find_package(
-  executorch REQUIRED HINTS "${ET_BUILD_DIR_PATH}/lib/cmake/ExecuTorch"
-)
+if(NOT CMAKE_SKIP_INSTALL_RULES AND TARGET ethosu_core_driver)
+  get_property(
+    _et_ethosu_core_driver_exported GLOBAL
+    PROPERTY ET_ETHOSU_CORE_DRIVER_EXPORTED
+  )
+  if(NOT _et_ethosu_core_driver_exported)
+    install(
+      TARGETS ethosu_core_driver
+      EXPORT ExecuTorchTargets
+      ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+      LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    )
+    set_property(GLOBAL PROPERTY ET_ETHOSU_CORE_DRIVER_EXPORTED TRUE)
+  endif()
+endif()
 
 # Convert pte to header
-if(NOT ${ET_MODEL_PTE_ADDR} AND NOT SEMIHOSTING)
+if(NOT "${ET_MODEL_PTE_ADDR}" AND NOT SEMIHOSTING)
   add_custom_target(
     gen_model_header DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/model_pte.h
   )
 
   add_custom_command(
     OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/model_pte.h
-    COMMAND ${PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/pte_to_header.py --pte
-            ${ET_PTE_FILE_PATH} --outdir ${CMAKE_CURRENT_BINARY_DIR}
+    COMMAND ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/pte_to_header.py
+            --pte ${ET_PTE_FILE_PATH} --outdir ${CMAKE_CURRENT_BINARY_DIR}
     DEPENDS ${ET_PTE_FILE_PATH}
-    WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   )
 endif()
 
@@ -199,12 +266,12 @@ endif()
 # Proceed with specific actions if either is found
 if(NOT U55_FOUND EQUAL -1)
   message(STATUS "SYSTEM_CONFIG contains 'U55'.")
-  set(LINK_FILE_IN "${CMAKE_SOURCE_DIR}/Corstone-300.ld")
+  set(LINK_FILE_IN "${CMAKE_CURRENT_LIST_DIR}/Corstone-300.ld")
 endif()
 
 if(NOT U85_FOUND EQUAL -1)
   message(STATUS "SYSTEM_CONFIG contains 'U85'.")
-  set(LINK_FILE_IN "${CMAKE_SOURCE_DIR}/Corstone-320.ld")
+  set(LINK_FILE_IN "${CMAKE_CURRENT_LIST_DIR}/Corstone-320.ld")
 endif()
 
 if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
@@ -213,9 +280,9 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
   set(COMPILER_PREPROCESSOR_OPTIONS -E -x c -P)
 endif()
 
-get_filename_component(LINK_FILE_OUT_BASE ${LINK_FILE} NAME)
+get_filename_component(LINK_FILE_OUT_BASE "${LINK_FILE}" NAME)
 set(LINK_FILE_OUT
-    ${CMAKE_CURRENT_BINARY_DIR}/${LINK_FILE_OUT_BASE}.${LINK_FILE_EXT}
+    "${CMAKE_CURRENT_BINARY_DIR}/${LINK_FILE_OUT_BASE}.${LINK_FILE_EXT}"
 )
 
 execute_process(
@@ -251,26 +318,31 @@ list(
 # (user-set)SELECT_OPS_MODEL variable. For normal build, use
 # EXECUTORCH_SELECT_OPS_MODEL to include ops automatically. If the pte contains
 # no undelegated ops, use neither.
-execute_process(
-  COMMAND
-    python "${ET_DIR_PATH}/codegen/tools/gen_oplist.py"
-    --model_file_path=${ET_PTE_FILE_PATH}
-    --output_path=${CMAKE_CURRENT_BINARY_DIR}/temp.yaml
-  OUTPUT_VARIABLE CMD_RESULT
+set(FOUND_OPS_IN_FILE FALSE)
+if(NOT SEMIHOSTING
+   AND NOT ET_MODEL_PTE_ADDR
+   AND NOT "${ET_PTE_FILE_PATH}" STREQUAL ""
+   AND EXISTS "${ET_PTE_FILE_PATH}"
 )
+  execute_process(
+    COMMAND
+      ${PYTHON_EXECUTABLE} "${EXECUTORCH_ROOT}/codegen/tools/gen_oplist.py"
+      --model_file_path=${ET_PTE_FILE_PATH}
+      --output_path=${CMAKE_CURRENT_BINARY_DIR}/temp.yaml
+    OUTPUT_VARIABLE CMD_RESULT
+  )
 
-if(CMD_RESULT MATCHES "aten::" OR CMD_RESULT MATCHES "dim_order_ops::")
-  set(FOUND_OPS_IN_FILE "true")
-else()
-  set(FOUND_OPS_IN_FILE "false")
+  if(CMD_RESULT MATCHES "aten::" OR CMD_RESULT MATCHES "dim_order_ops::")
+    set(FOUND_OPS_IN_FILE TRUE)
+  endif()
 endif()
 
-if(${SEMIHOSTING})
+if(SEMIHOSTING)
   set(EXECUTORCH_SELECT_OPS_MODEL "")
   message(
     "gen_oplist: Building with semihosting, no model is used to auto generate ops from will use EXECUTORCH_SELECT_OPS_LIST=${EXECUTORCH_SELECT_OPS_LIST}"
   )
-elseif(${FOUND_OPS_IN_FILE})
+elseif(FOUND_OPS_IN_FILE)
   set(EXECUTORCH_SELECT_OPS_LIST "")
   set(EXECUTORCH_SELECT_OPS_MODEL "${ET_PTE_FILE_PATH}")
   message(
@@ -289,10 +361,6 @@ endif()
 if(NOT ("${EXECUTORCH_SELECT_OPS_LIST}" STREQUAL ""
         AND "${EXECUTORCH_SELECT_OPS_MODEL}" STREQUAL "")
 )
-  set(EXECUTORCH_ROOT ${ET_DIR_PATH})
-  include(${ET_DIR_PATH}/tools/cmake/Utils.cmake)
-  include(${ET_DIR_PATH}/tools/cmake/Codegen.cmake)
-
   gen_selected_ops(
     LIB_NAME
     "arm_portable_ops_lib"
@@ -310,7 +378,7 @@ if(NOT ("${EXECUTORCH_SELECT_OPS_LIST}" STREQUAL ""
 
   generate_bindings_for_kernels(
     LIB_NAME "arm_portable_ops_lib" FUNCTIONS_YAML
-    ${ET_DIR_PATH}/kernels/portable/functions.yaml DTYPE_SELECTIVE_BUILD
+    ${EXECUTORCH_ROOT}/kernels/portable/functions.yaml DTYPE_SELECTIVE_BUILD
     "${EXECUTORCH_ENABLE_DTYPE_SELECTIVE_BUILD}"
   )
   gen_operators_lib(
@@ -333,16 +401,66 @@ if(EXECUTORCH_ENABLE_EVENT_TRACER)
 endif()
 
 if(ET_BUNDLE_IO)
-  list(APPEND arm_executor_runner_link bundled_program)
+  if(TARGET bundled_program)
+    list(APPEND arm_executor_runner_link bundled_program)
+    target_link_directories(
+      arm_executor_runner PRIVATE $
+    )
+  else()
+    set(_bundled_program_library "")
+    set(_bundled_program_search_paths "")
+    if(ET_BUILD_DIR_PATH)
+      list(
+        APPEND
+        _bundled_program_search_paths
+        "${ET_BUILD_DIR_PATH}"
+        "${ET_BUILD_DIR_PATH}/lib"
+        "${ET_BUILD_DIR_PATH}/devtools/bundled_program"
+        "${ET_BUILD_DIR_PATH}/devtools/bundled_program/lib"
+      )
+    endif()
+    if(BUNDLED_PROGRAM_LIBRARY_DIR)
+      list(APPEND _bundled_program_search_paths
+           "${BUNDLED_PROGRAM_LIBRARY_DIR}"
+      )
+    endif()
+    if(_bundled_program_search_paths)
+      list(REMOVE_DUPLICATES _bundled_program_search_paths)
+      # BundleIO can reuse a separate ExecuTorch build tree where
+      # bundled_program is not part of this CMake graph. Restrict the fallback
+      # lookup to the caller-provided build directories so we do not
+      # accidentally pick up an unrelated library from the host system.
+      find_library(
+        _bundled_program_library
+        NAMES bundled_program
+        PATHS ${_bundled_program_search_paths}
+        NO_DEFAULT_PATH
+      )
+    endif()
+    if(_bundled_program_library)
+      list(APPEND arm_executor_runner_link ${_bundled_program_library})
+    else()
+      message(
+        FATAL_ERROR
+          "ET_BUNDLE_IO enabled but bundled_program is unavailable. Either configure this build with EXECUTORCH_BUILD_DEVTOOLS=ON so the target exists or set BUNDLED_PROGRAM_LIBRARY_DIR/ET_BUILD_DIR_PATH to a build directory that contains libbundled_program.a."
+      )
+    endif()
+  endif()
 endif()
 
 # Need whole-archive to ensure C++ ctor's are called - this may be wasteful for
 # bin size as we link in a number of other symbols
 target_link_libraries(arm_executor_runner PUBLIC ${arm_executor_runner_link})
 
+# Ensure the ELF lands next to the CMake build tree so run.sh (and downstream
+# tooling) can locate it deterministically regardless of multi-config vs
+# single-config generators.
 target_link_options(
   arm_executor_runner PUBLIC LINKER:-Map=arm_executor_runner.map
 )
+# Reuse a parent build's output directory if it already set one; otherwise keep
+# the runner ELF next to this build tree so run.sh can find it predictably.
+arm_runner_configure_runtime_output(arm_executor_runner "${CMAKE_BINARY_DIR}")
 
 # Sanitizers
 if(CMAKE_BUILD_TYPE MATCHES "UndefinedSanitizer")
@@ -351,7 +469,7 @@ if(CMAKE_BUILD_TYPE MATCHES "UndefinedSanitizer")
   target_link_options(arm_executor_runner PRIVATE ${_et_runner_ubsan_flag})
   if(NOT TARGET executorch_ubsan)
     add_subdirectory(
-      ${ET_DIR_PATH}/examples/arm/ubsan
+      ${EXECUTORCH_ROOT}/examples/arm/ubsan
       ${CMAKE_CURRENT_BINARY_DIR}/ubsan_runtime
     )
   endif()
@@ -367,7 +485,8 @@ if(CMAKE_BUILD_TYPE MATCHES "AddressSanitizer")
   target_link_options(arm_executor_runner PRIVATE ${_et_runner_asan_flags})
   if(NOT TARGET executorch_asan)
     add_subdirectory(
-      ${ET_DIR_PATH}/examples/arm/asan ${CMAKE_CURRENT_BINARY_DIR}/asan_runtime
+      ${EXECUTORCH_ROOT}/examples/arm/asan
+      ${CMAKE_CURRENT_BINARY_DIR}/asan_runtime
     )
   endif()
   target_link_libraries(arm_executor_runner PRIVATE executorch_asan)
@@ -377,16 +496,18 @@ if(CMAKE_BUILD_TYPE MATCHES "AddressSanitizer")
 endif()
 
 # ET headers and generated headers includes
+set(_arm_runner_include_dirs
+    ${ET_INCLUDE_PATH} ${ET_INCLUDE_PATH}/runtime/core/portable_type/c10
+    ${CMAKE_CURRENT_BINARY_DIR}
+)
 target_include_directories(
-  arm_executor_runner
-  PRIVATE ${ET_INCLUDE_PATH} ${ET_DIR_PATH}/runtime/core/portable_type/c10
-          ${CMAKE_CURRENT_BINARY_DIR}
+  arm_executor_runner PRIVATE ${_arm_runner_include_dirs}
 )
 target_compile_definitions(
   arm_executor_runner PRIVATE C10_USING_CUSTOM_GENERATED_MACROS
 )
 
-if(NOT ${ET_MODEL_PTE_ADDR} AND NOT SEMIHOSTING)
+if(NOT "${ET_MODEL_PTE_ADDR}" AND NOT SEMIHOSTING)
   add_dependencies(arm_executor_runner gen_model_header)
 endif()
 
diff --git a/examples/arm/executor_runner/pte_to_header.py b/examples/arm/executor_runner/pte_to_header.py
index 65213bc729e..8656ac5abdf 100644
--- a/examples/arm/executor_runner/pte_to_header.py
+++ b/examples/arm/executor_runner/pte_to_header.py
@@ -1,6 +1,7 @@
+#!/usr/bin/env python3
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
-# Copyright 2023-2025 Arm Limited and/or its affiliates.
+# Copyright 2023-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
diff --git a/examples/arm/executor_runner/standalone/CMakeLists.txt b/examples/arm/executor_runner/standalone/CMakeLists.txt
new file mode 100644
index 00000000000..73493ca9e71
--- /dev/null
+++ b/examples/arm/executor_runner/standalone/CMakeLists.txt
@@ -0,0 +1,159 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+cmake_minimum_required(VERSION 3.20)
+project(arm_executor_runner LANGUAGES C CXX)
+
+get_filename_component(
+  _default_executorch_root "${CMAKE_CURRENT_LIST_DIR}/../../../.." ABSOLUTE
+)
+
+if(NOT DEFINED EXECUTORCH_ROOT)
+  set(EXECUTORCH_ROOT
+      "${_default_executorch_root}"
+      CACHE PATH "Path to an ExecuTorch checkout"
+  )
+endif()
+
+if(NOT EXISTS "${EXECUTORCH_ROOT}/CMakeLists.txt")
+  if(EXISTS "${_default_executorch_root}/CMakeLists.txt")
+    message(
+      WARNING
+        "EXECUTORCH_ROOT (${EXECUTORCH_ROOT}) does not contain an ExecuTorch CMake project. Falling back to ${_default_executorch_root}."
+    )
+    set(EXECUTORCH_ROOT
+        "${_default_executorch_root}"
+        CACHE PATH "Path to an ExecuTorch checkout" FORCE
+    )
+  else()
+    message(
+      FATAL_ERROR
+        "EXECUTORCH_ROOT (${EXECUTORCH_ROOT}) does not contain an ExecuTorch CMake project."
+    )
+  endif()
+endif()
+
+set(ARM_EXECUTOR_RUNNER_STANDALONE
+    ON
+    CACHE BOOL
+          "Indicates arm_executor_runner was configured as a standalone project"
+          FORCE
+)
+
+# Load the preset helper so standalone builds inherit the same defaults as the
+# superbuild (toolchains, delegated targets, devtools options, etc.).
+set(_executorch_preset_cmake
+    "${EXECUTORCH_ROOT}/tools/cmake/common/preset.cmake"
+)
+if(EXISTS "${_executorch_preset_cmake}")
+  include("${_executorch_preset_cmake}")
+  if(NOT DEFINED EXECUTORCH_BUILD_PRESET_FILE)
+    set(EXECUTORCH_BUILD_PRESET_FILE
+        "${EXECUTORCH_ROOT}/tools/cmake/preset/arm_baremetal.cmake"
+        CACHE PATH "Preset used when configuring the standalone runner"
+    )
+  endif()
+  load_build_preset()
+endif()
+include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)
+
+option(ARM_EXECUTOR_RUNNER_SKIP_INSTALL_RULES
+       "Skip install() rules for standalone arm_executor_runner builds" ON
+)
+if(DEFINED CMAKE_SKIP_INSTALL_RULES)
+  set(_arm_runner_skip_install_rules "${CMAKE_SKIP_INSTALL_RULES}")
+endif()
+if(ARM_EXECUTOR_RUNNER_SKIP_INSTALL_RULES)
+  set(CMAKE_SKIP_INSTALL_RULES ON)
+endif()
+
+foreach(
+  _opt
+  EXECUTORCH_BUILD_ARM_BAREMETAL EXECUTORCH_BUILD_CORTEX_M
+  EXECUTORCH_BUILD_KERNELS_QUANTIZED EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL
+)
+  if(NOT DEFINED ${_opt})
+    set(${_opt}
+        ON
+        CACHE BOOL "" FORCE
+    )
+  endif()
+endforeach()
+set(EXECUTORCH_SKIP_ARM_EXECUTOR_RUNNER
+    ON
+    CACHE BOOL "" FORCE
+)
+
+# examples/arm/executor_runner/CMakeLists.txt generates the runner-specific
+# portable-op registration based on the PTE or an explicit select-ops list.
+# Avoid feeding those cache entries into the top-level ExecuTorch configure,
+# otherwise executorch_core auto-right-sizes MAX_KERNEL_NUM from the runner's
+# placeholder/selective build inputs even though the runner also links
+# quantized/cortex-m registration libraries.
+set(_arm_runner_selective_cache_vars
+    EXECUTORCH_SELECT_OPS_LIST EXECUTORCH_SELECT_OPS_MODEL
+    EXECUTORCH_SELECT_OPS_YAML
+)
+foreach(_arm_runner_cache_var IN LISTS _arm_runner_selective_cache_vars)
+  if(DEFINED CACHE{${_arm_runner_cache_var}})
+    get_property(
+      _arm_runner_cache_type
+      CACHE ${_arm_runner_cache_var}
+      PROPERTY TYPE
+    )
+    if(NOT _arm_runner_cache_type OR _arm_runner_cache_type STREQUAL
+                                     "UNINITIALIZED"
+    )
+      set(_arm_runner_cache_type STRING)
+    endif()
+    set(_arm_runner_saved_type_${_arm_runner_cache_var}
+        "${_arm_runner_cache_type}"
+    )
+    set(_arm_runner_saved_value_${_arm_runner_cache_var}
+        "${${_arm_runner_cache_var}}"
+    )
+    set(_arm_runner_saved_defined_${_arm_runner_cache_var} TRUE)
+    set(${_arm_runner_cache_var}
+        ""
+        CACHE ${_arm_runner_cache_type} "" FORCE
+    )
+  endif()
+endforeach()
+
+if(NOT DEFINED CACHE{MAX_KERNEL_NUM} AND NOT DEFINED MAX_KERNEL_NUM)
+  set(MAX_KERNEL_NUM
+      2000
+      CACHE STRING
+            "Maximum number of kernels registered by the standalone Arm runner"
+  )
+endif()
+
+# Pull ExecuTorch in-tree so all required targets (delegates, kernels, runner
+# util, etc.) are built from this checkout.
+add_subdirectory(
+  ${EXECUTORCH_ROOT} ${CMAKE_BINARY_DIR}/executorch EXCLUDE_FROM_ALL
+)
+
+if(ARM_EXECUTOR_RUNNER_SKIP_INSTALL_RULES)
+  if(DEFINED _arm_runner_skip_install_rules)
+    set(CMAKE_SKIP_INSTALL_RULES "${_arm_runner_skip_install_rules}")
+  else()
+    unset(CMAKE_SKIP_INSTALL_RULES)
+  endif()
+endif()
+
+foreach(_arm_runner_cache_var IN LISTS _arm_runner_selective_cache_vars)
+  if(_arm_runner_saved_defined_${_arm_runner_cache_var})
+    set(${_arm_runner_cache_var}
+        "${_arm_runner_saved_value_${_arm_runner_cache_var}}"
+        CACHE ${_arm_runner_saved_type_${_arm_runner_cache_var}} "" FORCE
+    )
+  endif()
+endforeach()
+
+add_subdirectory(
+  ${EXECUTORCH_ROOT}/examples/arm/executor_runner
+  ${CMAKE_BINARY_DIR}/examples/arm/executor_runner
+)
diff --git a/examples/arm/image_classification_example_ethos_u/model_export/export_deit.py b/examples/arm/image_classification_example_ethos_u/model_export/export_deit.py
index 0291df7d044..ce015a4d457 100644
--- a/examples/arm/image_classification_example_ethos_u/model_export/export_deit.py
+++ b/examples/arm/image_classification_example_ethos_u/model_export/export_deit.py
@@ -6,8 +6,8 @@
 import argparse
 
 import torch
-import tqdm
-from datasets import DatasetDict, load_dataset
+import tqdm  # type: ignore[import]
+from datasets import DatasetDict, load_dataset  # type: ignore[import]
 
 from executorch.backends.arm.ethosu import EthosUCompileSpec, EthosUPartitioner
 from executorch.backends.arm.quantizer import (
@@ -21,7 +21,10 @@
 )
 from executorch.extension.export_util.utils import save_pte_program
 
-from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
+from torchao.quantization.pt2e.quantize_pt2e import (  # type: ignore[import]
+    convert_pt2e,
+    prepare_pt2e,
+)
 from transformers import AutoImageProcessor
 from transformers.models.vit.modeling_vit import ViTForImageClassification
 
diff --git a/examples/arm/image_classification_example_ethos_u/model_export/train_deit.py b/examples/arm/image_classification_example_ethos_u/model_export/train_deit.py
index bc9c322fd51..afc93c76663 100644
--- a/examples/arm/image_classification_example_ethos_u/model_export/train_deit.py
+++ b/examples/arm/image_classification_example_ethos_u/model_export/train_deit.py
@@ -9,8 +9,8 @@
 
 import numpy as np
 import torch
-from datasets import DatasetDict, load_dataset
-from evaluate import load as load_metric
+from datasets import DatasetDict, load_dataset  # type: ignore[import]
+from evaluate import load as load_metric  # type: ignore[import]
 from transformers import AutoImageProcessor, set_seed, Trainer, TrainingArguments
 
 from transformers.models.vit.modeling_vit import ViTForImageClassification
diff --git a/examples/arm/image_classification_example_vgf/model_export/export_deit.py b/examples/arm/image_classification_example_vgf/model_export/export_deit.py
index 859baab9f13..a46d157bfa5 100644
--- a/examples/arm/image_classification_example_vgf/model_export/export_deit.py
+++ b/examples/arm/image_classification_example_vgf/model_export/export_deit.py
@@ -6,8 +6,8 @@
 import argparse
 
 import torch
-import tqdm
-from datasets import DatasetDict, load_dataset
+import tqdm  # type: ignore[import]
+from datasets import DatasetDict, load_dataset  # type: ignore[import]
 
 from executorch.backends.arm.quantizer import (
     get_symmetric_quantization_config,
@@ -21,7 +21,10 @@
 )
 from executorch.extension.export_util.utils import save_pte_program
 
-from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
+from torchao.quantization.pt2e.quantize_pt2e import (  # type: ignore[import]
+    convert_pt2e,
+    prepare_pt2e,
+)
 from transformers import AutoImageProcessor
 from transformers.models.vit.modeling_vit import ViTForImageClassification
 
diff --git a/examples/arm/image_classification_example_vgf/model_export/train_deit.py b/examples/arm/image_classification_example_vgf/model_export/train_deit.py
index bc9c322fd51..afc93c76663 100644
--- a/examples/arm/image_classification_example_vgf/model_export/train_deit.py
+++ b/examples/arm/image_classification_example_vgf/model_export/train_deit.py
@@ -9,8 +9,8 @@
 
 import numpy as np
 import torch
-from datasets import DatasetDict, load_dataset
-from evaluate import load as load_metric
+from datasets import DatasetDict, load_dataset  # type: ignore[import]
+from evaluate import load as load_metric  # type: ignore[import]
 from transformers import AutoImageProcessor, set_seed, Trainer, TrainingArguments
 
 from transformers.models.vit.modeling_vit import ViTForImageClassification
diff --git a/examples/arm/pruning_minimal_example.ipynb b/examples/arm/pruning_minimal_example.ipynb
index db585b94158..a24c6626a15 100644
--- a/examples/arm/pruning_minimal_example.ipynb
+++ b/examples/arm/pruning_minimal_example.ipynb
@@ -453,7 +453,7 @@
     "      -DMEMORY_MODE=Shared_Sram \\\n",
     "      -DSYSTEM_CONFIG=Ethos_U85_SYS_DRAM_Mid \\\n",
     "      -Bethos_u_original_model \\\n",
-    "      executor_runner\n",
+    "      -S executor_runner/standalone\n",
     "cmake --build ethos_u_original_model -j$(nproc) -- arm_executor_runner"
    ]
   },
@@ -499,7 +499,7 @@
     "      -DMEMORY_MODE=Shared_Sram \\\n",
     "      -DSYSTEM_CONFIG=Ethos_U85_SYS_DRAM_Mid \\\n",
     "      -Bethos_u_pruned_model \\\n",
-    "      executor_runner\n",
+    "      -S executor_runner/standalone\n",
     "cmake --build ethos_u_pruned_model -j$(nproc) -- arm_executor_runner"
    ]
   },
diff --git a/examples/arm/run.sh b/examples/arm/run.sh
index b18115723b0..351eda14071 100755
--- a/examples/arm/run.sh
+++ b/examples/arm/run.sh
@@ -14,8 +14,9 @@ set -eu
 ########
 script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
 et_root_dir=$(cd ${script_dir}/../.. && pwd)
-et_root_dir=$(realpath ${et_root_dir})
-
+et_root_dir=$(realpath "${et_root_dir}")
+runner_source_dir="${et_root_dir}/examples/arm/executor_runner/standalone"
+runner_source_dir=$(realpath "${runner_source_dir}")
 
 model_name=""
 model_input_set=false
@@ -29,7 +30,7 @@ output_folder="."
 bundleio=false
 build_with_etdump=false
 build_type="Release"
-extra_build_flags=""
+build_dir=""
 build_only=false
 system_config=""
 config=""
@@ -38,8 +39,9 @@ pte_placement="elf"
 et_build_root="${et_root_dir}/arm_test"
 arm_scratch_dir=${script_dir}/arm-scratch
 scratch_dir_set=false
-toolchain=arm-none-eabi-gcc
+toolchain="arm-none-eabi-gcc"
 select_ops_list="aten::_softmax.out"
+select_ops_list_overridden=false
 qdq_fusion_op=false
 model_explorer=false
 perf_overlay=false
@@ -47,6 +49,12 @@ 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=""
+build_dir_initialized=false
+multi_config=false
+parallel_jobs=1
 
 function help() {
     echo "Usage: $(basename $0) [options]"
@@ -57,8 +65,7 @@ function help() {
     echo "  --aot_arm_compiler_flags=       Extra flags to pass to aot compiler"
     echo "  --no_delegate                          Do not delegate the model (can't override builtin models)"
     echo "  --no_quantize                          Do not quantize the model (can't override builtin models)"
-    echo "  --portable_kernels=               TO BE DEPRECATED: Alias to select_ops_list."
-    echo "  --select_ops_list=                Comma separated list of portable (non delagated) kernels to include Default: ${select_ops_list}"
+    echo "  --select_ops_list=                Comma separated list of portable (non delegated) kernels to include. Default: ${select_ops_list}"
     echo "                                           NOTE: This is only used when building for semihosting."
     echo "                                           See https://docs.pytorch.org/executorch/stable/kernel-library-selective-build.html for more information."
     echo "  --target=                      Target to build and run for Default: ${target}"
@@ -66,9 +73,10 @@ function help() {
     echo "  --bundleio                             Create Bundled pte using Devtools BundelIO with Input/RefOutput included"
     echo "  --etdump                               Adds Devtools etdump support to track timing, etdump area will be base64 encoded in the log"
     echo "  --build_type=                    Build with Release, Debug, RelWithDebInfo, UndefinedSanitizer or AddressSanitizer, default is ${build_type}"
-    echo "  --extra_build_flags=            Extra flags to pass to cmake like -DET_ARM_BAREMETAL_METHOD_ALLOCATOR_POOL_SIZE=60000 Default: none "
+    echo "  --build-dir=
                      Optional: reuse an existing arm_executor_runner build directory (configured via 'cmake -S examples/arm/executor_runner/standalone -B  ...'). If omitted, run.sh auto-configures one under ${et_build_root} for bare-metal targets."
     echo "  --build_only                           Only build, don't run"
-    echo "  --toolchain=                Ethos-U: Toolchain can be specified (e.g. bare metal as arm-none-eabi-gcc or zephyr as arm-zephyr-eabi-gcc Default: ${toolchain}"
+    echo "  --extra_build_flags=\"\"         Extra -D style flags to pass to cmake when run.sh auto-configures the build"
+    echo "  --toolchain=  Toolchain preset to use when run.sh auto-configures the build. Default: ${toolchain}"
     echo "  --system_config=               Ethos-U: System configuration to select from the Vela configuration file (see vela.ini). Default: Ethos_U55_High_End_Embedded for EthosU55 targets, Ethos_U85_SYS_DRAM_Mid for EthosU85 targets."
     echo "                                            NOTE: If given, this option must match the given target. This option also sets timing adapter values customized for specific hardware, see ./executor_runner/CMakeLists.txt."
     echo "  --config=                    Ethos-U: System configuration file that specifies system configurations (vela.ini)"
@@ -76,11 +84,11 @@ function help() {
     echo "  --pte_placement=             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=               Executorch build output root folder to use, defaults to ${et_build_root}"
-    echo "  --scratch-dir=                 Path to your Arm scrach dir if you not using default ${arm_scratch_dir}"
+    echo "  --scratch-dir=                 Path to your Ethos-U scratch dir if you not using default ${arm_scratch_dir}"
     echo "  --qdq_fusion_op                        Enable QDQ fusion op"
     echo "  --model_explorer                       Enable model explorer to visualize a TOSA or PTE model graph."
     echo "  --visualize_pte                        With --model_explorer, visualize PTE flatbuffer model and delegates. Cannot be used with --visualize_tosa"
-    echo "                                            NOTE:  If PTE contains an Ethos-U delegate, the Ethos-U subgraph will be visualized if aot_arm_compiler_flags is set with the -i flag to include intermediate tosa files."
+    echo "                                            NOTE: If PTE contains an Ethos-U delegate, the Ethos-U subgraph will be visualized if aot_arm_compiler_flags includes -i for TOSA dumps."
     echo "  --visualize_tosa                       With --model_explorer, visualize TOSA flatbuffer model. Cannot be used with --visualize_pte"
     echo "  --perf_overlay                         With --model_explorer and --visualize_tosa, include performance data from FVP PMU trace."
     exit 0
@@ -94,15 +102,18 @@ for arg in "$@"; do
       --aot_arm_compiler_flags=*) aot_arm_compiler_flags="${arg#*=}";;
       --no_delegate) aot_arm_compiler_flag_delegate="" ;;
       --no_quantize) aot_arm_compiler_flag_quantize="" ;;
-      --portable_kernels=*) select_ops_list="${arg#*=}" ; echo "WARNING: --portable_kernels is DEPRECATED use select_ops_list." ;;
-      --select_ops_list=*) select_ops_list="${arg#*=}";;
+      --select_ops_list=*)
+        select_ops_list="${arg#*=}"
+        select_ops_list_overridden=true
+        ;;
       --target=*) target="${arg#*=}";;
       --output=*) output_folder="${arg#*=}" ; output_folder_set=true ;;
       --bundleio) bundleio=true ;;
       --etdump) build_with_etdump=true ;;
       --build_type=*) build_type="${arg#*=}";;
-      --extra_build_flags=*) extra_build_flags="${arg#*=}";;
+      --build-dir=*) build_dir="${arg#*=}";;
       --build_only) build_only=true ;;
+      --extra_build_flags=*) extra_build_flags="${arg#*=}";;
       --toolchain=*) toolchain="${arg#*=}";;
       --system_config=*) system_config="${arg#*=}";;
       --config=*) config="${arg#*=}";;
@@ -121,6 +132,11 @@ for arg in "$@"; do
     esac
 done
 
+auto_configure=false
+if [[ -z "${build_dir}" ]]; then
+    auto_configure=true
+fi
+
 if [ "$perf_overlay" = true ] && [ "$model_explorer" != true ]; then
     echo "Error: --perf_overlay requires --model_explorer" >&2
     exit 1
@@ -141,10 +157,32 @@ if ! [[ ${pte_placement} == "elf" ]]; then
 fi
 
 # Default Ethos-u tool folder override with --scratch-dir=
-arm_scratch_dir=$(realpath ${arm_scratch_dir})
+arm_scratch_dir=$(realpath "${arm_scratch_dir}")
+ethos_u_root_dir="${arm_scratch_dir}/ethos-u"
+mkdir -p "${ethos_u_root_dir}"
+ethos_u_root_dir=$(realpath "${ethos_u_root_dir}")
+cmsis_nn_local_path=""
+if [[ -d "${ethos_u_root_dir}/core_software/cmsis-nn" ]]; then
+    cmsis_nn_local_path=$(realpath "${ethos_u_root_dir}/core_software/cmsis-nn")
+fi
 setup_path_script=${arm_scratch_dir}/setup_path.sh
 _setup_msg="please refer to ${script_dir}/setup.sh to properly install necessary tools."
 
+toolchain_cmake=""
+case "${toolchain}" in
+  arm-none-eabi-gcc)
+    toolchain_cmake="${et_root_dir}/examples/arm/ethos-u-setup/${toolchain}.cmake"
+    ;;
+  arm-zephyr-eabi-gcc)
+    toolchain_cmake="${et_root_dir}/examples/zephyr/x86_64-linux-arm-zephyr-eabi-gcc.cmake"
+    ;;
+  *)
+    echo "Error: Invalid toolchain selection '${toolchain}'. Valid options: arm-none-eabi-gcc, arm-zephyr-eabi-gcc" >&2
+    exit 1
+    ;;
+esac
+
+
 # Set target based variables
 if [[ ${system_config} == "" ]]
 then
@@ -169,26 +207,10 @@ then
     config="Arm/vela.ini"
 fi
 
-# Build executorch libraries
-cd $et_root_dir
-devtools_flag=""
-bundleio_flag=""
-etrecord_flag=""
-et_dump_flag=""
-qdq_fusion_op_flag=""
-fvp_pmu_flag=""
-if [ "$build_with_etdump" = true ] ; then
-    et_dump_flag="--etdump"
-    etrecord_flag="--etrecord"
-fi
-
-if [ "$bundleio" = true ] ; then
-    devtools_flag="--devtools"
-    bundleio_flag="--bundleio"
-fi
-
-if [ "$qdq_fusion_op" = true ] ; then
-    qdq_fusion_op_flag="--enable_qdq_fusion_pass"
+target_cpu="cortex-m85"
+if [[ ${target} =~ "ethos-u55" ]]
+then
+    target_cpu="cortex-m55"
 fi
 
 function check_setup () {
@@ -201,36 +223,31 @@ function check_setup () {
         echo "Could not find ${setup_path_script} file, ${_setup_msg}"
         return 1
     fi
-    # If setup_path_script was correct all these checks should now pass
-    if [[ ${target} =~ "ethos-u" ]]; then
-        if [[ ${toolchain} == "arm-none-eabi-gcc" ]]; then
-            toolchain_cmake=${et_root_dir}/examples/arm/ethos-u-setup/${toolchain}.cmake
-        elif [[ ${toolchain} == "arm-zephyr-eabi-gcc" ]]; then
-            toolchain_cmake=${et_root_dir}/examples/zephyr/x86_64-linux-arm-zephyr-eabi-gcc.cmake
-        else
-            echo "Error: Invalid toolchain selection, provided: ${toolchain}"
-            echo "    Valid options are {arm-none-eabi-gcc, arm-zephyr-eabi-gcc}"
-            exit 1;
+
+    [[ -f ${et_root_dir}/CMakeLists.txt ]] \
+        || { echo "Executorch repo doesn't contain CMakeLists.txt file at root level"; return 1; }
+
+    [[ -f ${preset_file} ]] \
+        || { echo "Could not find ${preset_file} file, ${_setup_msg}"; return 1; }
+
+    if [[ "${auto_configure}" == true && ${target} != *"TOSA"* ]]; then
+        if ! command -v "${toolchain}" >/dev/null 2>&1; then
+            echo "Could not find ${toolchain} toolchain on PATH, ${_setup_msg}"
+            return 1
         fi
-        toolchain_cmake=$(realpath ${toolchain_cmake})
-        hash ${toolchain} \
-            || { echo "Could not find ${toolchain} toolchain on PATH, ${_setup_msg}"; return 1; }
 
         [[ -f ${toolchain_cmake} ]] \
             || { echo "Could not find ${toolchain_cmake} file, ${_setup_msg}"; return 1; }
+    fi
 
-        [[ -f ${et_root_dir}/CMakeLists.txt ]] \
-            || { echo "Executorch repo doesn't contain CMakeLists.txt file at root level"; return 1; }
-
-    backends/arm/scripts/build_executorch.sh --et_build_root="${et_build_root}" --build_type=$build_type $devtools_flag $et_dump_flag --toolchain="${toolchain}"
-    elif [[ ${target} == cortex-m* ]]; then
+    if [[ ${target} == cortex-m* ]]; then
         # build_test_runner.sh handles toolchain setup; just validate it's on PATH.
         hash arm-none-eabi-gcc \
             || { echo "Could not find arm-none-eabi-gcc on PATH, ${_setup_msg}"; return 1; }
     elif [[ ${target} =~ "vgf" ]]; then
-        model_converter=$(which model-converter)
+        model_converter=$(which model-converter || true)
         echo "${model_converter}"
-        [[ "${model_converter}" == "model-converter not found" ]] \
+        [[ -z "${model_converter}" || "${model_converter}" == "model-converter not found" ]] \
             && { echo "Could not find model-converter, ${_setup_msg}"; return 1; }
     fi
 
@@ -252,25 +269,416 @@ print(size)
 PY
 }
 
+sanitize_for_path() {
+    local value="$1"
+    printf '%s' "${value}" | tr -c '[:alnum:]._-' '_'
+}
+
+set_default_build_dir_path() {
+    if [[ ${target} == *"vgf"* ]]; then
+        cat <&2
+Error: auto-configuring a build directory is only supported for Ethos-U bare-metal targets.
+Configure a host build manually, e.g.
+  cmake -S "${runner_source_dir}" -B  -DEXECUTORCH_ROOT="${et_root_dir}" -DEXECUTORCH_BUILD_VGF=ON
+and then pass --build-dir=.
+EOF
+        exit 1
+    fi
+    local sanitized_target
+    sanitized_target=$(sanitize_for_path "${target}")
+    local sanitized_build_type
+    sanitized_build_type=$(sanitize_for_path "${build_type}")
+    local sanitized_toolchain
+    sanitized_toolchain=$(sanitize_for_path "${toolchain}")
+    build_dir="${et_build_root}/${sanitized_target}_${sanitized_build_type}_${sanitized_toolchain}"
+}
+
+configure_runner_build_dir() {
+    local pte_source="$1"
+    if [[ -z "${build_dir}" ]]; then
+        echo "Error: build_dir is not set. Cannot configure runner." >&2
+        exit 1
+    fi
+    if [[ "${pte_placement}" == "elf" ]]; then
+        pte_source=$(realpath "${pte_source}")
+    fi
+    mkdir -p "${build_dir}"
+    local cmake_cmd=(
+        cmake -S "${runner_source_dir}" -B "${build_dir}"
+        -DEXECUTORCH_ROOT="${et_root_dir}"
+        -DCMAKE_TOOLCHAIN_FILE="${toolchain_cmake}"
+        -DCMAKE_BUILD_TYPE="${build_type}"
+        -DEXECUTORCH_BUILD_ARM_BAREMETAL=ON
+        -DEXECUTORCH_BUILD_CORTEX_M=ON
+        -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON
+        -DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=ON
+        -DEXECUTORCH_BUILD_PRESET_FILE="${preset_file}"
+        -DEXECUTORCH_BAREMETAL_SKIP_INSTALL=OFF
+        -DETHOSU_TARGET_NPU_CONFIG="${target}"
+        -DTARGET_CPU="${target_cpu}"
+        -DSYSTEM_CONFIG="${system_config}"
+        -DMEMORY_MODE="${memory_mode}"
+        -DETHOS_SDK_PATH:PATH="${ethos_u_root_dir}"
+        -DEXECUTORCH_SELECT_OPS_LIST="${select_ops_list}"
+    )
+    if [[ -n "${cmsis_nn_local_path}" ]]; then
+        cmake_cmd+=(-DCMSIS_NN_LOCAL_PATH:PATH="${cmsis_nn_local_path}")
+    fi
+    cmake_cmd+=(-DET_PTE_FILE_PATH:PATH="${pte_source}")
+    if [[ "${pte_placement}" == "elf" ]]; then
+        cmake_cmd+=(-DET_MODEL_PTE_ADDR=)
+    else
+        cmake_cmd+=(-DET_MODEL_PTE_ADDR="${pte_placement}")
+    fi
+    if [[ "${bundleio}" == true ]]; then
+        cmake_cmd+=(-DET_BUNDLE_IO=ON)
+    else
+        cmake_cmd+=(-DET_BUNDLE_IO=OFF)
+    fi
+    if [[ "${bundleio}" == true || "${build_with_etdump}" == true ]]; then
+        cmake_cmd+=(-DEXECUTORCH_BUILD_DEVTOOLS=ON)
+    else
+        cmake_cmd+=(-DEXECUTORCH_BUILD_DEVTOOLS=OFF)
+    fi
+    if [[ "${build_with_etdump}" == true ]]; then
+        cmake_cmd+=(-DEXECUTORCH_ENABLE_EVENT_TRACER=ON -DET_DUMP_INTERMEDIATE_OUTPUTS=ON)
+    else
+        cmake_cmd+=(-DEXECUTORCH_ENABLE_EVENT_TRACER=OFF -DET_DUMP_INTERMEDIATE_OUTPUTS=OFF)
+    fi
+    if [[ -n "${extra_build_flags}" ]]; then
+        # shellcheck disable=SC2206
+        local extra_args=(${extra_build_flags})
+        cmake_cmd+=("${extra_args[@]}")
+    fi
+    echo "[run.sh] Configuring ExecuTorch build at ${build_dir}"
+    "${cmake_cmd[@]}"
+    build_dir_initialized=false
+}
+
+cmake_cache_get() {
+    local key="$1"
+    if [[ ! -f ${cmake_cache_file} ]]; then
+        echo ""
+        return 0
+    fi
+    local line
+    line=$(grep -m1 "^${key}:" "${cmake_cache_file}" || true)
+    if [[ -z "${line}" ]]; then
+        echo ""
+    else
+        echo "${line#*=}"
+    fi
+}
+
+cmake_cache_has_key() {
+    local key="$1"
+    [[ -f ${cmake_cache_file} ]] && grep -q "^${key}:" "${cmake_cache_file}"
+}
+
+ensure_runner_build_dir() {
+    local standalone
+    standalone=$(cmake_cache_get ARM_EXECUTOR_RUNNER_STANDALONE)
+    local normalized
+    normalized=$(printf '%s' "${standalone}" | tr '[:lower:]' '[:upper:]')
+    if [[ "${normalized}" != "TRUE" && "${normalized}" != "ON" ]]; then
+        cat <&2
+Error: ${build_dir} is not a standalone arm_executor_runner build directory.
+Configure it via:
+  cmake -S ${runner_source_dir} -B ${build_dir} -DEXECUTORCH_ROOT=${et_root_dir} [...]
+and re-run run.sh.
+EOF
+        exit 1
+    fi
+}
+
+ensure_select_ops_list_setting() {
+    local expected="$1"
+    local cache_value
+    cache_value=$(cmake_cache_get EXECUTORCH_SELECT_OPS_LIST)
+    if [[ -z "${cache_value}" ]]; then
+        cat <&2
+Error: EXECUTORCH_SELECT_OPS_LIST is not configured in ${build_dir}.
+Reconfigure cmake -S ${runner_source_dir} -B ${build_dir} -DEXECUTORCH_SELECT_OPS_LIST=${expected}.
+EOF
+        exit 1
+    fi
+    if [[ "${cache_value}" != "${expected}" ]]; then
+        cat <&2
+Error: ${build_dir} was configured with EXECUTORCH_SELECT_OPS_LIST=${cache_value}, but run.sh requested ${expected}.
+Reconfigure cmake -S ${runner_source_dir} -B ${build_dir} -DEXECUTORCH_SELECT_OPS_LIST=${expected}, or omit --select_ops_list.
+EOF
+        exit 1
+    fi
+}
+
+require_cache_value() {
+    local key="$1"
+    local expected="$2"
+    local value
+    if ! cmake_cache_has_key "${key}"; then
+        echo "Error: ${key} not found in ${cmake_cache_file}. Reconfigure CMake with -D${key}=${expected}." >&2
+        exit 1
+    fi
+    value=$(cmake_cache_get "${key}")
+    if [[ "${value}" != "${expected}" ]]; then
+        echo "Error: ${key}=${value} in ${build_dir}. Reconfigure CMake with -D${key}=${expected} to use this run.sh invocation." >&2
+        exit 1
+    fi
+}
+
+require_cache_bool() {
+    local key="$1"
+    local expected="$2"
+    local value
+    value=$(cmake_cache_get "${key}")
+    if [[ -z "${value}" ]]; then
+        echo "Error: ${key} not found in ${cmake_cache_file}. Reconfigure CMake with -D${key}=${expected}." >&2
+        exit 1
+    fi
+    local value_upper
+    value_upper=$(printf '%s' "${value}" | tr '[:lower:]' '[:upper:]')
+    local expected_upper
+    expected_upper=$(printf '%s' "${expected}" | tr '[:lower:]' '[:upper:]')
+    if [[ "${value_upper}" != "${expected_upper}" ]]; then
+        echo "Error: ${key}=${value} in ${build_dir}. Reconfigure CMake with -D${key}=${expected} to use run.sh." >&2
+        exit 1
+    fi
+}
+
+is_cmake_false_value() {
+    local value_upper
+    value_upper=$(printf '%s' "$1" | tr '[:lower:]' '[:upper:]')
+    case "${value_upper}" in
+        ""|0|OFF|FALSE|NO|N|IGNORE|*-NOTFOUND)
+            return 0
+            ;;
+        *)
+            return 1
+            ;;
+    esac
+}
+
+ensure_pte_placement_setting() {
+    local cached_addr
+    cached_addr=$(cmake_cache_get ET_MODEL_PTE_ADDR)
+    if ! cmake_cache_has_key ET_MODEL_PTE_ADDR; then
+        echo "Error: ET_MODEL_PTE_ADDR not found in ${cmake_cache_file}. Reconfigure CMake for the requested --pte_placement=${pte_placement}." >&2
+        exit 1
+    fi
+    if [[ "${pte_placement}" == "elf" ]]; then
+        if ! is_cmake_false_value "${cached_addr}"; then
+            echo "Error: --pte_placement=elf requested, but ${build_dir} was configured with ET_MODEL_PTE_ADDR=${cached_addr}. Reconfigure CMake with -DET_MODEL_PTE_ADDR=." >&2
+            exit 1
+        fi
+        if ! cmake_cache_has_key ET_PTE_FILE_PATH; then
+            echo "Error: ET_PTE_FILE_PATH not found in ${cmake_cache_file}. Reconfigure CMake with -DET_PTE_FILE_PATH=." >&2
+            exit 1
+        fi
+        return
+    fi
+    if is_cmake_false_value "${cached_addr}"; then
+        echo "Error: --pte_placement=${pte_placement} requested, but ${build_dir} was configured for an embedded PTE. Reconfigure CMake with -DET_MODEL_PTE_ADDR=${pte_placement}, or use --pte_placement=elf." >&2
+        exit 1
+    fi
+    if [[ "${cached_addr}" != "${pte_placement}" ]]; then
+        echo "Error: --pte_placement=${pte_placement} requested, but ${build_dir} was configured with ET_MODEL_PTE_ADDR=${cached_addr}. Reconfigure CMake with -DET_MODEL_PTE_ADDR=${pte_placement}." >&2
+        exit 1
+    fi
+}
+
+get_parallel_jobs() {
+    if command -v nproc >/dev/null 2>&1; then
+        nproc
+    elif command -v sysctl >/dev/null 2>&1 && sysctl hw.logicalcpu >/dev/null 2>&1; then
+        sysctl -n hw.logicalcpu
+    elif command -v getconf >/dev/null 2>&1; then
+        getconf _NPROCESSORS_ONLN
+    elif [[ -n "${NUMBER_OF_PROCESSORS:-}" ]]; then
+        echo "${NUMBER_OF_PROCESSORS}"
+    else
+        echo 1
+    fi
+}
+
+build_runner_target() {
+    local cmake_target="$1"
+    local build_cmd=(cmake --build "${build_dir}" --target "${cmake_target}" --parallel "${parallel_jobs}")
+    if [[ "${multi_config}" == true ]]; then
+        build_cmd+=(--config "${build_type}")
+    fi
+    echo "[run.sh] Building target ${cmake_target} in ${build_dir}"
+    "${build_cmd[@]}"
+}
+
+locate_runner_binary() {
+    local binary_name="$1"
+    local candidates=()
+    if [[ "${multi_config}" == true ]]; then
+        candidates+=("${build_dir}/${build_type}/${binary_name}")
+        candidates+=("${build_dir}/examples/arm/executor_runner/${build_type}/${binary_name}")
+    fi
+    candidates+=("${build_dir}/${binary_name}")
+    candidates+=("${build_dir}/examples/arm/executor_runner/${binary_name}")
+    for candidate in "${candidates[@]}"; do
+        if [[ -f "${candidate}" ]]; then
+            echo "${candidate}"
+            return 0
+        fi
+    done
+    local found
+    found=$(find "${build_dir}" -name "${binary_name}" -type f 2>/dev/null | head -n 1 || true)
+    if [[ -n "${found}" ]]; then
+        echo "${found}"
+        return 0
+    fi
+    return 1
+}
+ensure_build_dir_ready() {
+    if [[ "${build_dir_initialized}" == true ]]; then
+        return
+    fi
+    if [[ -z "${build_dir}" ]]; then
+        echo "Error: build_dir is not set. Configure CMake first." >&2
+        exit 1
+    fi
+    build_dir=$(realpath "${build_dir}")
+    cmake_cache_file="${build_dir}/CMakeCache.txt"
+    if [[ ! -f ${cmake_cache_file} ]]; then
+        cat <&2
+Error: ${build_dir} does not contain a configured arm_executor_runner build (missing CMakeCache.txt).
+Run cmake -S ${runner_source_dir} -B ${build_dir} -DEXECUTORCH_ROOT=${et_root_dir} with the desired options first, then re-run run.sh.
+EOF
+        exit 1
+    fi
+    if [[ ${target} == *"vgf"* ]]; then
+        require_cache_bool EXECUTORCH_BUILD_VGF ON
+    else
+        ensure_runner_build_dir
+        require_cache_bool EXECUTORCH_BUILD_ARM_BAREMETAL ON
+        require_cache_bool EXECUTORCH_BAREMETAL_SKIP_INSTALL OFF
+        require_cache_value ETHOSU_TARGET_NPU_CONFIG "${target}"
+        require_cache_value TARGET_CPU "${target_cpu}"
+        require_cache_value SYSTEM_CONFIG "${system_config}"
+        require_cache_value MEMORY_MODE "${memory_mode}"
+        if [[ "${bundleio}" == true ]]; then
+            require_cache_bool ET_BUNDLE_IO ON
+        else
+            require_cache_bool ET_BUNDLE_IO OFF
+        fi
+        if [[ "${bundleio}" == true || "${build_with_etdump}" == true ]]; then
+            require_cache_bool EXECUTORCH_BUILD_DEVTOOLS ON
+        else
+            require_cache_bool EXECUTORCH_BUILD_DEVTOOLS OFF
+        fi
+        if [[ "${build_with_etdump}" == true ]]; then
+            require_cache_bool EXECUTORCH_ENABLE_EVENT_TRACER ON
+            require_cache_bool ET_DUMP_INTERMEDIATE_OUTPUTS ON
+        else
+            require_cache_bool EXECUTORCH_ENABLE_EVENT_TRACER OFF
+            require_cache_bool ET_DUMP_INTERMEDIATE_OUTPUTS OFF
+        fi
+    fi
+    if [[ ${target} != *"vgf"* ]]; then
+        ensure_select_ops_list_setting "${select_ops_list}"
+    fi
+    multi_config=false
+    if [[ -n "$(cmake_cache_get CMAKE_CONFIGURATION_TYPES)" ]]; then
+        multi_config=true
+    fi
+    parallel_jobs=$(get_parallel_jobs)
+    build_dir_initialized=true
+}
+
 #######
 ### Main
 #######
 if ! check_setup; then
     if [ "$scratch_dir_set" = false ] ; then
-        # check setup failed, no scratchdir given as parameter. trying to run setup.sh
-        if ${script_dir}/setup.sh; then
-            # and recheck setup. If this fails exit.
-            if ! check_setup; then
-                exit 1
-            fi
-        else
-            # setup.sh failed, it should print why
-            exit 1
-        fi
+	# check setup failed, no scratchdir given as parameter. trying to run setup.sh
+	if ${script_dir}/setup.sh; then
+	    # and recheck setup. If this fails exit.
+	    if ! check_setup; then
+		exit 1
+	    fi
+	else
+	    # setup.sh failed, it should print why
+	    exit 1
+	fi
     fi
 fi
 
+cd "${et_root_dir}"
+
+bundleio_flag=""
+etrecord_flag_template=""
+qdq_fusion_op_flag=""
+if [ "$build_with_etdump" = true ] ; then
+    etrecord_flag_template="--etrecord"
+fi
+
+if [ "$bundleio" = true ] ; then
+    bundleio_flag="--bundleio"
+fi
+
+if [ "$qdq_fusion_op" = true ] ; then
+    qdq_fusion_op_flag="--enable_qdq_fusion_pass"
+fi
+
+if [[ "${auto_configure}" == true ]]; then
+    set_default_build_dir_path
+else
+    if [[ -z "${build_dir}" ]]; then
+        echo "Error: --build-dir must not be empty." >&2
+        exit 1
+    fi
+    ensure_build_dir_ready
+fi
+
+stage_pte_into_cache() {
+    local new_pte="$1"
+    local cache_path
+    cache_path=$(cmake_cache_get ET_PTE_FILE_PATH)
+    if [[ -z "${cache_path}" ]]; then
+        cat <&2
+Error: --pte_placement=elf requires ET_PTE_FILE_PATH to be set when configuring CMake.
+Re-run cmake -S . -B ${build_dir} -DET_PTE_FILE_PATH=/absolute/path/to/model.pte (or use --pte_placement=).
+EOF
+        exit 1
+    fi
+    if [[ "${cache_path}" != /* ]]; then
+        cache_path="${build_dir}/${cache_path}"
+    fi
+    mkdir -p "$(dirname "${cache_path}")"
+    cp "${new_pte}" "${cache_path}"
+    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" || true)
+    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
     test_model=(
         "softmax"   # 0
@@ -302,7 +710,7 @@ for i in "${!test_model[@]}"; do
     printf "Running e2e flow for model '%s' with flags '%s'\n" "${model}" "${model_compiler_flags}"
     echo "--------------------------------------------------------------------------------"
 
-    cd $et_root_dir
+    cd "${et_root_dir}"
     # Remove path and file exetension to get model_short_name
     ext=${model##*.}
     model_short_name=$(basename -- "${model}" .$ext)
@@ -323,13 +731,14 @@ for i in "${!test_model[@]}"; do
         output_folder=${et_build_root}/${model_short_name}
     fi
 
+    local_fvp_pmu_flag=""
     if [ "$perf_overlay" = true ] ; then
         model_compiler_flags+="--enable_debug_mode tosa"
-        fvp_pmu_flag="--trace_file=${output_folder}/pmu_trace.gz"
+        local_fvp_pmu_flag="--trace_file=${output_folder}/pmu_trace.gz"
     fi
 
-    mkdir -p ${output_folder}
-    output_folder=$(realpath ${output_folder})
+    mkdir -p "${output_folder}"
+    output_folder=$(realpath "${output_folder}")
     pte_file="${output_folder}/${model_filename_ext}"
 
     # Remove old pte files
@@ -339,16 +748,17 @@ for i in "${!test_model[@]}"; do
         model_compiler_flags="${model_compiler_flags} --model_input=${model_input}"
     fi
 
-    ARM_AOT_CMD="python3 -m backends.arm.scripts.aot_arm_compiler --model_name=${model} --target=${target} ${model_compiler_flags} --intermediate=${output_folder} --output=${pte_file} --system_config=${system_config} --memory_mode=${memory_mode} $bundleio_flag ${etrecord_flag} --config=${config} $qdq_fusion_op_flag"
+    model_etrecord_flag="${etrecord_flag_template}"
+    ARM_AOT_CMD="python3 -m backends.arm.scripts.aot_arm_compiler --model_name=${model} --target=${target} ${model_compiler_flags} --intermediate=${output_folder} --output=${pte_file} --system_config=${system_config} --memory_mode=${memory_mode} $bundleio_flag ${model_etrecord_flag} --config=${config} $qdq_fusion_op_flag"
     echo "CALL ${ARM_AOT_CMD}" >&2
     ${ARM_AOT_CMD} 1>&2
 
-    pte_file=$(realpath ${pte_file})
+    pte_file=$(realpath "${pte_file}")
 
-    if [ "${etrecord_flag}" != "" ] ; then
+    if [ "${model_etrecord_flag}" != "" ] ; then
         etrecord_filename="${output_folder}/${model_filename}_etrecord.bin"
-        etrecord_filename=$(realpath ${etrecord_filename})
-        etrecord_flag="--etrecord=${etrecord_filename}"
+        etrecord_filename=$(realpath "${etrecord_filename}")
+        model_etrecord_flag="--etrecord=${etrecord_filename}"
     fi
 
     [[ -f ${pte_file} ]] || { >&2 echo "Failed to generate a pte file - ${pte_file}"; exit 1; }
@@ -357,6 +767,7 @@ for i in "${!test_model[@]}"; do
 
     if [[ ${target} == *"TOSA"*  ]]; then
         echo "Build for ${target} skip generating a .elf and running it"
+        continue
     elif [[ ${target} == cortex-m*  ]]; then
         # Cortex-M backend uses a shared semihosting executor_runner (built
         # by build_test_runner.sh) that loads the .bpte at runtime, rather
@@ -374,51 +785,54 @@ for i in "${!test_model[@]}"; do
         set +x
     elif [[ ${target} == *"vgf"*  ]]; then
         echo "Build and run for VKML, (target: ${target})"
-        set -x
-        backends/arm/scripts/build_executor_runner_vkml.sh --build_type=${build_type} \
-                                                           --extra_build_flags="${extra_build_flags}" \
-                                                           --output="${output_folder}" \
-                                                           ${bundleio_flag}
+        build_runner_target executor_runner
         if [ "$build_only" = false ] ; then
-            backends/arm/scripts/run_vkml.sh --model=${pte_file} --build_path=${output_folder}
+            backends/arm/scripts/run_vkml.sh --model=${pte_file} --build_path=${build_dir}
         fi
-        set +x
-
     else
-        # Build the application, the pte is imported as a header/c array or the address specified by --pte_placement
-        model_data=""
-        pte_file_or_mem="${pte_file}"
-        elf_file="${output_folder}/${model_filename}/cmake-out/arm_executor_runner"
-        if ! [[ ${pte_placement} == "elf" ]]; then
-            # Place PTE in memory specified by pte_placement
-            pte_file_or_mem="${pte_placement}"
-            model_data="--data=${pte_file}@${pte_placement}"
-            elf_file="${et_build_root}/${target}_${pte_placement}/cmake-out/arm_executor_runner"
+        if [[ "${auto_configure}" == true ]]; then
+            configure_runner_build_dir "${pte_file}"
         fi
+        ensure_build_dir_ready
+        ensure_pte_placement_setting
 
-        if [ "$specify_ethosu_scratch" = true ] && [[ ${target} =~ "ethos-u" ]]; then
-            scratch_size=$(get_ethosu_scratch_size "$pte_file")
-            if [ "$?" -eq 0 ] && [ -n "$scratch_size" ]; then
-                extra_build_flags="${extra_build_flags} -DET_ARM_BAREMETAL_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE=${scratch_size}"
+        model_data=""
+        if [[ ${pte_placement} == "elf" ]]; then
+            if [[ "${auto_configure}" == true ]]; then
+                staged_path=$(cmake_cache_get ET_PTE_FILE_PATH)
+                echo "ET_PTE_FILE_PATH payload: ${staged_path}"
             else
-                echo "WARNING: Failed to derive Ethos-U scratch size from ${pte_file}" >&2
+                staged_path=$(stage_pte_into_cache "${pte_file}")
+                echo "Updated ET_PTE_FILE_PATH payload: ${staged_path}"
             fi
+        else
+            model_data="--data=${pte_file}@${pte_placement}"
         fi
 
-        set -x
-        backends/arm/scripts/build_executor_runner.sh --et_build_root="${et_build_root}" --pte="${pte_file_or_mem}" --build_type=${build_type} --target=${target} --system_config=${system_config} --memory_mode=${memory_mode} ${bundleio_flag} ${et_dump_flag} --extra_build_flags="${extra_build_flags}" --ethosu_tools_dir="${arm_scratch_dir}" --toolchain="${toolchain}" --select_ops_list="${select_ops_list}"
-        if [ "$build_only" = false ] ; then
-            # Execute the executor_runner on FVP Simulator
+        configure_ethosu_scratch_if_requested "${pte_file}"
 
-            backends/arm/scripts/run_fvp.sh --elf=${elf_file} ${model_data} --target=$target ${etrecord_flag} ${fvp_pmu_flag}
+        build_runner_target arm_executor_runner
+        elf_file=$(locate_runner_binary arm_executor_runner) \
+            || { echo "Failed to locate arm_executor_runner in ${build_dir}." >&2; exit 1; }
+        if [ "$build_only" = false ] ; then
+            fvp_args=("--elf=${elf_file}" "--target=${target}")
+            if [[ -n "${model_data}" ]]; then
+                fvp_args+=("${model_data}")
+            fi
+            if [[ -n "${model_etrecord_flag}" ]]; then
+                fvp_args+=("${model_etrecord_flag}")
+            fi
+            if [[ -n "${local_fvp_pmu_flag}" ]]; then
+                fvp_args+=("${local_fvp_pmu_flag}")
+            fi
+            backends/arm/scripts/run_fvp.sh "${fvp_args[@]}"
         fi
-        set +x
     fi
 
     if [ "$model_explorer" = true ]; then
         perf_flags=""
         if [ "$perf_overlay" = true ]; then
-            perf_flags+="--trace ${output_folder}/pmu_trace.gz --tables ${output_folder}/output/out_debug.xml"
+            perf_flags+=" --trace ${output_folder}/pmu_trace.gz --tables ${output_folder}/output/out_debug.xml"
         fi
 
         visualization_file=""
diff --git a/examples/arm/super_resolution_example_vgf/README.md b/examples/arm/super_resolution_example_vgf/README.md
new file mode 100644
index 00000000000..cefe5c55f22
--- /dev/null
+++ b/examples/arm/super_resolution_example_vgf/README.md
@@ -0,0 +1,19 @@
+# Swin2SR Super-Resolution Example Application (VGF)
+
+This example shows how to export a Swin2SR image super-resolution model for the
+Arm VGF backend and run it on host using the generic `executor_runner` binary.
+It is a host-only workflow; a device-specific VGF runtime application is out of
+scope here.
+
+## Layout
+
+- `model_export/prepare_demo_assets.py` — Creates a deterministic text-heavy
+  demo input plus small LR/HR calibration and evaluation sets from a repo-local
+  screenshot.
+- `model_export/README.md` — Dataset-backed FP/INT8 export, PTQ
+  calibration and evaluation, and `.pte` generation.
+- `runtime/README.md` — Running the exported `.pte` on host using
+  `executor_runner` and converting the output tensor back into an image.
+
+Use `examples/arm/image_classification_example_vgf` for the image
+classification flow.
diff --git a/examples/arm/super_resolution_example_vgf/model_export/README.md b/examples/arm/super_resolution_example_vgf/model_export/README.md
new file mode 100644
index 00000000000..9489327ed1b
--- /dev/null
+++ b/examples/arm/super_resolution_example_vgf/model_export/README.md
@@ -0,0 +1,116 @@
+# Swin2SR VGF Export
+
+This example provides two scripts:
+
+- `prepare_demo_assets.py` — Creates a deterministic local demo dataset from a
+  repo-local screenshot so the export and runtime steps can be reproduced
+  without an external SR dataset.
+- `export_super_resolution.py` — Loads a checkpoint, applies optional
+  post-training quantization, evaluates PSNR/SSIM on paired LR/HR samples, and
+  exports a VGF-ready ExecuTorch program for host execution.
+
+## Requirements
+
+- Python 3.10+ with `executorch` and the dependencies in
+  `examples/arm/super_resolution_example_vgf/requirements.txt`.
+- ML SDK dependencies installed through `examples/arm/setup.sh
+  --disable-ethos-u-deps --enable-mlsdk-deps`.
+
+## Quick demo assets
+
+To generate the text-heavy demo crop used by the runtime walkthrough, along
+with the small LR/HR directories used for INT8 calibration and evaluation, run:
+
+```bash
+python examples/arm/super_resolution_example_vgf/model_export/prepare_demo_assets.py \
+  --output-dir ./demo_assets
+```
+
+This writes:
+
+```text
+demo_assets/
+  calibration/hr/
+  calibration/lr/
+  eval/hr/
+  eval/lr/
+  runtime/demo_hr_128.png
+  runtime/demo_lr_64.png
+  metadata.json
+```
+
+The export flow expects paired RGB image directories with matching relative
+paths on the LR and HR sides when evaluation metrics are requested. For each
+pair, the HR image must be exactly `upscale x` larger than the LR image. The
+exporter crops LR inputs to `--input-height` x `--input-width` and crops the HR
+target to the matching scaled patch.
+
+## Export and evaluate for VGF
+
+The concrete quick-demo commands below use the pinned revision currently cached
+for `caidas/swin2SR-classical-sr-x2-64`:
+`cee1c923c6a37361c6e5650b65dcf4be821e5d52`.
+
+### FP export
+
+```bash
+python examples/arm/super_resolution_example_vgf/model_export/export_super_resolution.py \
+  --model-name swin2sr \
+  --checkpoint caidas/swin2SR-classical-sr-x2-64 \
+  --checkpoint-revision cee1c923c6a37361c6e5650b65dcf4be821e5d52 \
+  --input-height 64 \
+  --input-width 64 \
+  --quantization-mode none \
+  --eval-lr-dir ./demo_assets/eval/lr \
+  --eval-hr-dir ./demo_assets/eval/hr \
+  --num-eval-samples 2 \
+  --output-path ./demo_assets/swin2sr_x2_vgf_fp.pte
+```
+
+### INT8 export
+
+```bash
+python examples/arm/super_resolution_example_vgf/model_export/export_super_resolution.py \
+  --model-name swin2sr \
+  --checkpoint caidas/swin2SR-classical-sr-x2-64 \
+  --checkpoint-revision cee1c923c6a37361c6e5650b65dcf4be821e5d52 \
+  --input-height 64 \
+  --input-width 64 \
+  --quantization-mode int8 \
+  --calibration-lr-dir ./demo_assets/calibration/lr \
+  --eval-lr-dir ./demo_assets/eval/lr \
+  --eval-hr-dir ./demo_assets/eval/hr \
+  --num-calibration-samples 4 \
+  --num-eval-samples 2 \
+  --output-path ./demo_assets/swin2sr_x2_vgf_int8.pte
+```
+
+For FP export, set `--quantization-mode none`. INT8 export requires
+`--calibration-lr-dir`; the exporter no longer falls back to random calibration
+inputs. The exporter first tries installed ExecuTorch quantized kernels and
+then local build outputs such as `cmake-out/kernels/quantized` or
+`arm_test/*/kernels/quantized` to register the quantized out-variant ops needed
+by `to_executorch()`.
+
+When `--eval-lr-dir` and `--eval-hr-dir` are provided, the exporter compares
+the exported program module against the paired HR images and writes PSNR/SSIM
+metrics. The quick-demo dataset is intentionally tiny, so these metrics are a
+smoke signal for gross quality regressions rather than a benchmark target; use
+a larger paired validation set when setting release-quality thresholds.
+
+In the OOTB smoke flow, FP export over the generated 2-sample eval set produced
+approximately PSNR 34.85 / SSIM 0.994, while INT8 PTQ produced approximately
+PSNR 22.71 / SSIM 0.870. The INT8 drop is expected; these numbers are included
+only as a smoke-test reference for the generated demo assets.
+
+## Output artifacts
+
+For an export path such as `./swin2sr_x2_vgf_int8.pte`, the exporter writes:
+
+- `swin2sr_x2_vgf_int8.pte` — The ExecuTorch program.
+- `swin2sr_x2_vgf_int8.json` — Static input/output metadata consumed by the
+  runtime helper.
+- `swin2sr_x2_vgf_int8_delegation.txt` — A summary of delegated and
+  non-delegated operators.
+- `swin2sr_x2_vgf_int8_metrics.json` — Optional PSNR/SSIM evaluation metrics
+  when `--eval-lr-dir` and `--eval-hr-dir` are provided.
diff --git a/examples/arm/super_resolution_example_vgf/model_export/common.py b/examples/arm/super_resolution_example_vgf/model_export/common.py
new file mode 100644
index 00000000000..617d5575402
--- /dev/null
+++ b/examples/arm/super_resolution_example_vgf/model_export/common.py
@@ -0,0 +1,452 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterable
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torch.utils.data import Dataset
+from transformers import Swin2SRForImageSuperResolution
+
+IMAGE_SUFFIXES = {".bmp", ".jpeg", ".jpg", ".png", ".webp"}
+SUPPORTED_MODELS = ("swin2sr",)
+
+
+class Swin2SRWrapper(torch.nn.Module):
+    def __init__(self, model: Swin2SRForImageSuperResolution):
+        super().__init__()
+        self.model = model
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        return self.model(pixel_values=pixel_values, return_dict=True).reconstruction
+
+
+@dataclass(frozen=True)
+class SuperResolutionModelBundle:
+    model_name: str
+    model: torch.nn.Module
+    example_inputs: tuple[torch.Tensor]
+    input_shape: tuple[int, ...]
+    output_shape: tuple[int, ...]
+    input_dtype: str
+    output_dtype: str
+    upscale: int
+    window_size: int
+
+
+def write_json(path: str | Path, payload: dict[str, Any]) -> Path:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+    return path
+
+
+def list_image_paths(path: str | Path) -> list[Path]:
+    path = Path(path)
+    if path.is_file():
+        if path.suffix.lower() not in IMAGE_SUFFIXES:
+            raise ValueError(f"Unsupported image file type: {path}")
+        return [path]
+
+    if not path.is_dir():
+        raise ValueError(f"Image path does not exist: {path}")
+
+    image_paths = sorted(
+        candidate
+        for candidate in path.rglob("*")
+        if candidate.is_file() and candidate.suffix.lower() in IMAGE_SUFFIXES
+    )
+    if not image_paths:
+        raise ValueError(f"No supported images found in: {path}")
+    return image_paths
+
+
+def load_image_tensor(image_path: str | Path) -> torch.Tensor:
+    image = Image.open(image_path).convert("RGB")
+    image_np = np.asarray(image, dtype=np.float32) / 255.0
+    return torch.from_numpy(image_np).permute(2, 0, 1).contiguous().clone()
+
+
+def crop_input_tensor(
+    image: torch.Tensor,
+    input_height: int,
+    input_width: int,
+    crop_mode: str,
+) -> torch.Tensor:
+    height, width = image.shape[1:]
+    if height < input_height or width < input_width:
+        raise ValueError(
+            "Image tensor is smaller than the requested crop size: "
+            f"{tuple(image.shape)} vs {(input_height, input_width)}."
+        )
+
+    if height == input_height and width == input_width:
+        return image
+
+    max_top = height - input_height
+    max_left = width - input_width
+    if crop_mode == "random":
+        top = int(torch.randint(max_top + 1, ()).item())
+        left = int(torch.randint(max_left + 1, ()).item())
+    elif crop_mode == "center":
+        top = max_top // 2
+        left = max_left // 2
+    else:
+        raise ValueError(f"Unsupported crop mode: {crop_mode}")
+
+    return image[:, top : top + input_height, left : left + input_width].contiguous()
+
+
+def load_calibration_inputs(
+    image_path: str | Path,
+    input_height: int,
+    input_width: int,
+    max_samples: int,
+) -> list[tuple[torch.Tensor]]:
+    if max_samples <= 0:
+        raise ValueError("max_samples must be positive.")
+
+    calibration_inputs = []
+    for candidate in list_image_paths(image_path)[:max_samples]:
+        image = load_image_tensor(candidate)
+        image = crop_input_tensor(image, input_height, input_width, crop_mode="center")
+        calibration_inputs.append((image.unsqueeze(0),))
+    return calibration_inputs
+
+
+def _collect_image_map(root: Path) -> dict[Path, Path]:
+    image_map = {
+        candidate.relative_to(root): candidate for candidate in list_image_paths(root)
+    }
+    if not image_map:
+        raise ValueError(f"No supported images found in: {root}")
+    return image_map
+
+
+def paired_image_paths(
+    lr_dir: str | Path,
+    hr_dir: str | Path,
+    max_samples: int | None = None,
+) -> list[tuple[Path, Path]]:
+    lr_root = Path(lr_dir)
+    hr_root = Path(hr_dir)
+    lr_map = _collect_image_map(lr_root)
+    hr_map = _collect_image_map(hr_root)
+
+    lr_keys = set(lr_map)
+    hr_keys = set(hr_map)
+    if lr_keys != hr_keys:
+        missing_lr = sorted(str(key) for key in hr_keys - lr_keys)
+        missing_hr = sorted(str(key) for key in lr_keys - hr_keys)
+        details = []
+        if missing_lr:
+            details.append(f"Missing LR files: {missing_lr[:5]}")
+        if missing_hr:
+            details.append(f"Missing HR files: {missing_hr[:5]}")
+        raise ValueError(
+            "LR/HR directories do not contain matching files. " + " ".join(details)
+        )
+
+    pairs = [(lr_map[key], hr_map[key]) for key in sorted(lr_map)]
+    if max_samples is not None:
+        if max_samples <= 0:
+            raise ValueError("max_samples must be positive.")
+        pairs = pairs[:max_samples]
+    if not pairs:
+        raise ValueError("No paired super-resolution samples were found.")
+    return pairs
+
+
+def _validate_pair_shapes(
+    lr_tensor: torch.Tensor,
+    hr_tensor: torch.Tensor,
+    upscale: int,
+) -> None:
+    if lr_tensor.shape[0] != 3 or hr_tensor.shape[0] != 3:
+        raise ValueError("Only RGB LR/HR image pairs are supported.")
+
+    expected_hr_shape = (
+        hr_tensor.shape[0],
+        lr_tensor.shape[1] * upscale,
+        lr_tensor.shape[2] * upscale,
+    )
+    if tuple(hr_tensor.shape) != expected_hr_shape:
+        raise ValueError(
+            "HR image shape does not match the LR image and upscale factor: "
+            f"LR {tuple(lr_tensor.shape)}, HR {tuple(hr_tensor.shape)}, upscale {upscale}."
+        )
+
+
+def crop_super_resolution_pair(
+    lr_tensor: torch.Tensor,
+    hr_tensor: torch.Tensor,
+    input_height: int,
+    input_width: int,
+    upscale: int,
+    crop_mode: str,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    _validate_pair_shapes(lr_tensor, hr_tensor, upscale)
+
+    lr_height, lr_width = lr_tensor.shape[1:]
+    if lr_height < input_height or lr_width < input_width:
+        raise ValueError(
+            "LR image is smaller than the requested crop size: "
+            f"{tuple(lr_tensor.shape)} vs {(input_height, input_width)}."
+        )
+
+    if lr_height == input_height and lr_width == input_width:
+        return lr_tensor.contiguous(), hr_tensor.contiguous()
+
+    max_top = lr_height - input_height
+    max_left = lr_width - input_width
+    if crop_mode == "random":
+        top = int(torch.randint(max_top + 1, ()).item())
+        left = int(torch.randint(max_left + 1, ()).item())
+    elif crop_mode == "center":
+        top = max_top // 2
+        left = max_left // 2
+    else:
+        raise ValueError(f"Unsupported crop mode: {crop_mode}")
+
+    hr_top = top * upscale
+    hr_left = left * upscale
+    return (
+        lr_tensor[:, top : top + input_height, left : left + input_width].contiguous(),
+        hr_tensor[
+            :,
+            hr_top : hr_top + input_height * upscale,
+            hr_left : hr_left + input_width * upscale,
+        ].contiguous(),
+    )
+
+
+class PairedSuperResolutionDataset(Dataset[tuple[torch.Tensor, torch.Tensor]]):
+    def __init__(
+        self,
+        lr_dir: str | Path,
+        hr_dir: str | Path,
+        input_height: int,
+        input_width: int,
+        upscale: int,
+        crop_mode: str,
+        max_samples: int | None = None,
+    ):
+        self.pairs = paired_image_paths(lr_dir, hr_dir, max_samples=max_samples)
+        self.input_height = input_height
+        self.input_width = input_width
+        self.upscale = upscale
+        self.crop_mode = crop_mode
+
+    def __len__(self) -> int:
+        return len(self.pairs)
+
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
+        lr_path, hr_path = self.pairs[idx]
+        lr_tensor = load_image_tensor(lr_path)
+        hr_tensor = load_image_tensor(hr_path)
+        return crop_super_resolution_pair(
+            lr_tensor,
+            hr_tensor,
+            self.input_height,
+            self.input_width,
+            self.upscale,
+            self.crop_mode,
+        )
+
+
+def _load_checkpointed_swin2sr(
+    checkpoint: str,
+    checkpoint_revision: str | None,
+    local_files_only: bool,
+) -> Swin2SRWrapper:
+    is_local_checkpoint = Path(checkpoint).expanduser().exists()
+    if checkpoint_revision is None and not (local_files_only or is_local_checkpoint):
+        raise ValueError(
+            "--checkpoint-revision is required when --checkpoint is a remote Hugging Face model id."
+        )
+    if checkpoint_revision is None:
+        model = Swin2SRForImageSuperResolution.from_pretrained(  # nosec B615
+            checkpoint,
+            local_files_only=local_files_only,
+        ).eval()
+    else:
+        model = Swin2SRForImageSuperResolution.from_pretrained(  # nosec B615
+            checkpoint,
+            revision=checkpoint_revision,
+            local_files_only=local_files_only,
+        ).eval()
+    return Swin2SRWrapper(model)
+
+
+def create_model_bundle(
+    model_name: str,
+    input_height: int,
+    input_width: int,
+    checkpoint: str | None = None,
+    checkpoint_revision: str | None = None,
+    local_files_only: bool = False,
+) -> SuperResolutionModelBundle:
+    if model_name not in SUPPORTED_MODELS:
+        raise ValueError(
+            f"Unsupported model '{model_name}'. Supported models: {SUPPORTED_MODELS}"
+        )
+    if input_height <= 0 or input_width <= 0:
+        raise ValueError("Input dimensions must be positive.")
+
+    if checkpoint is None:
+        raise ValueError("--checkpoint is required when --model-name=swin2sr.")
+    model = _load_checkpointed_swin2sr(
+        checkpoint,
+        checkpoint_revision,
+        local_files_only,
+    )
+
+    example_input = torch.rand((1, 3, input_height, input_width), dtype=torch.float32)
+    with torch.no_grad():
+        example_output = model(example_input)
+
+    if example_output.dim() != 4:
+        raise ValueError(
+            f"Expected a 4D reconstruction tensor, got {tuple(example_output.shape)}."
+        )
+
+    model_impl = model.model
+    return SuperResolutionModelBundle(
+        model_name=model_name,
+        model=model,
+        example_inputs=(example_input,),
+        input_shape=tuple(example_input.shape),
+        output_shape=tuple(example_output.shape),
+        input_dtype=str(example_input.dtype).replace("torch.", ""),
+        output_dtype=str(example_output.dtype).replace("torch.", ""),
+        upscale=int(model_impl.config.upscale),
+        window_size=int(model_impl.config.window_size),
+    )
+
+
+def _gaussian_window(
+    channels: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    height: int,
+    width: int,
+) -> torch.Tensor:
+    window_size = min(11, height, width)
+    if window_size % 2 == 0:
+        window_size -= 1
+    window_size = max(window_size, 1)
+    sigma = max(window_size / 6.0, 1e-3)
+
+    coords = torch.arange(window_size, device=device, dtype=dtype) - window_size // 2
+    kernel_1d = torch.exp(-(coords**2) / (2 * sigma**2))
+    kernel_1d = kernel_1d / kernel_1d.sum()
+    kernel_2d = torch.outer(kernel_1d, kernel_1d)
+    return kernel_2d.expand(channels, 1, window_size, window_size).contiguous()
+
+
+def batch_psnr(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    mse = F.mse_loss(prediction, target, reduction="none").mean(dim=(1, 2, 3))
+    return 10.0 * torch.log10(1.0 / mse.clamp_min(1e-12))
+
+
+def batch_ssim(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    channels = prediction.shape[1]
+    kernel = _gaussian_window(
+        channels,
+        prediction.device,
+        prediction.dtype,
+        prediction.shape[-2],
+        prediction.shape[-1],
+    )
+    padding = kernel.shape[-1] // 2
+    c1 = 0.01**2
+    c2 = 0.03**2
+
+    mu_x = F.conv2d(prediction, kernel, padding=padding, groups=channels)
+    mu_y = F.conv2d(target, kernel, padding=padding, groups=channels)
+
+    mu_x_sq = mu_x.pow(2)
+    mu_y_sq = mu_y.pow(2)
+    mu_xy = mu_x * mu_y
+
+    sigma_x_sq = (
+        F.conv2d(prediction * prediction, kernel, padding=padding, groups=channels)
+        - mu_x_sq
+    )
+    sigma_y_sq = (
+        F.conv2d(target * target, kernel, padding=padding, groups=channels) - mu_y_sq
+    )
+    sigma_xy = (
+        F.conv2d(prediction * target, kernel, padding=padding, groups=channels) - mu_xy
+    )
+
+    ssim_map = ((2 * mu_xy + c1) * (2 * sigma_xy + c2)) / (
+        (mu_x_sq + mu_y_sq + c1) * (sigma_x_sq + sigma_y_sq + c2)
+    )
+    return ssim_map.mean(dim=(1, 2, 3))
+
+
+def _model_reconstruction(
+    model: torch.nn.Module,
+    input_tensor: torch.Tensor,
+) -> torch.Tensor:
+    output = model(input_tensor)
+    if hasattr(output, "reconstruction"):
+        output = output.reconstruction
+    if not isinstance(output, torch.Tensor):
+        raise TypeError(f"Expected tensor output, got {type(output)}")
+    return output
+
+
+def evaluate_super_resolution_model(
+    model: torch.nn.Module,
+    dataset: Iterable[tuple[torch.Tensor, torch.Tensor]],
+    device: torch.device,
+) -> dict[str, float]:
+    if hasattr(model, "eval"):
+        try:
+            model.eval()
+        except NotImplementedError:
+            pass
+
+    total_l1 = 0.0
+    total_psnr = 0.0
+    total_ssim = 0.0
+    total_examples = 0
+
+    with torch.no_grad():
+        for lr_tensor, hr_tensor in dataset:
+            lr_tensor = lr_tensor.to(device)
+            hr_tensor = hr_tensor.to(device)
+            prediction = _model_reconstruction(model, lr_tensor).clamp(0.0, 1.0)
+
+            l1_values = F.l1_loss(prediction, hr_tensor, reduction="none").mean(
+                dim=(1, 2, 3)
+            )
+            psnr_values = batch_psnr(prediction, hr_tensor)
+            ssim_values = batch_ssim(prediction, hr_tensor)
+
+            batch_size = lr_tensor.shape[0]
+            total_l1 += l1_values.sum().item()
+            total_psnr += psnr_values.sum().item()
+            total_ssim += ssim_values.sum().item()
+            total_examples += batch_size
+
+    if total_examples == 0:
+        raise ValueError("Evaluation dataset is empty.")
+
+    return {
+        "l1": total_l1 / total_examples,
+        "psnr": total_psnr / total_examples,
+        "ssim": total_ssim / total_examples,
+        "num_samples": float(total_examples),
+    }
diff --git a/examples/arm/super_resolution_example_vgf/model_export/export_super_resolution.py b/examples/arm/super_resolution_example_vgf/model_export/export_super_resolution.py
new file mode 100644
index 00000000000..a8d3d7dfecf
--- /dev/null
+++ b/examples/arm/super_resolution_example_vgf/model_export/export_super_resolution.py
@@ -0,0 +1,366 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+
+from executorch.backends.arm.quantizer import (
+    get_symmetric_quantization_config,
+    VgfQuantizer,
+)
+from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner
+from executorch.devtools.backend_debug import get_delegation_info
+from executorch.exir import (
+    EdgeCompileConfig,
+    ExecutorchBackendConfig,
+    to_edge_transform_and_lower,
+)
+from executorch.extension.export_util.utils import save_pte_program
+from torch.utils.data import DataLoader
+from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
+
+# Keep script-compatible imports without requiring package execution.
+if __package__ is None or __package__ == "":
+    import sys
+
+    sys.path.append(str(Path(__file__).resolve().parent))
+    from common import (  # type: ignore[import-not-found, no-redef]
+        create_model_bundle,
+        evaluate_super_resolution_model,
+        load_calibration_inputs,
+        PairedSuperResolutionDataset,
+        SUPPORTED_MODELS,
+        write_json,
+    )
+else:
+    from .common import (
+        create_model_bundle,
+        evaluate_super_resolution_model,
+        load_calibration_inputs,
+        PairedSuperResolutionDataset,
+        SUPPORTED_MODELS,
+        write_json,
+    )
+
+CALIBRATION_MAX_SAMPLES = 1000
+EVAL_MAX_SAMPLES = 1000
+
+
+def has_quantized_out_variants() -> bool:
+    try:
+        _ = torch.ops.quantized_decomposed.quantize_per_tensor.out
+        _ = torch.ops.quantized_decomposed.dequantize_per_tensor.out
+        return True
+    except AttributeError:
+        return False
+
+
+def ensure_quantized_ops_loaded() -> Path | None:
+    if has_quantized_out_variants():
+        return None
+
+    quantized_kernels_available = False
+    try:
+        import executorch.kernels.quantized  # noqa: F401
+    except ImportError:
+        quantized_kernels_available = False
+    else:
+        quantized_kernels_available = True
+
+    if quantized_kernels_available and has_quantized_out_variants():
+        return None
+
+    repo_root = Path(__file__).resolve().parents[4]
+    search_patterns = (
+        "cmake-out/kernels/quantized/libquantized_ops_aot_lib.*",
+        "arm_test/*/kernels/quantized/libquantized_ops_aot_lib.*",
+    )
+    for pattern in search_patterns:
+        for candidate in sorted(repo_root.glob(pattern)):
+            if not candidate.is_file():
+                continue
+            torch.ops.load_library(str(candidate))
+            if has_quantized_out_variants():
+                return candidate
+
+    raise RuntimeError(
+        "INT8 export requires the quantized ops out-variant library. "
+        "Build or install ExecuTorch quantized kernels so that "
+        "`quantized_decomposed::quantize_per_tensor.out` and "
+        "`quantized_decomposed::dequantize_per_tensor.out` are available."
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Export a Swin2SR model for VGF.")
+    parser.add_argument(
+        "--model-name",
+        choices=SUPPORTED_MODELS,
+        default="swin2sr",
+        help="Model profile to export.",
+    )
+    parser.add_argument(
+        "--checkpoint",
+        default=None,
+        help="Checkpoint directory or Hugging Face model id.",
+    )
+    parser.add_argument(
+        "--local-files-only",
+        action="store_true",
+        help="Resolve Hugging Face assets from local cache only.",
+    )
+    parser.add_argument(
+        "--checkpoint-revision",
+        default=None,
+        help=(
+            "Pinned Hugging Face revision to use when --checkpoint points to a model id. "
+            "Required for remote checkpoints; ignored for local checkpoint paths."
+        ),
+    )
+    parser.add_argument(
+        "--input-height",
+        type=int,
+        default=64,
+        help="Static low-resolution input height used for export.",
+    )
+    parser.add_argument(
+        "--input-width",
+        type=int,
+        default=64,
+        help="Static low-resolution input width used for export.",
+    )
+    parser.add_argument(
+        "--output-path",
+        type=str,
+        required=True,
+        help="Destination .pte path.",
+    )
+    parser.add_argument(
+        "--quantization-mode",
+        choices=("none", "int8"),
+        default="none",
+        help="Quantization mode used before lowering to VGF.",
+    )
+    parser.add_argument(
+        "--calibration-lr-dir",
+        default=None,
+        help=(
+            "Directory of low-resolution images used for PTQ calibration. "
+            "Required when --quantization-mode=int8."
+        ),
+    )
+    parser.add_argument(
+        "--num-calibration-samples",
+        type=int,
+        default=32,
+        help="Maximum number of calibration images to use.",
+    )
+    parser.add_argument(
+        "--eval-lr-dir",
+        default=None,
+        help="Optional directory of low-resolution evaluation images.",
+    )
+    parser.add_argument(
+        "--eval-hr-dir",
+        default=None,
+        help="Optional directory of high-resolution evaluation targets.",
+    )
+    parser.add_argument(
+        "--num-eval-samples",
+        type=int,
+        default=100,
+        help="Maximum number of evaluation image pairs to use.",
+    )
+    parser.add_argument(
+        "--artifact-dir",
+        default=None,
+        help="Optional directory for intermediate VGF/TOSA artifacts.",
+    )
+    return parser.parse_args()
+
+
+def quantize_model(
+    model: torch.nn.Module,
+    quantizer: VgfQuantizer,
+    example_inputs: tuple[torch.Tensor],
+    calibration_samples: list[tuple[torch.Tensor]],
+) -> torch.export.ExportedProgram:
+    exported_program = torch.export.export(model, example_inputs)
+    graph_module = exported_program.module(check_guards=False)
+
+    prepared = prepare_pt2e(graph_module, quantizer)
+    for sample in calibration_samples:
+        prepared(*sample)
+
+    quantized_model = convert_pt2e(prepared)
+    return torch.export.export(quantized_model, example_inputs)
+
+
+def write_delegation_report(edge_program_manager, report_path: Path) -> None:
+    delegation_info = get_delegation_info(
+        edge_program_manager.exported_program().graph_module
+    )
+    report_path.write_text(delegation_info.get_summary() + "\n")
+
+
+def maybe_make_eval_loader(
+    eval_lr_dir: str | None,
+    eval_hr_dir: str | None,
+    input_height: int,
+    input_width: int,
+    upscale: int,
+    num_eval_samples: int,
+) -> DataLoader[tuple[torch.Tensor, torch.Tensor]] | None:
+    if eval_lr_dir is None and eval_hr_dir is None:
+        return None
+    if (eval_lr_dir is None) != (eval_hr_dir is None):
+        raise ValueError("--eval-lr-dir and --eval-hr-dir must be provided together.")
+
+    eval_dataset = PairedSuperResolutionDataset(
+        eval_lr_dir,
+        eval_hr_dir,
+        input_height,
+        input_width,
+        upscale,
+        crop_mode="center",
+        max_samples=min(num_eval_samples, EVAL_MAX_SAMPLES),
+    )
+    return DataLoader(eval_dataset, batch_size=1, shuffle=False, num_workers=0)
+
+
+def main() -> None:
+    args = parse_args()
+
+    output_path = Path(args.output_path).resolve()
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    metadata_path = output_path.with_suffix(".json")
+    delegation_path = output_path.with_name(f"{output_path.stem}_delegation.txt")
+    metrics_path = output_path.with_name(f"{output_path.stem}_metrics.json")
+
+    bundle = create_model_bundle(
+        model_name=args.model_name,
+        input_height=args.input_height,
+        input_width=args.input_width,
+        checkpoint=args.checkpoint,
+        checkpoint_revision=args.checkpoint_revision,
+        local_files_only=args.local_files_only,
+    )
+
+    quantize = args.quantization_mode != "none"
+    compile_spec = VgfCompileSpec("TOSA-1.0+INT" if quantize else "TOSA-1.0+FP")
+    if args.artifact_dir is not None:
+        artifact_dir = Path(args.artifact_dir).resolve()
+        artifact_dir.mkdir(parents=True, exist_ok=True)
+        compile_spec.dump_intermediate_artifacts_to(str(artifact_dir))
+
+    calibration_samples: list[tuple[torch.Tensor]] = []
+    if quantize:
+        quantized_ops_library = ensure_quantized_ops_loaded()
+        if quantized_ops_library is not None:
+            print(f"Loaded quantized ops library from {quantized_ops_library}")
+        if args.calibration_lr_dir is None:
+            raise ValueError(
+                "--calibration-lr-dir is required when --quantization-mode=int8."
+            )
+        calibration_samples = load_calibration_inputs(
+            args.calibration_lr_dir,
+            args.input_height,
+            args.input_width,
+            min(args.num_calibration_samples, CALIBRATION_MAX_SAMPLES),
+        )
+
+    exported_program: torch.export.ExportedProgram
+    if quantize:
+        quantizer = VgfQuantizer(compile_spec)
+        quantizer.set_global(get_symmetric_quantization_config(is_per_channel=True))
+        exported_program = quantize_model(
+            bundle.model,
+            quantizer,
+            bundle.example_inputs,
+            calibration_samples,
+        )
+    else:
+        exported_program = torch.export.export(bundle.model, bundle.example_inputs)
+
+    eval_loader = maybe_make_eval_loader(
+        args.eval_lr_dir,
+        args.eval_hr_dir,
+        args.input_height,
+        args.input_width,
+        bundle.upscale,
+        args.num_eval_samples,
+    )
+    evaluation_metrics: dict[str, float] | None = None
+    if eval_loader is not None:
+        eval_module = exported_program.module(check_guards=False)
+        metrics = evaluate_super_resolution_model(
+            eval_module,
+            eval_loader,
+            torch.device("cpu"),
+        )
+        evaluation_metrics = metrics
+        write_json(metrics_path, metrics)
+        print(
+            "Evaluation metrics: "
+            f"l1={metrics['l1']:.6f} "
+            f"psnr={metrics['psnr']:.4f} "
+            f"ssim={metrics['ssim']:.4f}"
+        )
+
+    partitioner = VgfPartitioner(compile_spec)
+    edge_program_manager = to_edge_transform_and_lower(
+        programs=exported_program,
+        partitioner=[partitioner],
+        compile_config=EdgeCompileConfig(_check_ir_validity=False),
+    )
+    write_delegation_report(edge_program_manager, delegation_path)
+
+    executorch_program_manager = edge_program_manager.to_executorch(
+        config=ExecutorchBackendConfig(extract_delegate_segments=False)
+    )
+    save_pte_program(
+        executorch_program_manager,
+        str(output_path),
+        output_dir=str(output_path.parent),
+    )
+    if not output_path.is_file():
+        raise RuntimeError(f"Expected exported model at {output_path}")
+
+    write_json(
+        metadata_path,
+        {
+            "model_name": bundle.model_name,
+            "checkpoint": args.checkpoint,
+            "checkpoint_revision": args.checkpoint_revision,
+            "input_shape": list(bundle.input_shape),
+            "output_shape": list(bundle.output_shape),
+            "input_dtype": bundle.input_dtype,
+            "output_dtype": bundle.output_dtype,
+            "num_outputs": 1,
+            "upscale": bundle.upscale,
+            "window_size": bundle.window_size,
+            "quantization_mode": args.quantization_mode,
+            "num_calibration_samples": len(calibration_samples),
+            "num_eval_samples": (
+                int(evaluation_metrics["num_samples"])
+                if evaluation_metrics is not None
+                else 0
+            ),
+        },
+    )
+
+    print(f"Exported model saved to {output_path}")
+    print(f"Metadata saved to {metadata_path}")
+    print(f"Delegation summary saved to {delegation_path}")
+    if evaluation_metrics is not None:
+        print(f"Evaluation metrics saved to {metrics_path}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/arm/super_resolution_example_vgf/model_export/prepare_demo_assets.py b/examples/arm/super_resolution_example_vgf/model_export/prepare_demo_assets.py
new file mode 100644
index 00000000000..130b72b80d8
--- /dev/null
+++ b/examples/arm/super_resolution_example_vgf/model_export/prepare_demo_assets.py
@@ -0,0 +1,149 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+from PIL import Image
+
+RUNTIME_DEMO_CROP = ("demo", 60, 420)
+CALIBRATION_CROPS = (
+    ("calib_0", 60, 280),
+    ("calib_1", 60, 420),
+    ("calib_2", 60, 560),
+    ("calib_3", 60, 700),
+)
+EVAL_CROPS = (
+    ("eval_0", 60, 420),
+    ("eval_1", 60, 700),
+)
+DEFAULT_SOURCE_IMAGE = (
+    Path(__file__).resolve().parents[4] / "docs/source/_static/img/ios_demo_app.jpg"
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Prepare a reproducible Swin2SR demo dataset from a repo-local image."
+    )
+    parser.add_argument(
+        "--output-dir",
+        required=True,
+        help="Directory where calibration, evaluation, and runtime demo assets are written.",
+    )
+    parser.add_argument(
+        "--source-image",
+        default=str(DEFAULT_SOURCE_IMAGE),
+        help="Source image used for the fixed text-heavy crops.",
+    )
+    parser.add_argument(
+        "--hr-size",
+        type=int,
+        default=128,
+        help="High-resolution crop size.",
+    )
+    parser.add_argument(
+        "--lr-size",
+        type=int,
+        default=64,
+        help="Low-resolution size written for runtime/calibration inputs.",
+    )
+    return parser.parse_args()
+
+
+def downsample_to_lr(image: Image.Image, lr_size: int) -> Image.Image:
+    return image.resize((lr_size, lr_size), Image.Resampling.BICUBIC)
+
+
+def save_image(path: Path, image: Image.Image) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    image.save(path)
+
+
+def crop_square(image: Image.Image, left: int, top: int, size: int) -> Image.Image:
+    right = left + size
+    bottom = top + size
+    if right > image.width or bottom > image.height:
+        raise ValueError(
+            f"Crop {(left, top, size)} exceeds source image bounds {image.size}."
+        )
+    return image.crop((left, top, right, bottom))
+
+
+def write_metadata(
+    output_dir: Path,
+    source_image: Path,
+    hr_size: int,
+    lr_size: int,
+) -> None:
+    metadata = {
+        "source_image": str(source_image),
+        "hr_size": hr_size,
+        "lr_size": lr_size,
+        "runtime_demo_crop": {
+            "name": RUNTIME_DEMO_CROP[0],
+            "left": RUNTIME_DEMO_CROP[1],
+            "top": RUNTIME_DEMO_CROP[2],
+        },
+        "calibration_crops": [
+            {"name": name, "left": left, "top": top}
+            for name, left, top in CALIBRATION_CROPS
+        ],
+        "eval_crops": [
+            {"name": name, "left": left, "top": top} for name, left, top in EVAL_CROPS
+        ],
+    }
+    metadata_path = output_dir / "metadata.json"
+    metadata_path.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n")
+
+
+def main() -> None:
+    args = parse_args()
+    if args.lr_size <= 0 or args.hr_size <= 0:
+        raise ValueError("--lr-size and --hr-size must be positive.")
+    if args.hr_size != args.lr_size * 2:
+        raise ValueError("This demo helper expects x2 super-resolution sizes.")
+
+    output_dir = Path(args.output_dir).resolve()
+    source_image = Path(args.source_image).resolve()
+    if not source_image.is_file():
+        raise FileNotFoundError(f"Source image not found: {source_image}")
+
+    image = Image.open(source_image).convert("RGB")
+
+    for name, left, top in CALIBRATION_CROPS:
+        hr_crop = crop_square(image, left, top, args.hr_size)
+        lr_crop = downsample_to_lr(hr_crop, args.lr_size)
+        save_image(output_dir / "calibration/hr" / f"{name}.png", hr_crop)
+        save_image(output_dir / "calibration/lr" / f"{name}.png", lr_crop)
+
+    for name, left, top in EVAL_CROPS:
+        hr_crop = crop_square(image, left, top, args.hr_size)
+        lr_crop = downsample_to_lr(hr_crop, args.lr_size)
+        save_image(output_dir / "eval/hr" / f"{name}.png", hr_crop)
+        save_image(output_dir / "eval/lr" / f"{name}.png", lr_crop)
+
+    demo_name, demo_left, demo_top = RUNTIME_DEMO_CROP
+    demo_hr = crop_square(image, demo_left, demo_top, args.hr_size)
+    demo_lr = downsample_to_lr(demo_hr, args.lr_size)
+    save_image(output_dir / "runtime" / f"{demo_name}_hr_{args.hr_size}.png", demo_hr)
+    save_image(output_dir / "runtime" / f"{demo_name}_lr_{args.lr_size}.png", demo_lr)
+
+    write_metadata(output_dir, source_image, args.hr_size, args.lr_size)
+
+    print(f"Prepared demo assets under {output_dir}")
+    print(
+        f"Runtime input: {output_dir / 'runtime' / f'{demo_name}_lr_{args.lr_size}.png'}"
+    )
+    print(
+        f"Runtime reference: {output_dir / 'runtime' / f'{demo_name}_hr_{args.hr_size}.png'}"
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/arm/super_resolution_example_vgf/requirements.txt b/examples/arm/super_resolution_example_vgf/requirements.txt
new file mode 100644
index 00000000000..ba8fbdf938a
--- /dev/null
+++ b/examples/arm/super_resolution_example_vgf/requirements.txt
@@ -0,0 +1,8 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+numpy == 2.1.3
+Pillow == 12.0.0
+transformers[torch] == 4.56.1
diff --git a/examples/arm/super_resolution_example_vgf/runtime/README.md b/examples/arm/super_resolution_example_vgf/runtime/README.md
new file mode 100644
index 00000000000..a5ab742fb81
--- /dev/null
+++ b/examples/arm/super_resolution_example_vgf/runtime/README.md
@@ -0,0 +1,57 @@
+# VGF Host Runtime (executor_runner)
+
+This flow runs the VGF-exported `.pte` on host using the portable
+`executor_runner` binary built at the repo root. The runtime helper script
+serializes the input image, invokes `executor_runner`, then reconstructs the
+output image from the generated tensor bytes.
+
+For the smallest reproducible demo in this repo, first create the fixed
+text-heavy input and the small LR/HR export set:
+
+```bash
+python examples/arm/super_resolution_example_vgf/model_export/prepare_demo_assets.py \
+  --output-dir ./demo_assets
+```
+
+1. Install ML SDK dependencies and set up the environment:
+
+```bash
+examples/arm/setup.sh --disable-ethos-u-deps --enable-mlsdk-deps
+source examples/arm/arm-scratch/setup_path.sh
+```
+
+2. Build the runner with VGF enabled:
+
+```bash
+cmake -B cmake-out \
+  -DCMAKE_BUILD_TYPE=Debug \
+  -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
+  -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
+  -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \
+  -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \
+  -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \
+  -DEXECUTORCH_BUILD_VULKAN=ON \
+  -DEXECUTORCH_BUILD_VGF=ON \
+  -DEXECUTORCH_ENABLE_LOGGING=ON \
+  .
+
+cmake --build cmake-out --target executor_runner
+```
+
+3. Run the exported model on an image that matches the static export size:
+
+```bash
+python examples/arm/super_resolution_example_vgf/runtime/run_super_resolution.py \
+  --model-path ./demo_assets/swin2sr_x2_vgf_fp.pte \
+  --input-image ./demo_assets/runtime/demo_lr_64.png \
+  --output-image ./demo_assets/runtime/demo_fp_128.png
+```
+
+The runtime helper reads the metadata emitted by the exporter to reconstruct the
+output tensor and save it as an image. Because the export uses static shapes,
+the input image must match the exported low-resolution dimensions exactly.
+
+Use the same runtime input with `./demo_assets/swin2sr_x2_vgf_int8.pte` to
+validate the INT8 path once the quantized export has been generated. If the
+host VKML emulation layer rejects quantized shaders, rerun that runtime step on
+Linux or on the target Vulkan driver stack.
diff --git a/examples/arm/super_resolution_example_vgf/runtime/run_super_resolution.py b/examples/arm/super_resolution_example_vgf/runtime/run_super_resolution.py
new file mode 100644
index 00000000000..ff6b8f30480
--- /dev/null
+++ b/examples/arm/super_resolution_example_vgf/runtime/run_super_resolution.py
@@ -0,0 +1,252 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import subprocess  # nosec B404 - executes the trusted local executor_runner binary
+import tempfile
+from pathlib import Path
+from typing import Any, Sequence
+
+import numpy as np
+import torch
+from PIL import Image
+
+STRING_TO_NUMPY_DTYPE = {
+    "float16": np.float16,
+    "float32": np.float32,
+    "int8": np.int8,
+    "uint8": np.uint8,
+}
+
+
+def read_json(path: str | Path) -> dict[str, Any]:
+    return json.loads(Path(path).read_text())
+
+
+def load_image_tensor(
+    image_path: str | Path,
+    expected_shape: Sequence[int] | None = None,
+) -> torch.Tensor:
+    image = Image.open(image_path).convert("RGB")
+    image_np = np.asarray(image, dtype=np.float32) / 255.0
+    tensor = (
+        torch.from_numpy(image_np).permute(2, 0, 1).contiguous().unsqueeze(0).clone()
+    )
+    if expected_shape is not None and tuple(tensor.shape) != tuple(expected_shape):
+        raise ValueError(
+            f"Image {image_path} produces tensor shape {tuple(tensor.shape)}, "
+            f"expected {tuple(expected_shape)}."
+        )
+    return tensor
+
+
+def save_tensor_bytes(path: str | Path, tensor: torch.Tensor) -> Path:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    array = tensor.detach().cpu().contiguous().numpy()
+    path.write_bytes(array.tobytes())
+    return path
+
+
+def load_tensor_bytes(
+    path: str | Path,
+    shape: Sequence[int],
+    dtype_name: str,
+) -> torch.Tensor:
+    np_dtype = STRING_TO_NUMPY_DTYPE.get(dtype_name)
+    if np_dtype is None:
+        raise ValueError(f"Unsupported tensor dtype in metadata: {dtype_name}")
+
+    array = np.fromfile(path, dtype=np_dtype)
+    expected_numel = math.prod(shape)
+    if array.size != expected_numel:
+        raise ValueError(
+            f"Tensor file {path} contains {array.size} values, expected {expected_numel}."
+        )
+
+    reshaped = array.reshape(tuple(shape)).copy()
+    return torch.from_numpy(reshaped)
+
+
+def save_image_tensor(path: str | Path, tensor: torch.Tensor) -> Path:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    image = tensor.detach().cpu()
+    if image.dim() == 4:
+        if image.shape[0] != 1:
+            raise ValueError(
+                "Only batch size 1 is supported when writing image output."
+            )
+        image = image[0]
+
+    if image.dim() != 3:
+        raise ValueError(f"Expected CHW image tensor, got shape {tuple(image.shape)}.")
+
+    channels = image.shape[0]
+    if channels not in {1, 3}:
+        raise ValueError(f"Expected 1 or 3 channels, got {channels}.")
+
+    image = image.clamp(0.0, 1.0)
+    image_np = (
+        image.permute(1, 2, 0).mul(255.0).round().to(torch.uint8).contiguous().numpy()
+    )
+
+    if channels == 1:
+        pil_image = Image.fromarray(image_np[..., 0], mode="L")
+    else:
+        pil_image = Image.fromarray(image_np, mode="RGB")
+    pil_image.save(path)
+    return path
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run a VGF-exported Swin2SR model with executor_runner."
+    )
+    parser.add_argument(
+        "--model-path",
+        required=True,
+        help="Path to the exported .pte file.",
+    )
+    parser.add_argument(
+        "--metadata-path",
+        default=None,
+        help="Optional metadata JSON path. Defaults to .json.",
+    )
+    parser.add_argument(
+        "--runner",
+        default="./cmake-out/executor_runner",
+        help="Path to the host executor_runner binary built with VGF support.",
+    )
+    parser.add_argument(
+        "--input-image",
+        required=True,
+        help="Low-resolution input image.",
+    )
+    parser.add_argument(
+        "--output-image",
+        required=True,
+        help="High-resolution output image path.",
+    )
+    parser.add_argument(
+        "--working-dir",
+        default=None,
+        help="Optional directory for temporary input/output tensor files.",
+    )
+    return parser.parse_args()
+
+
+def resolve_metadata_path(model_path: Path, metadata_path: str | None) -> Path:
+    if metadata_path is not None:
+        return Path(metadata_path).resolve()
+    return model_path.with_suffix(".json")
+
+
+def run_executor_runner(
+    runner: Path,
+    model_path: Path,
+    input_file: Path,
+    output_base: Path,
+) -> subprocess.CompletedProcess[str]:
+    command = [
+        str(runner),
+        "--model_path",
+        str(model_path),
+        "--inputs",
+        str(input_file),
+        "--output_file",
+        str(output_base),
+        "--print_output",
+        "none",
+    ]
+    return subprocess.run(  # nosec B603 - command list is assembled without a shell
+        command,
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+
+
+def main() -> None:
+    args = parse_args()
+    model_path = Path(args.model_path).resolve()
+    metadata_path = resolve_metadata_path(model_path, args.metadata_path)
+    runner_path = Path(args.runner).resolve()
+
+    if not model_path.is_file():
+        raise FileNotFoundError(f"Model file not found: {model_path}")
+    if not metadata_path.is_file():
+        raise FileNotFoundError(f"Metadata file not found: {metadata_path}")
+    if not runner_path.is_file():
+        raise FileNotFoundError(f"executor_runner not found: {runner_path}")
+
+    metadata = read_json(metadata_path)
+    if metadata.get("num_outputs") != 1:
+        raise ValueError(
+            "The runtime helper currently supports single-output models only."
+        )
+
+    if args.working_dir is None:
+        with tempfile.TemporaryDirectory(prefix="executorch-sr-vgf-") as tmp_dir:
+            workdir = Path(tmp_dir)
+            run_once(
+                model_path=model_path,
+                metadata=metadata,
+                runner_path=runner_path,
+                input_image=Path(args.input_image),
+                output_image=Path(args.output_image),
+                working_dir=workdir,
+            )
+    else:
+        workdir = Path(args.working_dir).resolve()
+        workdir.mkdir(parents=True, exist_ok=True)
+        run_once(
+            model_path=model_path,
+            metadata=metadata,
+            runner_path=runner_path,
+            input_image=Path(args.input_image),
+            output_image=Path(args.output_image),
+            working_dir=workdir,
+        )
+
+
+def run_once(
+    model_path: Path,
+    metadata: dict,
+    runner_path: Path,
+    input_image: Path,
+    output_image: Path,
+    working_dir: Path,
+) -> None:
+    input_tensor = load_image_tensor(input_image, metadata["input_shape"])
+    input_path = save_tensor_bytes(working_dir / "input0.bin", input_tensor)
+    output_base = working_dir / "output"
+
+    result = run_executor_runner(runner_path, model_path, input_path, output_base)
+    if result.stdout:
+        print(result.stdout.strip())
+    if result.returncode != 0:
+        raise RuntimeError(
+            "executor_runner failed.\n"
+            f"stdout:\n{result.stdout}\n"
+            f"stderr:\n{result.stderr}"
+        )
+
+    output_tensor = load_tensor_bytes(
+        output_base.with_name(f"{output_base.name}-0.bin"),
+        metadata["output_shape"],
+        metadata["output_dtype"],
+    )
+    save_image_tensor(output_image, output_tensor)
+    print(f"Saved super-resolved image to {output_image.resolve()}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/arm/visualize.py b/examples/arm/visualize.py
index fc410f60de1..f7a3b28d9e6 100644
--- a/examples/arm/visualize.py
+++ b/examples/arm/visualize.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -17,8 +17,11 @@
 from executorch.devtools.visualization.visualization_utils import (
     visualize_model_explorer,
 )
-from model_explorer import config as model_explorer_config, node_data_builder as ndb
-from model_explorer.config import ModelSource
+from model_explorer import (  # type: ignore[import]
+    config as model_explorer_config,
+    node_data_builder as ndb,
+)
+from model_explorer.config import ModelSource  # type: ignore[import]
 
 COMPILER_OP_ID = "scheduled_id"
 
@@ -254,7 +257,7 @@ def validate_perf_mode_args(trace: str, tables: str) -> None:
 
 
 def set_pte_model_explorer_config(model_file, tosa_files, config):
-    from pte_adapter_model_explorer.main import PTEAdapter
+    from pte_adapter_model_explorer.main import PTEAdapter  # type: ignore[import]
 
     pte_adapter = PTEAdapter()
 
diff --git a/examples/models/gemma4_31b/CMakeLists.txt b/examples/models/gemma4_31b/CMakeLists.txt
index 8d536a47fc5..52419eb95bc 100644
--- a/examples/models/gemma4_31b/CMakeLists.txt
+++ b/examples/models/gemma4_31b/CMakeLists.txt
@@ -42,14 +42,17 @@ list(
   extension_flat_tensor
 )
 
-# CUDA backend (the only supported backend for this example for now)
+# Backend: CUDA or MLX (exactly one required)
 if(EXECUTORCH_BUILD_CUDA)
   find_package(CUDAToolkit REQUIRED)
   list(APPEND link_libraries aoti_cuda_backend)
   executorch_target_link_options_shared_lib(aoti_cuda_backend)
   add_compile_definitions(EXECUTORCH_BUILD_CUDA)
+elseif(TARGET mlxdelegate)
+  list(APPEND link_libraries mlxdelegate mlx)
+  executorch_target_link_options_shared_lib(mlxdelegate)
 else()
-  message(FATAL_ERROR "Set EXECUTORCH_BUILD_CUDA=ON")
+  message(FATAL_ERROR "Set EXECUTORCH_BUILD_CUDA=ON or EXECUTORCH_BUILD_MLX=ON")
 endif()
 
 # Tokenizer (HuggingFace tokenizer.json)
@@ -63,5 +66,11 @@ target_link_libraries(gemma4_31b_runner PUBLIC ${link_libraries})
 
 if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug")
   target_link_options_gc_sections(gemma4_31b_runner)
-  target_link_options(gemma4_31b_runner PRIVATE "LINKER:-s")
+  if(NOT APPLE AND NOT MSVC)
+    target_link_options(gemma4_31b_runner PRIVATE "LINKER:-s")
+  endif()
+endif()
+
+if(TARGET mlxdelegate)
+  executorch_target_copy_mlx_metallib(gemma4_31b_runner)
 endif()
diff --git a/examples/models/gemma4_31b/CMakePresets.json b/examples/models/gemma4_31b/CMakePresets.json
index 97ba7f4c57a..23a7d42e035 100644
--- a/examples/models/gemma4_31b/CMakePresets.json
+++ b/examples/models/gemma4_31b/CMakePresets.json
@@ -23,6 +23,17 @@
                 "string": "${hostSystemName}",
                 "list": ["Linux", "Windows"]
             }
+        },
+        {
+            "name": "gemma4-31b-mlx",
+            "displayName": "Gemma 4 31B runner (MLX)",
+            "inherits": ["gemma4-31b-base"],
+            "cacheVariables": {},
+            "condition": {
+                "type": "equals",
+                "lhs": "${hostSystemName}",
+                "rhs": "Darwin"
+            }
         }
     ],
     "buildPresets": [
@@ -31,6 +42,12 @@
             "displayName": "Build Gemma 4 31B runner (CUDA)",
             "configurePreset": "gemma4-31b-cuda",
             "targets": ["gemma4_31b_runner"]
+        },
+        {
+            "name": "gemma4-31b-mlx",
+            "displayName": "Build Gemma 4 31B runner (MLX)",
+            "configurePreset": "gemma4-31b-mlx",
+            "targets": ["gemma4_31b_runner"]
         }
     ],
     "workflowPresets": [
@@ -47,6 +64,20 @@
                     "name": "gemma4-31b-cuda"
                 }
             ]
+        },
+        {
+            "name": "gemma4-31b-mlx",
+            "displayName": "Configure and build Gemma 4 31B runner (MLX)",
+            "steps": [
+                {
+                    "type": "configure",
+                    "name": "gemma4-31b-mlx"
+                },
+                {
+                    "type": "build",
+                    "name": "gemma4-31b-mlx"
+                }
+            ]
         }
     ]
 }
diff --git a/examples/models/gemma4_31b/README.md b/examples/models/gemma4_31b/README.md
index 6f567d739b7..da4aa893079 100644
--- a/examples/models/gemma4_31b/README.md
+++ b/examples/models/gemma4_31b/README.md
@@ -1,7 +1,7 @@
 # Gemma 4 31B-IT
 
 Text-only export of Google's Gemma 4 31B-IT to ExecuTorch with INT4/INT8
-weight quantization. Currently supports the CUDA backend.
+weight quantization. Supports CUDA and MLX (Apple Silicon) backends.
 
 For architecture and design notes see [model.md](model.md).
 
@@ -67,6 +67,8 @@ recipe. Writes `model.safetensors`, `config.json`, and `tokenizer.json` into
 
 ## Export to ExecuTorch
 
+### CUDA
+
 ```bash
 python examples/models/gemma4_31b/export.py \
     --prequantized ./gemma4_31b_int4 \
@@ -75,10 +77,26 @@ python examples/models/gemma4_31b/export.py \
     --backend cuda
 ```
 
-Writes `model.pte` and `model.ptd` into `--output-dir`.
+### MLX (Apple Silicon)
+
+```bash
+python examples/models/gemma4_31b/export.py \
+    --prequantized ./gemma4_31b_int4 \
+    --output-dir ./gemma4_31b_exports_mlx \
+    --max-seq-len 4096 \
+    --backend mlx
+```
+
+The same quantized checkpoint works for both backends. MLX exports a single
+method with dynamic sequence length and host-side sampling.
+
+Writes `model.pte` (and optionally `model.ptd`) into `--output-dir`.
 
 ## Eager inference
 
+The prompt is automatically wrapped with the Gemma 4 IT chat template.
+Pass `--raw-prompt` to skip template wrapping for pre-formatted input.
+
 ```bash
 python examples/models/gemma4_31b/inference.py \
     --prequantized ./gemma4_31b_int4 \
@@ -102,13 +120,17 @@ model produces sensible text.
 ## Build the runner
 
 ```bash
-make gemma4_31b-cuda
+make gemma4_31b-cuda   # Linux — CUDA backend
+make gemma4_31b-mlx    # macOS — MLX backend (Apple Silicon)
 ```
 
 The binary lands at `cmake-out/examples/models/gemma4_31b/gemma4_31b_runner`.
 
 ## Run the .pte
 
+The prompt is automatically wrapped with the Gemma 4 IT chat template.
+Pass `--raw_prompt` to skip template wrapping for pre-formatted input.
+
 ```bash
 ./gemma4_31b_runner \
     --model_path  ./gemma4_31b_exports/model.pte \
diff --git a/examples/models/gemma4_31b/export.py b/examples/models/gemma4_31b/export.py
index a96dba0d512..046e365947b 100644
--- a/examples/models/gemma4_31b/export.py
+++ b/examples/models/gemma4_31b/export.py
@@ -19,6 +19,8 @@
 
 Backends:
   --backend cuda            (default) CUDA via tinygemm INT4 + CudaPartitioner.
+  --backend mlx             Apple Silicon via MLXPartitioner (single method,
+                            dynamic seq_len, host-side sampling).
 """
 
 import argparse
@@ -98,12 +100,21 @@ def load_and_quantize(
 # Backend dispatch helpers
 
 
+_SUPPORTED_BACKENDS = ("cuda", "mlx")
+
+
 def _get_packers(backend: str) -> dict:
     if backend == "cuda":
         from executorch.examples.models.gemma4_31b.quant import DEFAULT_CUDA_PACKERS
 
         return DEFAULT_CUDA_PACKERS
-    raise ValueError(f"Unsupported backend: {backend!r}. Supported: 'cuda'.")
+    if backend == "mlx":
+        from executorch.examples.models.gemma4_31b.quant import DEFAULT_MLX_PACKERS
+
+        return DEFAULT_MLX_PACKERS
+    raise ValueError(
+        f"Unsupported backend: {backend!r}. Supported: {_SUPPORTED_BACKENDS}."
+    )
 
 
 def _pack_for_backend(model: nn.Module, path: str, backend: str) -> None:
@@ -111,8 +122,14 @@ def _pack_for_backend(model: nn.Module, path: str, backend: str) -> None:
         from executorch.examples.models.gemma4_31b.quant import load_and_pack_for_cuda
 
         load_and_pack_for_cuda(path, model)
+    elif backend == "mlx":
+        from executorch.examples.models.gemma4_31b.quant import load_and_pack_for_mlx
+
+        load_and_pack_for_mlx(path, model)
     else:
-        raise ValueError(f"Unsupported backend: {backend!r}. Supported: 'cuda'.")
+        raise ValueError(
+            f"Unsupported backend: {backend!r}. Supported: {_SUPPORTED_BACKENDS}."
+        )
 
 
 # ---------------------------------------------------------------------------
@@ -128,8 +145,12 @@ def export_and_lower(
     """Export and lower the model to ExecuTorch for the given backend."""
     if backend == "cuda":
         _export_cuda(model, config, output_dir)
+    elif backend == "mlx":
+        _export_mlx(model, config, output_dir)
     else:
-        raise ValueError(f"Unsupported backend: {backend!r}. Supported: 'cuda'.")
+        raise ValueError(
+            f"Unsupported backend: {backend!r}. Supported: {_SUPPORTED_BACKENDS}."
+        )
 
 
 def _export_cuda(model: Gemma4_31B, config: Gemma4_31BConfig, output_dir: str) -> None:
@@ -258,6 +279,98 @@ def _export_cuda(model: Gemma4_31B, config: Gemma4_31BConfig, output_dir: str) -
     print("Done.")
 
 
+def _export_mlx(model: Gemma4_31B, config: Gemma4_31BConfig, output_dir: str) -> None:
+    """Export to .pte via torch.export + MLX backend.
+
+    Unlike CUDA (which exports separate decode/prefill methods with an
+    Int4Tensor dispatch override), MLX uses a single method with dynamic
+    sequence length.  No int4_dispatch import — IntxUnpackedToInt8Tensor's
+    default dispatch produces the ``dequantize_affine → linear`` pattern
+    that MLX's QuantizedLinearHandler matches.
+    """
+    import gc
+
+    from executorch.backends.mlx import MLXPartitioner
+    from executorch.backends.mlx.passes import get_default_passes
+
+    from executorch.examples.models.gemma4_31b.mlx_source_transformations import (
+        mlx_source_transformations,
+    )
+    from executorch.exir import (
+        EdgeCompileConfig,
+        ExecutorchBackendConfig,
+        to_edge_transform_and_lower,
+    )
+    from executorch.exir.passes import MemoryPlanningPass
+    from torch.export import Dim, export
+
+    mlx_source_transformations(model, dtype=torch.bfloat16)
+    materialize_runtime_buffers(model, dtype=torch.bfloat16)
+
+    max_prefill = min(config.max_seq_len - 1, config.sliding_window * 2)
+    seq_dim = Dim("seq_len", min=1, max=max_prefill)
+
+    print(f"Exporting (T in [1, {max_prefill}])...")
+    with torch.no_grad():
+        exported = export(
+            model,
+            (
+                torch.tensor([[0, 1]], dtype=torch.long),
+                torch.tensor([0, 1], dtype=torch.long),
+            ),
+            dynamic_shapes=({1: seq_dim}, {0: seq_dim}),
+            strict=True,
+        )
+
+    del model
+    gc.collect()
+
+    print("Lowering to ExecuTorch with MLX backend...")
+    et_prog = to_edge_transform_and_lower(
+        exported,
+        transform_passes=get_default_passes(),
+        partitioner=[MLXPartitioner()],
+        compile_config=EdgeCompileConfig(
+            _check_ir_validity=False,
+            _skip_dim_order=True,
+        ),
+        constant_methods={
+            "get_max_seq_len": config.max_seq_len,
+            "get_vocab_size": config.vocab_size,
+            "get_n_layers": config.num_hidden_layers,
+            "get_max_prefill_chunk": max_prefill,
+            "use_kv_cache": True,
+            "use_sdpa_with_kv_cache": False,
+            "enable_dynamic_shape": True,
+        },
+    )
+
+    del exported
+    gc.collect()
+
+    et_program = et_prog.to_executorch(
+        config=ExecutorchBackendConfig(
+            extract_delegate_segments=True,
+            memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False),
+        ),
+    )
+
+    del et_prog
+    gc.collect()
+
+    os.makedirs(output_dir, exist_ok=True)
+    pte_path = os.path.join(output_dir, "model.pte")
+    print(f"Saving to {pte_path}...")
+    with open(pte_path, "wb") as f:
+        et_program.write_to_file(f)
+    print(f"  {os.path.getsize(pte_path) / 1024**2:.1f} MB")
+
+    if et_program._tensor_data:
+        et_program.write_tensor_data_to_file(output_dir)
+        print(f"  Saved tensor data (.ptd) to {output_dir}/")
+    print("Done.")
+
+
 # ---------------------------------------------------------------------------
 # CLI
 
@@ -302,7 +415,7 @@ def main() -> None:
     parser.add_argument(
         "--backend",
         default="cuda",
-        choices=["cuda"],
+        choices=list(_SUPPORTED_BACKENDS),
         help="Target backend for export.",
     )
     args = parser.parse_args()
diff --git a/examples/models/gemma4_31b/inference.py b/examples/models/gemma4_31b/inference.py
index 12785450d8c..92654fca5f2 100644
--- a/examples/models/gemma4_31b/inference.py
+++ b/examples/models/gemma4_31b/inference.py
@@ -6,12 +6,15 @@
 
 """Eager inference on Gemma 4 31B-IT (CUDA + torch.compile).
 
-Two input paths:
+Three input paths:
   --prequantized    Load a quantized checkpoint (from quantize_and_save.py).
   --gguf           Load a GGUF file (e.g., Q4_K_M from the community).
+  --bf16            Load the bf16 HF safetensors checkpoint via from_hf_checkpoint.
 
-Packs for the target backend (--backend cuda), materializes runtime buffers,
-optionally compiles with ``torch.compile``, and generates text autoregressively.
+Gemma 4 31B-IT is instruction-tuned and requires chat-template formatting.
+The ``--prompt`` is automatically wrapped with the Gemma 4 chat template
+(``<|turn>user\\n{prompt}\\n<|turn>model\\n<|channel>thought\\n``; BOS is prepended separately).
+Pass ``--raw-prompt`` to skip template wrapping (e.g., for pre-formatted input).
 
 Usage:
     python inference.py \\
@@ -33,7 +36,10 @@
 import torch
 
 from executorch.examples.models.gemma4_31b.export import load_prequantized_model
-from executorch.examples.models.gemma4_31b.model import materialize_runtime_buffers
+from executorch.examples.models.gemma4_31b.model import (
+    Gemma4_31B,
+    materialize_runtime_buffers,
+)
 
 
 def _move_to_cuda(model, config) -> None:
@@ -63,12 +69,24 @@ def _move_to_cuda(model, config) -> None:
     materialize_runtime_buffers(model, dtype=torch.bfloat16, device="cuda")
 
 
+def apply_chat_template(prompt: str) -> str:
+    """Wrap a user prompt in the Gemma 4 IT chat template.
+
+    Does not include BOS — ``generate()`` prepends it at the token-ID level.
+    """
+    return (
+        "<|turn>user\n"
+        + prompt
+        + "\n<|turn>model\n<|channel>thought\n"
+    )
+
+
 def generate(
     model,
     tokenizer,
     prompt: str,
     max_new_tokens: int = 128,
-    temperature: float = 0.0,
+    temperature: float = 0.8,
     eos_token_ids=None,
     bos_token_id: int = 2,
 ) -> str:
@@ -131,6 +149,11 @@ def main() -> None:
         default=None,
         help="Path to a GGUF file (e.g., gemma-4-31B-it-Q4_K_M.gguf).",
     )
+    src.add_argument(
+        "--bf16",
+        default=None,
+        help="Path to a bf16 hf directory (e.g., gemma-4-31B).",
+    )
     parser.add_argument(
         "--tokenizer-path",
         default=None,
@@ -155,6 +178,11 @@ def main() -> None:
         default=4096,
         help="KV cache length to allocate for this run.",
     )
+    parser.add_argument(
+        "--raw-prompt",
+        action="store_true",
+        help="Skip chat-template wrapping (use if the prompt is already formatted).",
+    )
     parser.add_argument(
         "--no-compile",
         action="store_true",
@@ -171,12 +199,34 @@ def main() -> None:
     if args.backend == "cuda" and not torch.cuda.is_available():
         parser.error("CUDA is required for the cuda backend.")
 
+    # ---- Tokenizer ----
+    if args.tokenizer_path:
+        tokenizer_path = args.tokenizer_path
+    elif args.prequantized:
+        tokenizer_path = os.path.join(args.prequantized, "tokenizer.json")
+    elif args.bf16:
+        tokenizer_path = os.path.join(args.bf16, "tokenizer.json")
+    else:
+        parser.error("--tokenizer-path is required with --gguf.")
+    from tokenizers import Tokenizer
+
+    tokenizer = Tokenizer.from_file(tokenizer_path)
+
+    prompt_str = args.prompt if args.raw_prompt else apply_chat_template(args.prompt)
+
+    # Gemma 4 EOS tokens (from generation_config.json: ids 1, 50, 106).
+    eos_token_ids = {1, 50, 106}
+
     if args.gguf:
         from executorch.examples.models.gemma4_31b.gguf_loader import load_gguf_model
 
         model, config = load_gguf_model(
             args.gguf, args.max_seq_len, backend=args.backend
         )
+    elif args.bf16:
+        model, config = Gemma4_31B.from_hf_checkpoint(
+            args.bf16, max_seq_len=args.max_seq_len
+        )
     else:
         print(f"Loading prequantized model from {args.prequantized}...")
         model, config = load_prequantized_model(
@@ -191,19 +241,6 @@ def main() -> None:
         print("Compiling model with torch.compile...")
         model = torch.compile(model, mode="default")
 
-    if args.tokenizer_path:
-        tokenizer_path = args.tokenizer_path
-    elif args.prequantized:
-        tokenizer_path = os.path.join(args.prequantized, "tokenizer.json")
-    else:
-        parser.error("--tokenizer-path is required with --gguf.")
-    from tokenizers import Tokenizer
-
-    tokenizer = Tokenizer.from_file(tokenizer_path)
-
-    # Gemma 4 EOS tokens (from generation_config.json: ids 1, 50, 106).
-    eos_token_ids = {1, 50, 106}
-
     print(f"\nPrompt: {args.prompt}")
     print("-" * 40)
 
@@ -211,7 +248,7 @@ def main() -> None:
     output = generate(
         model,
         tokenizer,
-        args.prompt,
+        prompt_str,
         max_new_tokens=args.max_new_tokens,
         temperature=args.temperature,
         eos_token_ids=eos_token_ids,
diff --git a/examples/models/gemma4_31b/main.cpp b/examples/models/gemma4_31b/main.cpp
index 0be2fef517c..6cf65cc8246 100644
--- a/examples/models/gemma4_31b/main.cpp
+++ b/examples/models/gemma4_31b/main.cpp
@@ -6,18 +6,21 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// Gemma 4 31B-IT runner for the CUDA ExecuTorch backend.
-//
-// Drives the prefill + decode methods produced by export.py.
-// The exported model performs Gumbel-max sampling on-device and returns a
-// single float token ID per call, so this runner only has to feed tokens
-// in and decode them via the HuggingFace tokenizer.
+// Gemma 4 31B-IT runner for ExecuTorch. Supports two backends:
+//   CUDA  — exports ``prefill`` (T>=2, dynamic) + ``decode`` (T=1, static)
+//           methods sharing KV-cache buffers; on-device Gumbel-max sampling
+//           with temperature passed as a third input; returns a scalar
+//           float token id.
+//   MLX   — exports a single ``forward`` method with dynamic seq_len;
+//           returns last-token logits; the runner samples on the host via
+//           ``llm::logits_to_token`` with the same temperature semantics.
 
 #include 
 
 #include 
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -65,6 +68,10 @@ DEFINE_double(temperature, 0.8, "Sampling temperature (0 = near-greedy).");
 DEFINE_int32(max_new_tokens, 128, "Maximum tokens to generate.");
 DEFINE_int32(bos_id, 2, "BOS token id to prepend (Gemma convention: 2).");
 DEFINE_int32(eos_id, 1, "EOS token id (Gemma convention: 1).");
+DEFINE_bool(
+    raw_prompt,
+    false,
+    "Skip chat-template wrapping (use if the prompt is already formatted).");
 DEFINE_bool(
     cuda_graph,
     false,
@@ -78,6 +85,7 @@ using ::executorch::runtime::EValue;
 
 using SizesType = executorch::aten::SizesType;
 
+// Read a sampled token ID from a scalar float output (CUDA path).
 static uint64_t read_token(const executorch::aten::Tensor& output) {
   const void* ptr = output.const_data_ptr();
   float val = 0.0f;
@@ -139,8 +147,7 @@ int main(int argc, char** argv) {
     return 1;
   }
 
-  // Module: share_memory_arenas=true so prefill and decode see the same
-  // KV-cache memory (we exported with share_mutable_buffers=True).
+  // Module
   std::vector data_files;
   if (!FLAGS_data_path.empty()) {
     data_files.push_back(FLAGS_data_path);
@@ -148,7 +155,7 @@ int main(int argc, char** argv) {
   auto module = std::make_unique(
       FLAGS_model_path,
       data_files,
-      Module::LoadMode::File,
+      Module::LoadMode::MmapUseMlockIgnoreErrors,
       /*event_tracer=*/nullptr,
       /*memory_allocator=*/nullptr,
       /*temp_allocator=*/nullptr,
@@ -161,6 +168,19 @@ int main(int argc, char** argv) {
     return 1;
   }
 
+  int64_t max_prefill_chunk = (*metadata_result)[llm::kMaxSeqLen] - 1;
+  {
+    auto get_result = module->get("get_max_prefill_chunk");
+    if (get_result.ok()) {
+      max_prefill_chunk = get_result->toScalar().to();
+    }
+  }
+
+  auto S = [](int64_t v) -> SizesType { return static_cast(v); };
+
+  float temp_val =
+      FLAGS_temperature <= 0.0 ? 1e-6f : static_cast(FLAGS_temperature);
+
 #ifdef EXECUTORCH_BUILD_CUDA
   if (FLAGS_cuda_graph) {
     executorch::runtime::BackendOptions<2> cuda_opts;
@@ -168,11 +188,6 @@ int main(int argc, char** argv) {
     executorch::runtime::set_option("CudaBackend", cuda_opts.view());
     printf("CUDA graph enabled for decode method\n");
   }
-
-  // Cross-method per-FQN weight sharing: prefill + decode share the same
-  // weight tensors and (more importantly) the same KV-cache buffers, so
-  // without this flag we would allocate them twice. MUST be set before
-  // load_method.
   {
     executorch::runtime::BackendOptions<1> backend_options;
     auto set_err =
@@ -180,7 +195,7 @@ int main(int argc, char** argv) {
     if (set_err != Error::Ok) {
       ET_LOG(
           Error,
-          "Failed to construct weight_sharing_across_methods option: %d",
+          "Failed to set weight_sharing_across_methods: %d",
           static_cast(set_err));
       return 1;
     }
@@ -194,12 +209,6 @@ int main(int argc, char** argv) {
       return 1;
     }
   }
-#else
-  if (FLAGS_cuda_graph) {
-    ET_LOG(Info, "--cuda_graph ignored on non-CUDA build");
-  }
-#endif
-
   printf("Loading methods...\n");
   if (module->load_method("prefill") != Error::Ok) {
     ET_LOG(Error, "Failed to load prefill method");
@@ -209,6 +218,19 @@ int main(int argc, char** argv) {
     ET_LOG(Error, "Failed to load decode method");
     return 1;
   }
+  auto temp_tensor =
+      from_blob(&temp_val, {1}, executorch::aten::ScalarType::Float);
+#else
+  if (FLAGS_cuda_graph) {
+    ET_LOG(Info, "--cuda_graph ignored on non-CUDA build");
+  }
+  printf("Loading model...\n");
+  if (module->load_method("forward") != Error::Ok) {
+    ET_LOG(Error, "Failed to load forward method");
+    return 1;
+  }
+#endif
+
   stats.model_load_end_ms = llm::time_in_ms();
 
 #ifdef EXECUTORCH_BUILD_CUDA
@@ -218,8 +240,12 @@ int main(int argc, char** argv) {
 
   auto eos_ids = llm::get_eos_ids(tokenizer.get(), module.get());
   eos_ids.insert(static_cast(FLAGS_eos_id));
+  auto turn_ids = tokenizer->encode("", /*bos=*/0, /*eos=*/0);
+  if (turn_ids.ok() && turn_ids->size() == 1) {
+    eos_ids.insert(turn_ids.get()[0]);
+  }
 
-  // Read prompt from file or flag
+  // Read prompt
   std::string prompt_text = FLAGS_prompt;
   if (!FLAGS_prompt_file.empty()) {
     std::ifstream f(FLAGS_prompt_file);
@@ -232,6 +258,14 @@ int main(int argc, char** argv) {
         (std::istreambuf_iterator(f)), std::istreambuf_iterator());
   }
 
+  // Wrap with Gemma 4 IT chat template unless --raw_prompt is set.
+  // BOS is prepended separately below; this adds the turn structure and the
+  // empty thought block required by the instruction-tuned model.
+  if (!FLAGS_raw_prompt) {
+    prompt_text = "<|turn>user\n" + prompt_text +
+        "\n<|turn>model\n<|channel>thought\n";
+  }
+
   // Encode prompt
   auto encode_result = tokenizer->encode(prompt_text);
   if (!encode_result.ok()) {
@@ -248,38 +282,15 @@ int main(int argc, char** argv) {
 
   stats.inference_start_ms = llm::time_in_ms();
 
-  auto S = [](int64_t v) -> SizesType { return static_cast(v); };
-
-#ifdef EXECUTORCH_BUILD_CUDA
-  // CUDA build: model fuses the sampler. Pass temperature as a third input.
-  float temp_val =
-      FLAGS_temperature <= 0.0 ? 1e-6f : static_cast(FLAGS_temperature);
-  auto temp_tensor =
-      from_blob(&temp_val, {1}, executorch::aten::ScalarType::Float);
-#endif
-
   // ---------------------------------------------------------------
   // Prefill (chunked to respect ring-buffer KV cache limit)
   // ---------------------------------------------------------------
-  // Sliding layers use a ring buffer sized to 2×sliding_window. A single
-  // prefill call must not exceed this size, otherwise index_copy_ with
-  // wrapped indices produces non-deterministic results on CUDA.
-  int64_t max_prefill_chunk = (*metadata_result)[llm::kMaxSeqLen] - 1;
-  {
-    auto get_result = module->get("get_max_prefill_chunk");
-    if (get_result.ok()) {
-      max_prefill_chunk = get_result->toScalar().to();
-    }
-  }
-
   uint64_t cur_token = 0;
   int64_t prefill_pos = 0;
   while (prefill_pos < num_prompt_tokens) {
     int64_t chunk_len =
         std::min(num_prompt_tokens - prefill_pos, max_prefill_chunk);
 
-    std::string run_method = (chunk_len == 1) ? "decode" : "prefill";
-
     std::vector token_data(
         prompt_tokens.begin() + prefill_pos,
         prompt_tokens.begin() + prefill_pos + chunk_len);
@@ -294,39 +305,52 @@ int main(int argc, char** argv) {
     auto pos_tensor = from_blob(
         pos_data.data(), {S(chunk_len)}, executorch::aten::ScalarType::Long);
 
-    std::vector prefill_inputs;
-    prefill_inputs.push_back(EValue(tokens_tensor));
-    prefill_inputs.push_back(EValue(pos_tensor));
+    std::vector inputs;
+    inputs.push_back(EValue(tokens_tensor));
+    inputs.push_back(EValue(pos_tensor));
+
 #ifdef EXECUTORCH_BUILD_CUDA
-    prefill_inputs.push_back(EValue(temp_tensor));
+    inputs.push_back(EValue(temp_tensor));
+    std::string method = (chunk_len == 1) ? "decode" : "prefill";
+#else
+    std::string method = "forward";
 #endif
 
-    auto prefill_result = module->execute(run_method, prefill_inputs);
-    if (prefill_result.error() != Error::Ok) {
-      ET_LOG(
-          Error, "%s failed at pos %" PRId64, run_method.c_str(), prefill_pos);
+    auto result = module->execute(method, inputs);
+    if (result.error() != Error::Ok) {
+      ET_LOG(Error, "%s failed at pos %" PRId64, method.c_str(), prefill_pos);
       return 1;
     }
-    cur_token = read_token(prefill_result.get()[0].toTensor());
+
+#ifdef EXECUTORCH_BUILD_CUDA
+    cur_token = read_token(result.get()[0].toTensor());
+#else
+    cur_token = static_cast(
+        llm::logits_to_token(result.get()[0].toTensor(), temp_val));
+#endif
+
     prefill_pos += chunk_len;
   }
 
   stats.prompt_eval_end_ms = llm::time_in_ms();
-  double prefill_ms =
-      static_cast(stats.prompt_eval_end_ms - stats.inference_start_ms);
-  printf(
-      "Prefill: %" PRId64 " tokens in %.1f ms (%.1f tok/s)\n",
-      num_prompt_tokens,
-      prefill_ms,
-      num_prompt_tokens * 1000.0 / prefill_ms);
+  // First generated token came from the last prefill chunk; TTFT is prefill.
+  stats.first_token_ms = stats.prompt_eval_end_ms;
 
 #ifdef EXECUTORCH_BUILD_CUDA
-  // Synchronize CUDA device to ensure prefill's writes to shared mutable
-  // buffers (KV cache) are visible to the decode method, which may run on
-  // a different CUDA stream.
   cudaDeviceSynchronize();
 #endif
 
+  // Print the first generated token (from the last prefill chunk).
+  // Use the last prompt token as the streaming-decode prefix so any BPE
+  // partial-character handling stays correct.
+  {
+    auto first_str = tokenizer->decode(prompt_tokens.back(), cur_token);
+    if (first_str.ok()) {
+      printf("%s", first_str->c_str());
+      fflush(stdout);
+    }
+  }
+
   // ---------------------------------------------------------------
   // Decode loop
   // ---------------------------------------------------------------
@@ -339,29 +363,34 @@ int main(int argc, char** argv) {
       decode_pos_data.data(), {1}, executorch::aten::ScalarType::Long);
 
   uint64_t prev_token = cur_token;
-  for (int32_t step = 0; step < FLAGS_max_new_tokens; step++) {
+  bool hit_eos = eos_ids.find(cur_token) != eos_ids.end();
+  for (int32_t step = 0; step < FLAGS_max_new_tokens && !hit_eos; step++) {
     decode_token_data[0] = static_cast(cur_token);
     decode_pos_data[0] = pos;
 
-    std::vector decode_inputs;
-    decode_inputs.push_back(EValue(decode_tokens));
-    decode_inputs.push_back(EValue(decode_pos));
+    std::vector inputs;
+    inputs.push_back(EValue(decode_tokens));
+    inputs.push_back(EValue(decode_pos));
+
 #ifdef EXECUTORCH_BUILD_CUDA
-    decode_inputs.push_back(EValue(temp_tensor));
+    inputs.push_back(EValue(temp_tensor));
+    auto result = module->execute("decode", inputs);
+#else
+    auto result = module->execute("forward", inputs);
 #endif
 
-    auto decode_result = module->execute("decode", decode_inputs);
-    if (decode_result.error() != Error::Ok) {
+    if (result.error() != Error::Ok) {
       ET_LOG(Error, "Decode step %d failed", step);
       return 1;
     }
 
     prev_token = cur_token;
-    cur_token = read_token(decode_result.get()[0].toTensor());
-
-    if (step == 0) {
-      stats.first_token_ms = llm::time_in_ms();
-    }
+#ifdef EXECUTORCH_BUILD_CUDA
+    cur_token = read_token(result.get()[0].toTensor());
+#else
+    cur_token = static_cast(
+        llm::logits_to_token(result.get()[0].toTensor(), temp_val));
+#endif
     pos++;
 
     auto decode_str = tokenizer->decode(prev_token, cur_token);
@@ -370,25 +399,12 @@ int main(int argc, char** argv) {
       fflush(stdout);
     }
 
-    if (eos_ids.find(cur_token) != eos_ids.end()) {
-      printf("\n");
-      break;
-    }
+    hit_eos = eos_ids.find(cur_token) != eos_ids.end();
   }
-
-  stats.inference_end_ms = llm::time_in_ms();
   printf("\n");
 
-  int64_t num_generated = pos - num_prompt_tokens;
-  stats.num_generated_tokens = num_generated;
-  double decode_ms =
-      static_cast(stats.inference_end_ms - stats.prompt_eval_end_ms);
-  printf(
-      "Decode: %" PRId64 " tokens in %.1f ms (%.1f tok/s)\n",
-      num_generated,
-      decode_ms,
-      num_generated * 1000.0 / decode_ms);
-  printf("Prompt tokens: %" PRId64 "\n", num_prompt_tokens);
+  stats.inference_end_ms = llm::time_in_ms();
+  stats.num_generated_tokens = pos - num_prompt_tokens;
 
 #ifdef EXECUTORCH_BUILD_CUDA
   cudaMemGetInfo(&gpu_free_bytes, &gpu_total_bytes);
diff --git a/examples/models/gemma4_31b/mlx_source_transformations.py b/examples/models/gemma4_31b/mlx_source_transformations.py
new file mode 100644
index 00000000000..3a8ae4420e3
--- /dev/null
+++ b/examples/models/gemma4_31b/mlx_source_transformations.py
@@ -0,0 +1,192 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""MLX source transformations for Gemma 4 31B-IT.
+
+Replaces the generic PyTorch ops in the model with MLX custom ops that lower
+to optimized Metal kernels:
+
+- ``torch.ops.mlx.rope`` for rotary position embeddings
+- ``torch.ops.mlx.kv_cache_update`` for KV cache scatter (via MLX cache modules)
+- ``torch.ops.mlx.custom_sdpa`` for scaled dot-product attention with GQA
+
+Applied at export time before ``torch.export`` — the model code in ``model.py``
+stays backend-agnostic.
+"""
+
+import executorch.backends.mlx.custom_ops  # noqa: F401 — registers mlx:: ops
+import torch
+import torch.nn as nn
+from executorch.backends.mlx.llm.cache import (
+    KVCache as MLXKVCache,
+    RingBufferKVCache as MLXRingKVCache,
+)
+
+
+def _replace_attention_forward(attn: nn.Module) -> None:
+    """Replace a Gemma4Attention's forward with one that uses MLX custom ops."""
+    import types
+
+    def _mlx_forward(self, x: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
+        B, T, _ = x.shape
+        start_pos = input_pos[0].item()
+
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim)
+        raw_k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim)
+        if self.k_eq_v:
+            raw_v = raw_k
+        else:
+            raw_v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim)
+
+        q = self.q_norm(q)
+        k = self.k_norm(raw_k)
+        v = self.v_norm(raw_v)
+
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        # RoPE via mlx::rope.
+        if self.is_sliding:
+            q = torch.ops.mlx.rope(
+                q, self.head_dim, start_pos, False, self.rope_theta, 1.0, None
+            )
+            k = torch.ops.mlx.rope(
+                k, self.head_dim, start_pos, False, self.rope_theta, 1.0, None
+            )
+        else:
+            # Full-attention layers use proportional partial RoPE: only
+            # rotary_dim out of head_dim dimensions are rotated. Pass
+            # dims=rotary_dim and the non-zero frequencies as 1D freqs.
+            # MLX computes inv_freq = 1/freqs internally.
+            rotary_dim = int(self.head_dim * self.partial_rotary)
+            rotary_inv_freq = self.inv_freq[: rotary_dim // 2]
+            mlx_freqs = 1.0 / rotary_inv_freq
+            q = torch.ops.mlx.rope(q, rotary_dim, start_pos, False, 0.0, 1.0, mlx_freqs)
+            k = torch.ops.mlx.rope(k, rotary_dim, start_pos, False, 0.0, 1.0, mlx_freqs)
+
+        k_cache, v_cache = self.kv_cache.update(start_pos, k, v)
+
+        if self.is_sliding:
+            sdpa_mask = self.kv_cache.create_sliding_window_mask(start_pos, T)
+            y = torch.ops.mlx.custom_sdpa(
+                q,
+                k_cache,
+                v_cache,
+                start_pos=self.kv_cache.buffer_size - T,
+                attn_mask=sdpa_mask,
+                dropout_p=0.0,
+                is_causal=False,
+                scale=self.scaling,
+            )
+        else:
+            y = torch.ops.mlx.custom_sdpa(
+                q,
+                k_cache,
+                v_cache,
+                start_pos=start_pos,
+                dropout_p=0.0,
+                is_causal=True,
+                scale=self.scaling,
+            )
+
+        y = y.transpose(1, 2).contiguous().view(B, T, self.n_heads * self.head_dim)
+        return self.o_proj(y)
+
+    attn.forward = types.MethodType(_mlx_forward, attn)
+
+
+def _replace_layer_forward(layer: nn.Module) -> None:
+    """Replace Gemma4DecoderLayer's forward to remove mask parameters."""
+    import types
+
+    def _mlx_layer_forward(
+        self, x: torch.Tensor, input_pos: torch.Tensor
+    ) -> torch.Tensor:
+        residual = x
+        h = self.input_layernorm(x)
+        h = self.self_attn(h, input_pos)
+        h = self.post_attention_layernorm(h)
+        x = residual + h
+
+        residual = x
+        h = self.pre_feedforward_layernorm(x)
+        h = self.mlp(h)
+        h = self.post_feedforward_layernorm(h)
+        x = residual + h
+
+        return x * self.layer_scalar
+
+    layer.forward = types.MethodType(_mlx_layer_forward, layer)
+
+
+def _replace_model_forward(model: nn.Module) -> None:
+    """Replace the top-level Gemma4_31B forward with a sampler-free, mask-free
+    ``(tokens, input_pos) → (B, 1, V)`` variant.
+
+    MLX samples on the host, so the on-device sampler and temperature input
+    are dropped.  Each MLX attention builds its own mask via ``custom_sdpa``,
+    so ``_build_masks`` and the per-layer mask arguments are removed.
+    """
+    import types
+
+    def _mlx_model_forward(
+        self, tokens: torch.Tensor, input_pos: torch.Tensor
+    ) -> torch.Tensor:
+        x = self.embed_tokens(tokens) * self.embed_normalizer
+        for layer in self.layers:
+            x = layer(x, input_pos)
+        x = self.norm(x)
+        last = self.lm_head(x[:, -1, :]).float()
+        cap = self.logit_softcap.float()
+        return torch.tanh(last / cap) * cap
+
+    model.forward = types.MethodType(_mlx_model_forward, model)
+
+
+def mlx_source_transformations(
+    model: nn.Module,
+    dtype: torch.dtype = torch.bfloat16,
+) -> None:
+    """Apply MLX source transformations to a Gemma 4 31B model in-place.
+
+    Self-contained MLX adaptation. After calling this, the model has
+    signature ``(tokens, input_pos) → (B, 1, V)`` logits — no temperature,
+    no sampler, no attention masks.
+
+    - Replaces KV caches with MLX-optimized versions using ``mlx.kv_cache_update``
+    - Rewrites attention forward to use ``mlx.rope`` and ``mlx.custom_sdpa``
+    - Rewrites layer forward to drop mask parameters (each attention builds
+      its own mask via ``custom_sdpa``)
+    - Rewrites model forward to drop the sampler and ``_build_masks``
+    """
+    config = model.config
+
+    for layer in model.layers:
+        attn = layer.self_attn
+
+        if attn.is_sliding:
+            attn.kv_cache = MLXRingKVCache(
+                max_batch_size=1,
+                max_context_length=config.sliding_window,
+                n_heads=attn.n_kv_heads,
+                head_dim=attn.head_dim,
+                dtype=dtype,
+            )
+        else:
+            attn.kv_cache = MLXKVCache(
+                max_batch_size=1,
+                max_context_length=config.max_seq_len,
+                n_heads=attn.n_kv_heads,
+                head_dim=attn.head_dim,
+                enable_dynamic_shape=True,
+                dtype=dtype,
+            )
+
+        _replace_attention_forward(attn)
+        _replace_layer_forward(layer)
+
+    _replace_model_forward(model)
diff --git a/examples/models/gemma4_31b/model.md b/examples/models/gemma4_31b/model.md
index 8233b6d430e..13207bdbb06 100644
--- a/examples/models/gemma4_31b/model.md
+++ b/examples/models/gemma4_31b/model.md
@@ -102,6 +102,8 @@ Decoder norms per layer: `input_layernorm`, `post_attention_layernorm`,
 
 ## Methods exported (`export.py`)
 
+### CUDA (`--backend cuda`)
+
 | Method    | Input                                                      | Output (sampled) |
 |-----------|------------------------------------------------------------|------------------|
 | `decode`  | tokens `(1, 1)` + input_pos `(1,)` + temperature `(1,)`    | `(1, 1)` float   |
@@ -113,6 +115,23 @@ Both methods share the same KV-cache buffers via
 sampling on-device and returns a single token ID per call so the C++ runner
 only has to feed tokens.
 
+### MLX (`--backend mlx`)
+
+| Method    | Input                                    | Output           |
+|-----------|------------------------------------------|------------------|
+| `forward` | tokens `(1, T)` + input_pos `(T,)`, T∈[1, min(max_seq_len-1, 2×sliding_window)] | `(1, V)` logits |
+
+Single method with dynamic sequence length. Only the last token's logits
+are returned. The C++ runner samples on the host via `logits_to_token`
+with temperature support. Int4Tensor weights are converted to
+IntxUnpackedToInt8Tensor at pack time so the default `dequantize_affine →
+linear` dispatch produces the pattern MLX's `QuantizedLinearHandler` fuses
+into `QuantizedMatmulNode`. Source transforms (`mlx_source_transformations.py`)
+replace generic PyTorch ops with `mlx.rope`, `mlx.kv_cache_update`, and
+`mlx.custom_sdpa` for optimized Metal kernels.
+
+### Shared
+
 Prefill length is capped to the ring-buffer KV cache size
 (`2 × sliding_window`) to avoid duplicate wrapped indices in
 `index_copy_`. The C++ runner chunks longer prompts automatically using
@@ -130,9 +149,11 @@ Modules in `quant/`:
   `IntxUnpackedToInt8Tensor`) from fp weights.
 - **Serialization**: callers use torchao's safetensors integration
   (`torchao.prototype.safetensors`) directly — no wrapper module needed.
-- **Pack** (`pack.py` + `pack_cuda.py`): `pack_model` groups weights by
-  parent module, `pack_one` handles single weights. Per-module packers
-  dispatch by module type (`nn.Linear`, `nn.Embedding`, extensible for MoE).
+- **Pack** (`pack.py` + `pack_cuda.py` + `pack_mlx.py`): `pack_model` groups
+  weights by parent module, `pack_one` handles single weights. Per-module
+  packers dispatch by module type (`nn.Linear`, `nn.Embedding`). CUDA passes
+  Int4Tensor through (dispatch handled by `int4_dispatch.py`); MLX converts
+  Int4Tensor → IntxUnpackedToInt8Tensor and regroups per-axis embeddings.
 - **GGUF** (`gguf.py`): `unpack_gguf_tensor` / `iter_gguf_tensors` for
   loading community-quantized GGUF files (Q4_K, Q6_K).
 
@@ -145,11 +166,12 @@ quantize_and_save.py                    export.py / inference.py
      |                                       |
   quantize_weight()                     load (torchao safetensors)
      |                                       |
-  Int4Tensor / IntxUnpacked             Int4Tensor / IntxUnpacked (used directly)
-     |                                       |
-  save (torchao safetensors)            int4_dispatch routes to int4_plain_mm
+  Int4Tensor / IntxUnpacked             pack for backend:
      |                                       |
-  model.safetensors                     dp4a decode / dequant+cuBLAS prefill
+  save (torchao safetensors)            CUDA: Int4Tensor passed through
+     |                                    → int4_dispatch → dp4a / dequant+cuBLAS
+  model.safetensors                     MLX:  Int4Tensor → IntxUnpacked(int4)
+                                          → dequantize_affine → QuantizedMatmulNode
 ```
 
 `embed_tokens` and `lm_head` start tied; they are untied before
@@ -165,7 +187,7 @@ RoPE inv_freq buffers, and scalar constants are still on the meta device.
 them with real tensors:
 
 - KV caches → zeros in `dtype` (bf16 for inference, bf16 for export)
-- `inv_freq` → moved to target device (cos/sin computed on the fly per forward)
+- `inv_freq` → recomputed on target device (cos/sin computed on the fly per forward)
 - `embed_normalizer`, `logit_softcap`, `cache_positions` → scalar constants
 
 Called by `export.py` (device="cpu" for tracing) and `inference.py`
@@ -181,14 +203,17 @@ These exist solely to make the model exportable / efficient under ExecuTorch:
   `2 × sliding_window`) saves memory for long sequences — positions wrap
   via modulo and the attention mask reconstructs which slots are valid.
   Full-attention layers use a flat `Gemma4KVCache` sized to `max_seq_len`.
-  Both use `index_copy_(dim=2, ...)` for trace-friendly updates.
+  CUDA uses `index_copy_` for trace-friendly updates; MLX source transforms
+  replace both caches with `mlx.kv_cache_update`-backed equivalents.
 - **On-the-fly RoPE**: stores only `inv_freq` per layer, computes cos/sin
   via `torch.outer(positions, inv_freq)` each forward. Saves memory vs
   precomputed `[max_seq_len, head_dim]` tables (sliding uses full RoPE,
   full uses proportional partial RoPE — head_dim and θ differ).
-- **On-device Gumbel-max sampling** so the exported program emits a token
-  rather than a full logits tensor — keeps the runner GPU↔CPU traffic to a
-  single float per step.
+- **Last-logits-only**: `lm_head` always runs on `x[:, -1, :]`, avoiding a
+  `(1, T, 262144)` matmul during prefill.
+- **On-device Gumbel-max sampling** (CUDA) so the exported program emits a
+  token rather than logits — keeps GPU↔CPU traffic to a single float per
+  step. MLX samples on the host via `logits_to_token`.
 - **Final-logit softcap baked into the graph**, applied before sampling.
 - **Meta-device construction + assign-load** keeps peak memory small enough
   to load the 31B-parameter checkpoint on one machine.
diff --git a/examples/models/gemma4_31b/model.py b/examples/models/gemma4_31b/model.py
index b0eb4004c52..a690bd79230 100644
--- a/examples/models/gemma4_31b/model.py
+++ b/examples/models/gemma4_31b/model.py
@@ -251,25 +251,7 @@ def __init__(self, config: Gemma4_31BConfig, layer_idx: int):
         self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
         self.v_norm = RMSNormNoWeight(self.head_dim, eps=config.rms_norm_eps)
 
-        # Precomputed RoPE table for this layer (per-layer because head_dim
-        # and theta differ between sliding and full attention). For full
-        # attention layers we pass freq_base_dim=head_dim so the zero-padded
-        # On-the-fly RoPE: store only inv_freq, compute cos/sin per forward.
-        # Saves memory vs precomputed [max_seq_len, head_dim] tables.
-        if self.is_sliding:
-            rotary_dim = self.head_dim
-        else:
-            rotary_dim = int(self.head_dim * self.partial_rotary)
-        rope_angles = rotary_dim // 2
-        inv_freq_rotated = 1.0 / (
-            self.rope_theta ** (torch.arange(0, rotary_dim, 2).float() / self.head_dim)
-        )
-        nope_angles = self.head_dim // 2 - rope_angles
-        if nope_angles > 0:
-            inv_freq = torch.cat([inv_freq_rotated, torch.zeros(nope_angles)])
-        else:
-            inv_freq = inv_freq_rotated
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False)
 
         # KV cache. Sliding layers use a ring buffer (2x window) to save
         # memory; full layers use a flat buffer (max_seq_len).
@@ -289,6 +271,30 @@ def __init__(self, config: Gemma4_31BConfig, layer_idx: int):
                 use_index_copy=True,
             )
 
+    def _compute_inv_freq(self, device: Optional[torch.device] = None) -> torch.Tensor:
+        """Compute RoPE inverse-frequency table for this layer."""
+        if self.is_sliding:
+            rotary_dim = self.head_dim
+        else:
+            rotary_dim = int(self.head_dim * self.partial_rotary)
+        rope_angles = rotary_dim // 2
+        inv_freq_rotated = 1.0 / (
+            self.rope_theta
+            ** (
+                torch.arange(0, rotary_dim, 2, device=device, dtype=torch.float32)
+                / self.head_dim
+            )
+        )
+        nope_angles = self.head_dim // 2 - rope_angles
+        if nope_angles > 0:
+            return torch.cat(
+                [
+                    inv_freq_rotated,
+                    torch.zeros(nope_angles, device=device, dtype=torch.float32),
+                ]
+            )
+        return inv_freq_rotated
+
     def forward(
         self,
         x: torch.Tensor,
@@ -464,20 +470,17 @@ def forward(
         self,
         tokens: torch.LongTensor,
         input_pos: torch.LongTensor,
-        temperature: Optional[torch.Tensor] = None,
+        temperature: torch.Tensor,
     ) -> torch.Tensor:
         """Run the model.
 
         Args:
             tokens: (B, T) token IDs.
             input_pos: (T,) absolute positions for RoPE / KV cache.
-            temperature: optional 1-D float tensor controlling on-device sampling.
-                When provided, returns sampled tokens (B, 1) via Gumbel-max;
-                when None (e.g. eager eval), returns full logits (B, T, V) with
-                soft-capping applied so callers see post-cap values.
+            temperature: 1-D float tensor for Gumbel-max sampling.
 
         Returns:
-            (B, 1) token IDs when sampling, else (B, T, V) float32 logits.
+            (B, 1) sampled token IDs as float.
         """
         x = self.embed_tokens(tokens) * self.embed_normalizer
 
@@ -486,13 +489,6 @@ def forward(
             x = layer(x, input_pos, sliding_mask, full_mask)
 
         x = self.norm(x)
-
-        if temperature is None:
-            logits = self.lm_head(x).float()
-            cap = self.logit_softcap.float()
-            return torch.tanh(logits / cap) * cap
-
-        # Decode-time fast path: only materialize logits for the last token.
         last = self.lm_head(x[:, -1, :]).float()
         cap = self.logit_softcap.float()
         last = torch.tanh(last / cap) * cap
@@ -675,7 +671,9 @@ def materialize_runtime_buffers(
 
     for layer in model.layers:
         attn = layer.self_attn
-        attn.inv_freq = attn.inv_freq.to(device)
+        attn.register_buffer(
+            "inv_freq", attn._compute_inv_freq(device=device), persistent=False
+        )
 
     model.register_buffer(
         "embed_normalizer",
diff --git a/examples/models/gemma4_31b/quant/README.md b/examples/models/gemma4_31b/quant/README.md
index 31b1c43d574..2eacced4387 100644
--- a/examples/models/gemma4_31b/quant/README.md
+++ b/examples/models/gemma4_31b/quant/README.md
@@ -9,7 +9,8 @@ Quantization framework: **recipe → quantize → pack**.
 | `recipe.py` | **Policy** — what to quantize, what precision, which layers | nothing |
 | `quantize.py` | **Computation** — produces torchao subclass tensors | recipe, torchao |
 | `pack.py` | **Packing dispatch** — `pack_model` (bulk) and `pack_one` (streaming) | — |
-| `pack_cuda.py` | **CUDA packing** — converts Int4Tensor to tinygemm format | pack |
+| `pack_cuda.py` | **CUDA packing** — passes Int4Tensor/IntxUnpacked through for CUDA dispatch | pack |
+| `pack_mlx.py` | **MLX packing** — converts Int4Tensor → IntxUnpacked, regroups per-axis embeddings | pack |
 | `gguf.py` | **GGUF import** — unpacks Q4_K/Q6_K blocks to torchao subclasses | torchao |
 
 ## Data flow
@@ -48,7 +49,6 @@ The format is compatible with torchao's `save_pretrained` / `load_pretrained`.
 ## TODO
 
 - `pack_metal.py` — Metal backend packer.
-- `pack_mlx.py` — MLX backend packer.
 - `gguf.py` — extend with Q5_K, Q8_0 GGUF quant types.
 - Upstream `Int4TilePackedTo4dTensor.from_int4_tensor()` to torchao
   to replace the manual conversion in `pack_int4_for_cuda`.
diff --git a/examples/models/gemma4_31b/quant/__init__.py b/examples/models/gemma4_31b/quant/__init__.py
index 93efb69865f..7e9ab97a1bb 100644
--- a/examples/models/gemma4_31b/quant/__init__.py
+++ b/examples/models/gemma4_31b/quant/__init__.py
@@ -6,5 +6,6 @@
 
 from .pack import ModulePackerFn, pack_model, pack_one  # noqa: F401
 from .pack_cuda import DEFAULT_CUDA_PACKERS, load_and_pack_for_cuda  # noqa: F401
+from .pack_mlx import DEFAULT_MLX_PACKERS, load_and_pack_for_mlx  # noqa: F401
 from .quantize import dequantize_weight, quantize_model, quantize_weight  # noqa: F401
 from .recipe import QuantConfig, QuantRecipe, QuantRule  # noqa: F401
diff --git a/examples/models/gemma4_31b/quant/pack_mlx.py b/examples/models/gemma4_31b/quant/pack_mlx.py
new file mode 100644
index 00000000000..63aeca426a8
--- /dev/null
+++ b/examples/models/gemma4_31b/quant/pack_mlx.py
@@ -0,0 +1,198 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""MLX packer: convert quantized weights to MLX-compatible format.
+
+MLX's ``QuantizedLinearHandler`` matches ``dequantize_affine → linear``
+in the exported graph.  ``IntxUnpackedToInt8Tensor`` produces this
+pattern naturally, but ``Int4Tensor`` does not (its dispatch calls
+CUDA-specific mslk kernels).  So INT4 weights are converted to
+``IntxUnpackedToInt8Tensor(target_dtype=torch.int4)`` at pack time.
+
+The backend-agnostic ``pack_model`` dispatcher lives in ``pack.py``.
+"""
+
+import json
+
+import torch
+import torch.nn as nn
+
+from .pack import ModulePackerFn, pack_model  # noqa: F401
+
+_MLX_SUPPORTED_GROUP_SIZES = (128, 64, 32)
+
+
+# ---------------------------------------------------------------------------
+# Int4Tensor → IntxUnpackedToInt8Tensor conversion
+
+
+def _int4_to_intx_unpacked(w: torch.Tensor) -> torch.Tensor:
+    """Convert an ``Int4Tensor`` to ``IntxUnpackedToInt8Tensor``.
+
+    Int4Tensor stores qdata as nibble-packed uint8 ``(N, K/2)`` with
+    scale/zero transposed to ``(K//gs, N)``.  IntxUnpackedToInt8Tensor
+    stores qdata as int8 ``(N, K)`` with scale/zero as ``(N, K//gs)``.
+    """
+    from torchao.quantization import IntxUnpackedToInt8Tensor
+
+    # Unpack nibbles: packed = even | (odd << 4), unsigned [0, 15]
+    p = w.qdata.to(torch.uint8)
+    low = (p & 0x0F).to(torch.int8)
+    high = ((p >> 4) & 0x0F).to(torch.int8)
+    qdata = torch.stack([low, high], dim=-1).reshape(w.shape)
+
+    # Shift unsigned [0, 15] → signed [-8, 7]
+    qdata = qdata - 8
+
+    gs = w.block_size[-1]
+
+    # Transpose scale/zero from (K//gs, N) → (N, K//gs)
+    scale = w.scale.t().contiguous()
+    zero_point = (w.zero_point - 8).t().contiguous()
+
+    return IntxUnpackedToInt8Tensor(
+        qdata=qdata,
+        scale=scale,
+        zero_point=zero_point,
+        target_dtype=torch.int4,
+        block_size=(1, gs),
+        dtype=scale.dtype,
+        activation_quantization=None,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Embedding group_size regrouping
+
+
+def _mlx_group_size(gs: int, K: int) -> int:
+    """Find an MLX-compatible group_size for the given weight group_size.
+
+    If ``gs`` is already in {32, 64, 128}, return it.  Otherwise find the
+    largest supported group_size that divides ``gs`` so per-axis scales can
+    be repeated to fill finer groups.
+    """
+    if gs in _MLX_SUPPORTED_GROUP_SIZES:
+        return gs
+    for candidate in _MLX_SUPPORTED_GROUP_SIZES:
+        if gs % candidate == 0 and K % candidate == 0:
+            return candidate
+    raise ValueError(
+        f"MLX requires group_size in {set(_MLX_SUPPORTED_GROUP_SIZES)} "
+        f"(or a multiple thereof), got {gs}"
+    )
+
+
+def _regroup_intx(w: torch.Tensor, new_gs: int) -> torch.Tensor:
+    """Regroup an ``IntxUnpackedToInt8Tensor`` to a finer group_size."""
+    from torchao.quantization import IntxUnpackedToInt8Tensor
+
+    old_gs = w.block_size[-1]
+    if old_gs % new_gs != 0:
+        raise ValueError(
+            f"new group_size {new_gs} must evenly divide old group_size {old_gs}"
+        )
+    repeat_factor = old_gs // new_gs
+    N = w.qdata.shape[0]
+    n_groups = w.qdata.shape[-1] // new_gs
+
+    scale = w.scale.repeat_interleave(repeat_factor, dim=-1).reshape(N, n_groups)
+    zero_point = w.zero_point.repeat_interleave(repeat_factor, dim=-1).reshape(
+        N, n_groups
+    )
+
+    return IntxUnpackedToInt8Tensor(
+        qdata=w.qdata,
+        scale=scale,
+        zero_point=zero_point,
+        target_dtype=w.target_dtype,
+        block_size=(1, new_gs),
+        dtype=w.dtype,
+        activation_quantization=w.activation_quantization,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Per-module packer
+
+
+def pack_for_mlx(module: nn.Module, weights: dict[str, torch.Tensor]) -> None:
+    """Pack a quantized weight for MLX.
+
+    ``Int4Tensor`` is converted to ``IntxUnpackedToInt8Tensor`` so the
+    default dispatch produces the ``dequantize_affine → linear`` pattern
+    MLX expects.  Regroups to a compatible group_size when needed (e.g.
+    per-axis group_size=5376 → group_size=128) since MLX's
+    ``parse_dequant_node`` only accepts group_size in {32, 64, 128}.
+    """
+    from torchao.quantization import IntxUnpackedToInt8Tensor
+    from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor
+
+    w = weights["weight"]
+    if isinstance(w, Int4Tensor):
+        w = _int4_to_intx_unpacked(w)
+    if isinstance(w, IntxUnpackedToInt8Tensor):
+        gs = w.block_size[-1]
+        K = w.qdata.shape[-1]
+        target_gs = _mlx_group_size(gs, K)
+        if target_gs != gs:
+            w = _regroup_intx(w, target_gs)
+    module.weight = nn.Parameter(w, requires_grad=False)
+
+
+DEFAULT_MLX_PACKERS: dict[type, ModulePackerFn] = {
+    nn.Linear: pack_for_mlx,
+    nn.Embedding: pack_for_mlx,
+}
+
+
+# ---------------------------------------------------------------------------
+# Load + pack (I/O wrapper)
+
+
+def load_and_pack_for_mlx(
+    path: str,
+    model: nn.Module,
+    packers: dict[type, ModulePackerFn] | None = None,
+) -> None:
+    """Load a quantized safetensors file and pack for MLX.
+
+    Streams one weight at a time via torchao's safetensors support.
+    """
+    from safetensors import safe_open
+    from torchao.prototype.safetensors.safetensors_support import (
+        unflatten_tensor_state_dict,
+    )
+
+    from .pack import pack_one
+
+    _packers = packers or DEFAULT_MLX_PACKERS
+    with safe_open(path, framework="pt", device="cpu") as f:
+        metadata = f.metadata()
+        all_keys = list(f.keys())
+        tensor_names = json.loads(metadata.get("tensor_names", "[]"))
+
+        for name in tensor_names:
+            parts = name.rsplit(".", 1)
+            module_fqn = parts[0] if len(parts) > 1 else ""
+            weight_name = parts[-1]
+            prefix = (
+                f"{module_fqn}._{weight_name}_" if module_fqn else f"_{weight_name}_"
+            )
+            partial = {}
+            for key in all_keys:
+                if key.startswith(prefix) or key == name:
+                    partial[key] = f.get_tensor(key)
+            result, _ = unflatten_tensor_state_dict(partial, metadata)
+            for fqn, value in result.items():
+                pack_one(model, fqn, value, _packers)
+
+    for fqn, p in model.named_parameters():
+        if p.device.type == "meta":
+            raise RuntimeError(
+                f"Weight '{fqn}' not found in checkpoint "
+                f"(model/checkpoint version mismatch?)"
+            )
diff --git a/examples/models/gemma4_31b/quant/tests/test_pack_mlx.py b/examples/models/gemma4_31b/quant/tests/test_pack_mlx.py
new file mode 100644
index 00000000000..ffb2e0e2dd3
--- /dev/null
+++ b/examples/models/gemma4_31b/quant/tests/test_pack_mlx.py
@@ -0,0 +1,219 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Unit tests for quant/pack_mlx.py. No CUDA or MLX hardware required."""
+
+import unittest
+
+import torch
+import torch.nn as nn
+
+from executorch.examples.models.gemma4_31b.quant.pack import pack_model
+from executorch.examples.models.gemma4_31b.quant.pack_mlx import (
+    _int4_to_intx_unpacked,
+    _mlx_group_size,
+    DEFAULT_MLX_PACKERS,
+    pack_for_mlx,
+)
+from executorch.examples.models.gemma4_31b.quant.quantize import (
+    dequantize_weight,
+    quantize_weight,
+)
+from executorch.examples.models.gemma4_31b.quant.recipe import QuantConfig
+
+
+class TestInt4ToIntxConversion(unittest.TestCase):
+    """Int4Tensor → IntxUnpackedToInt8Tensor conversion."""
+
+    def test_symmetric_dequant_matches(self):
+        """Converted weight dequantizes to same values as original."""
+        torch.manual_seed(0)
+        weight = torch.randn(64, 128, dtype=torch.bfloat16)
+        config = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max")
+        int4_w = quantize_weight(weight, config)
+        intx_w = _int4_to_intx_unpacked(int4_w)
+
+        int4_dense = dequantize_weight(int4_w, torch.float32)
+        intx_dense = dequantize_weight(intx_w, torch.float32)
+        self.assertTrue(
+            torch.allclose(int4_dense, intx_dense, atol=1e-5),
+            f"max diff: {(int4_dense - intx_dense).abs().max():.6g}",
+        )
+
+    def test_asymmetric_dequant_matches(self):
+        torch.manual_seed(0)
+        weight = torch.randn(64, 128, dtype=torch.bfloat16)
+        config = QuantConfig(bits=4, group_size=32, symmetric=False, method="min_max")
+        int4_w = quantize_weight(weight, config)
+        intx_w = _int4_to_intx_unpacked(int4_w)
+
+        int4_dense = dequantize_weight(int4_w, torch.float32)
+        intx_dense = dequantize_weight(intx_w, torch.float32)
+        self.assertTrue(
+            torch.allclose(int4_dense, intx_dense, atol=1e-5),
+            f"max diff: {(int4_dense - intx_dense).abs().max():.6g}",
+        )
+
+    def test_output_type_and_shape(self):
+        from torchao.quantization import IntxUnpackedToInt8Tensor
+
+        torch.manual_seed(0)
+        config = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max")
+        int4_w = quantize_weight(torch.randn(128, 256, dtype=torch.bfloat16), config)
+        intx_w = _int4_to_intx_unpacked(int4_w)
+
+        self.assertIsInstance(intx_w, IntxUnpackedToInt8Tensor)
+        self.assertEqual(intx_w.shape, torch.Size([128, 256]))
+        self.assertEqual(intx_w.qdata.shape, torch.Size([128, 256]))
+        self.assertEqual(intx_w.target_dtype, torch.int4)
+
+    def test_different_group_sizes(self):
+        torch.manual_seed(0)
+        for gs in (32, 64, 128):
+            with self.subTest(group_size=gs):
+                config = QuantConfig(
+                    bits=4, group_size=gs, symmetric=True, method="min_max"
+                )
+                int4_w = quantize_weight(
+                    torch.randn(64, 256, dtype=torch.bfloat16), config
+                )
+                intx_w = _int4_to_intx_unpacked(int4_w)
+                self.assertEqual(intx_w.shape, torch.Size([64, 256]))
+
+    def test_matmul_approximates_original(self):
+        torch.manual_seed(0)
+        weight = torch.randn(256, 128, dtype=torch.bfloat16)
+        x = torch.randn(1, 128, dtype=torch.bfloat16)
+        original_out = torch.nn.functional.linear(x, weight)
+
+        config = QuantConfig(bits=4, group_size=32, symmetric=False, method="min_max")
+        int4_w = quantize_weight(weight, config)
+        intx_w = _int4_to_intx_unpacked(int4_w)
+        packed_out = torch.nn.functional.linear(x, intx_w.dequantize())
+
+        rel_error = (
+            packed_out.float() - original_out.float()
+        ).abs().mean() / original_out.float().abs().mean()
+        self.assertLess(rel_error.item(), 0.15)
+
+
+class TestPackLinearForMlx(unittest.TestCase):
+    def test_int4_converts_to_intx(self):
+        from torchao.quantization import IntxUnpackedToInt8Tensor
+
+        module = nn.Linear(128, 64, bias=False)
+        config = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max")
+        w = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), config)
+        pack_for_mlx(module, {"weight": w})
+
+        self.assertIsInstance(module.weight.data, IntxUnpackedToInt8Tensor)
+        self.assertEqual(module.weight.shape, torch.Size([64, 128]))
+        self.assertFalse(module.weight.requires_grad)
+
+    def test_int8_passes_through(self):
+        from torchao.quantization import IntxUnpackedToInt8Tensor
+
+        module = nn.Linear(128, 64, bias=False)
+        config = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max")
+        w = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), config)
+        self.assertIsInstance(w, IntxUnpackedToInt8Tensor)
+        pack_for_mlx(module, {"weight": w})
+
+        self.assertIsInstance(module.weight.data, IntxUnpackedToInt8Tensor)
+        self.assertEqual(module.weight.shape, torch.Size([64, 128]))
+
+    def test_regroup_preserves_dequant(self):
+        """Linear with non-standard group_size regroups and dequantizes correctly."""
+        torch.manual_seed(0)
+        weight = torch.randn(64, 256, dtype=torch.bfloat16)
+        config = QuantConfig(bits=8, group_size=256, symmetric=True, method="min_max")
+        w = quantize_weight(weight, config)
+        before = dequantize_weight(w, torch.float32)
+
+        module = nn.Linear(256, 64, bias=False)
+        pack_for_mlx(module, {"weight": w})
+
+        self.assertEqual(module.weight.data.block_size, (1, 128))
+        after = dequantize_weight(module.weight.data, torch.float32)
+        self.assertTrue(
+            torch.allclose(before, after, atol=1e-5),
+            f"max diff: {(before - after).abs().max():.6g}",
+        )
+
+
+class TestMlxGroupSize(unittest.TestCase):
+    def test_passthrough(self):
+        for gs in (32, 64, 128):
+            self.assertEqual(_mlx_group_size(gs, 256), gs)
+
+    def test_regroup_5376(self):
+        self.assertEqual(_mlx_group_size(5376, 5376), 128)
+
+    def test_regroup_256(self):
+        self.assertEqual(_mlx_group_size(256, 256), 128)
+
+    def test_rejects_indivisible(self):
+        with self.assertRaises(ValueError):
+            _mlx_group_size(48, 48)
+
+
+class TestPackEmbeddingForMlx(unittest.TestCase):
+    def test_compatible_passes_through(self):
+        module = nn.Embedding(100, 64)
+        config = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max")
+        w = quantize_weight(torch.randn(100, 64, dtype=torch.bfloat16), config)
+        pack_for_mlx(module, {"weight": w})
+        self.assertEqual(module.weight.shape, torch.Size([100, 64]))
+
+    def test_per_axis_regroups(self):
+        module = nn.Embedding(50, 256)
+        config = QuantConfig(bits=8, group_size=256, symmetric=True, method="min_max")
+        w = quantize_weight(torch.randn(50, 256, dtype=torch.bfloat16), config)
+        pack_for_mlx(module, {"weight": w})
+        self.assertEqual(module.weight.shape, torch.Size([50, 256]))
+        self.assertEqual(module.weight.data.block_size, (1, 128))
+
+    def test_int4_converts_to_intx(self):
+        from torchao.quantization import IntxUnpackedToInt8Tensor
+
+        module = nn.Embedding(100, 64)
+        config = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max")
+        w = quantize_weight(torch.randn(100, 64, dtype=torch.bfloat16), config)
+        pack_for_mlx(module, {"weight": w})
+        self.assertIsInstance(module.weight.data, IntxUnpackedToInt8Tensor)
+        self.assertEqual(module.weight.shape, torch.Size([100, 64]))
+
+
+class TestPackModelMlx(unittest.TestCase):
+    def test_mixed_precision(self):
+        q4 = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max")
+        q8 = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max")
+        w4 = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), q4)
+        w8 = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), q8)
+
+        state_dict = {
+            "q_proj.weight": w4,
+            "v_proj.weight": w8,
+            "norm.weight": torch.randn(64, dtype=torch.bfloat16),
+        }
+
+        with torch.device("meta"):
+            model = nn.ModuleDict(
+                {
+                    "q_proj": nn.Linear(128, 64, bias=False),
+                    "v_proj": nn.Linear(128, 64, bias=False),
+                    "norm": nn.LayerNorm(64, bias=False),
+                }
+            )
+        pack_model(model, state_dict, DEFAULT_MLX_PACKERS)
+
+        self.assertEqual(model.q_proj.weight.shape, torch.Size([64, 128]))
+        self.assertEqual(model.v_proj.weight.shape, torch.Size([64, 128]))
+        self.assertEqual(model.norm.weight.shape, torch.Size([64]))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/examples/models/gemma4_31b/sampler.py b/examples/models/gemma4_31b/sampler.py
index 45e4e17887a..690344fd2e4 100644
--- a/examples/models/gemma4_31b/sampler.py
+++ b/examples/models/gemma4_31b/sampler.py
@@ -8,33 +8,26 @@
 
 Mirrors ``examples/models/qwen3_5_moe/sampler.py``: a single-output sampler
 that lets one exported program be re-driven with different temperatures
-without re-export. ``temperature=None`` is a no-op (returns logits).
+without re-export.
 """
 
-from typing import Optional
-
 import torch
 
 
 def sample(
     logits: torch.Tensor,
-    temperature: Optional[torch.Tensor] = None,
+    temperature: torch.Tensor,
 ) -> torch.Tensor:
     """Draw a single token per batch row using the Gumbel-max trick.
 
     Args:
         logits: ``[B, V]`` float32 logits (already soft-capped if applicable).
         temperature: 0-D or 1-D float tensor; clamped to >= 1e-6 so a 0
-            temperature still works ("near-greedy"). When ``None`` the call
-            short-circuits and returns ``logits`` unchanged.
+            temperature still works ("near-greedy").
 
     Returns:
-        ``[B, 1]`` float32 token IDs (``argmax(logits/T + gumbel_noise)``),
-        or the unmodified logits when ``temperature`` is ``None``.
+        ``[B, 1]`` float32 token IDs (``argmax(logits/T + gumbel_noise)``).
     """
-    if temperature is None:
-        return logits
-
     logits = logits / temperature.clamp(min=1e-6)
     noise = torch.rand_like(logits)
     gumbel = -torch.log(-torch.log(noise + 1e-20) + 1e-20)
diff --git a/examples/models/gemma4_31b/tests/test_cuda_pipeline.py b/examples/models/gemma4_31b/tests/test_cuda_pipeline.py
index 0ff28aac415..505d6f7bdc1 100644
--- a/examples/models/gemma4_31b/tests/test_cuda_pipeline.py
+++ b/examples/models/gemma4_31b/tests/test_cuda_pipeline.py
@@ -108,30 +108,27 @@ def test_chunked_prefill_matches_sequential(self):
         torch.manual_seed(0)
         prompt = torch.randint(0, config.vocab_size, (1, prompt_len), device="cuda")
 
+        temp = torch.tensor([1e-6], dtype=torch.float32, device="cuda")
+
         with torch.no_grad():
             for i in range(prompt_len):
                 tok = prompt[:, i : i + 1]
                 pos = torch.tensor([i], dtype=torch.long, device="cuda")
-                logits_seq = model_seq(tok, pos, None)
+                token_seq = model_seq(tok, pos, temp)
 
         with torch.no_grad():
             chunk1 = prompt[:, :buf_size]
             pos1 = torch.arange(buf_size, dtype=torch.long, device="cuda")
-            model_chunk(chunk1, pos1, None)
+            model_chunk(chunk1, pos1, temp)
 
             chunk2 = prompt[:, buf_size:]
             pos2 = torch.arange(buf_size, prompt_len, dtype=torch.long, device="cuda")
-            logits_chunk = model_chunk(chunk2, pos2, None)
-
-        max_diff = (logits_seq[0, -1].float() - logits_chunk[0, -1].float()).abs().max()
-        self.assertTrue(
-            torch.allclose(
-                logits_seq[0, -1].float(),
-                logits_chunk[0, -1].float(),
-                atol=1e-2,
-                rtol=1e-3,
-            ),
-            f"Chunked prefill diverged: max_diff={max_diff:.4g}",
+            token_chunk = model_chunk(chunk2, pos2, temp)
+
+        self.assertEqual(
+            int(token_seq.item()),
+            int(token_chunk.item()),
+            "Chunked prefill produced different token than sequential",
         )
 
 
diff --git a/examples/models/gemma4_31b/tests/test_mlx_pipeline.py b/examples/models/gemma4_31b/tests/test_mlx_pipeline.py
new file mode 100644
index 00000000000..0e62ab88e4b
--- /dev/null
+++ b/examples/models/gemma4_31b/tests/test_mlx_pipeline.py
@@ -0,0 +1,248 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""End-to-end MLX backend tests for the Gemma 4 31B-IT pipeline.
+
+Tests quantize → save → load → pack-for-MLX on a tiny model.
+No CUDA or MLX hardware required.
+
+Usage:
+    python -m pytest examples/models/gemma4_31b/tests/test_mlx_pipeline.py -v
+"""
+
+import json
+import os
+import tempfile
+import unittest
+
+import torch
+import torch.nn as nn
+
+from executorch.examples.models.gemma4_31b.model import Gemma4_31B
+from executorch.examples.models.gemma4_31b.quant import (
+    DEFAULT_MLX_PACKERS,
+    pack_model,
+    QuantConfig,
+    quantize_model,
+    QuantRecipe,
+    QuantRule,
+)
+from executorch.examples.models.gemma4_31b.tests.test_pipeline import (
+    build_random_tiny_model,
+    config_dict,
+    save_checkpoint,
+    TINY_CONFIG,
+)
+
+_INT4 = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max")
+_INT8 = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max")
+_INT8_PER_AXIS = QuantConfig(
+    bits=8, group_size=TINY_CONFIG.hidden_size, symmetric=True, method="min_max"
+)
+_EDGE_LAYERS = set(range(3))
+
+TINY_SENSITIVE_RECIPE = QuantRecipe(
+    rules=[
+        QuantRule(r"embed_tokens\.weight", _INT8_PER_AXIS),
+        QuantRule(r".*norm\.weight", None),
+        QuantRule(r".*\.(v_proj|down_proj)\.weight", _INT8, layers=_EDGE_LAYERS),
+        QuantRule(r".*\.weight", _INT4),
+    ]
+)
+
+
+class TestMlxPipeline(unittest.TestCase):
+    """End-to-end: quantize → pack for MLX → forward."""
+
+    def test_pack_for_mlx(self):
+        """Quantize with sensitive recipe, pack for MLX, no meta weights."""
+        model = build_random_tiny_model()
+        model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone())
+        state_dict = quantize_model(model, TINY_SENSITIVE_RECIPE)
+
+        with torch.device("meta"):
+            model = Gemma4_31B(TINY_CONFIG)
+        model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone())
+        pack_model(model, state_dict, DEFAULT_MLX_PACKERS)
+
+        for fqn, p in model.named_parameters():
+            self.assertNotEqual(p.device.type, "meta", f"Weight '{fqn}' still on meta")
+
+    def test_forward_after_pack(self):
+        """Model produces valid output after MLX packing."""
+        model = build_random_tiny_model()
+        model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone())
+        state_dict = quantize_model(model, TINY_SENSITIVE_RECIPE)
+
+        with torch.device("meta"):
+            model = Gemma4_31B(TINY_CONFIG)
+        model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone())
+        pack_model(model, state_dict, DEFAULT_MLX_PACKERS)
+        model.eval()
+
+        from executorch.examples.models.gemma4_31b.model import (
+            materialize_runtime_buffers,
+        )
+
+        materialize_runtime_buffers(model, dtype=torch.bfloat16)
+
+        tokens = torch.randint(0, TINY_CONFIG.vocab_size, (1, 1))
+        input_pos = torch.tensor([0], dtype=torch.long)
+        temp = torch.tensor([1e-6], dtype=torch.float32)
+
+        with torch.no_grad():
+            out = model(tokens, input_pos, temp)
+
+        self.assertEqual(out.shape, torch.Size([1, 1]))
+        self.assertFalse(torch.isnan(out).any())
+
+    def test_multi_token_forward(self):
+        model = build_random_tiny_model()
+        model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone())
+        state_dict = quantize_model(model, TINY_SENSITIVE_RECIPE)
+
+        with torch.device("meta"):
+            model = Gemma4_31B(TINY_CONFIG)
+        model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone())
+        pack_model(model, state_dict, DEFAULT_MLX_PACKERS)
+        model.eval()
+
+        from executorch.examples.models.gemma4_31b.model import (
+            materialize_runtime_buffers,
+        )
+
+        materialize_runtime_buffers(model, dtype=torch.bfloat16)
+
+        seq_len = 4
+        tokens = torch.randint(0, TINY_CONFIG.vocab_size, (1, seq_len))
+        input_pos = torch.arange(seq_len, dtype=torch.long)
+        temp = torch.tensor([1e-6], dtype=torch.float32)
+
+        with torch.no_grad():
+            out = model(tokens, input_pos, temp)
+
+        self.assertEqual(out.shape, torch.Size([1, 1]))
+        self.assertFalse(torch.isnan(out).any())
+
+    def test_source_transforms_forward(self):
+        """Model produces valid output after MLX source transforms."""
+        model = build_random_tiny_model()
+        model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone())
+        state_dict = quantize_model(model, TINY_SENSITIVE_RECIPE)
+
+        with torch.device("meta"):
+            model = Gemma4_31B(TINY_CONFIG)
+        model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone())
+        pack_model(model, state_dict, DEFAULT_MLX_PACKERS)
+        model.eval()
+
+        from executorch.examples.models.gemma4_31b.mlx_source_transformations import (
+            mlx_source_transformations,
+        )
+        from executorch.examples.models.gemma4_31b.model import (
+            materialize_runtime_buffers,
+        )
+
+        mlx_source_transformations(model, dtype=torch.bfloat16)
+        materialize_runtime_buffers(model, dtype=torch.bfloat16)
+
+        # After source transforms: signature is (tokens, input_pos) → (B, 1, V)
+        # Single-token decode
+        tokens = torch.randint(0, TINY_CONFIG.vocab_size, (1, 1))
+        input_pos = torch.tensor([0], dtype=torch.long)
+        with torch.no_grad():
+            out = model(tokens, input_pos)
+        self.assertEqual(out.shape, torch.Size([1, TINY_CONFIG.vocab_size]))
+        self.assertFalse(torch.isnan(out).any())
+        self.assertFalse(torch.isinf(out).any())
+
+        # Multi-token prefill
+        seq_len = 4
+        tokens = torch.randint(0, TINY_CONFIG.vocab_size, (1, seq_len))
+        input_pos = torch.arange(seq_len, dtype=torch.long)
+        with torch.no_grad():
+            out = model(tokens, input_pos)
+        self.assertEqual(out.shape, torch.Size([1, TINY_CONFIG.vocab_size]))
+        self.assertFalse(torch.isnan(out).any())
+
+    def test_source_transforms_use_mlx_ops(self):
+        """Verify the traced graph contains the expected MLX custom ops.
+
+        Each attention layer should produce:
+          - 2× ``mlx.rope`` (q and k)
+          - 2× ``mlx.kv_cache_update`` (k and v)
+          - 1× ``mlx.custom_sdpa``
+        """
+        from executorch.examples.models.gemma4_31b.mlx_source_transformations import (
+            mlx_source_transformations,
+        )
+        from executorch.examples.models.gemma4_31b.model import (
+            materialize_runtime_buffers,
+        )
+        from torch.export import Dim, export
+
+        model = build_random_tiny_model()
+        model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone())
+        state_dict = quantize_model(model, TINY_SENSITIVE_RECIPE)
+
+        with torch.device("meta"):
+            model = Gemma4_31B(TINY_CONFIG)
+        model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone())
+        pack_model(model, state_dict, DEFAULT_MLX_PACKERS)
+        model.eval()
+
+        mlx_source_transformations(model, dtype=torch.bfloat16)
+        materialize_runtime_buffers(model, dtype=torch.bfloat16)
+
+        # Trace with dynamic seq_len matching the MLX export shape.
+        seq_dim = Dim("seq", min=1, max=8)
+        ep = export(
+            model,
+            (torch.tensor([[1, 2]]), torch.tensor([0, 1])),
+            dynamic_shapes=({1: seq_dim}, {0: seq_dim}),
+            strict=True,
+        )
+
+        op_counts = {"rope": 0, "kv_cache_update": 0, "custom_sdpa": 0}
+        for node in ep.graph.nodes:
+            if node.op != "call_function":
+                continue
+            name = str(node.target)
+            for op in op_counts:
+                if f"mlx.{op}" in name:
+                    op_counts[op] += 1
+
+        n_layers = TINY_CONFIG.num_hidden_layers
+        self.assertEqual(op_counts["rope"], 2 * n_layers, f"got {op_counts}")
+        self.assertEqual(op_counts["kv_cache_update"], 2 * n_layers, f"got {op_counts}")
+        self.assertEqual(op_counts["custom_sdpa"], n_layers, f"got {op_counts}")
+
+    def test_export_to_pte(self):
+        """Full export: quantize → pack → export with MLXPartitioner."""
+        try:
+            from executorch.backends.mlx import MLXPartitioner  # noqa: F401
+        except ImportError:
+            self.skipTest("MLX backend not available")
+
+        from executorch.examples.models.gemma4_31b.export import (
+            export_and_lower,
+            load_prequantized_model,
+        )
+
+        with tempfile.TemporaryDirectory() as ckpt_dir, tempfile.TemporaryDirectory() as out_dir:
+            save_checkpoint(ckpt_dir)
+            with open(os.path.join(ckpt_dir, "config.json"), "w") as f:
+                json.dump(config_dict(), f)
+
+            model, config = load_prequantized_model(
+                ckpt_dir, max_seq_len=TINY_CONFIG.max_seq_len, backend="mlx"
+            )
+            export_and_lower(model, config, out_dir, backend="mlx")
+            self.assertTrue(os.path.exists(os.path.join(out_dir, "model.pte")))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/examples/models/llama/norm.py b/examples/models/llama/norm.py
index e424ee0361a..ec92b353eb4 100644
--- a/examples/models/llama/norm.py
+++ b/examples/models/llama/norm.py
@@ -41,20 +41,65 @@ def forward(self, x):
         return output * self.weight.type_as(x)
 
 
-class ScalelessRMSNorm(torch.nn.RMSNorm):
-    """RMSNorm with weight hardcoded to ones and not trainable.
+class ScalelessRMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def forward(self, x):
+        x_float = x.float()
+        return (
+            x_float * torch.rsqrt((x_float * x_float).mean(-1, keepdim=True) + self.eps)
+        ).type_as(x)
 
-    Equivalent to a scaleless RMSNorm (no learnable scaling) but implemented as a
-    torch.nn.RMSNorm so the op composes/decomposes cleanly for backends like QNN
-    instead of being expressed as a hand-rolled decomposition.
-    """
 
+class RMSNormCoreML(torch.nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__(dim, eps)
+        """
+        CoreML-friendly RMSNorm — uses `torch.linalg.vector_norm` so the op is
+        preserved in the CoreML graph for numerical stability.
+
+        Args:
+            dim (int): The dimension of the input tensor.
+            eps (float, optional): Floor on the L2-norm denominator
+                (`clamp_min(‖x‖₂, √(dim·eps))`). Prevents `0/0 = NaN` on
+                zero-padded positions and matches standard RMSNorm's
+                `rsqrt(mean(x²) + eps)` semantics on a zero input. Must be > 0.
+
+        Attributes:
+            eps (float): Floor coefficient consumed by `_norm`.
+            weight (nn.Parameter): Learnable scaling parameter.
+        """
+        super().__init__()
+        assert eps > 0, (
+            "RMSNormCoreML requires eps > 0; eps=0 collapses the denominator "
+            "floor and produces NaN on zero-padded positions"
+        )
         self.dim = dim
-        with torch.no_grad():
-            self.weight.fill_(1.0)
-        self.weight.requires_grad = False
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        # Floor the denominator to avoid 0 / 0 = NaN on zero-padded positions
+        # (chunked prefill in StaticAttentionIOManager pads each chunk to
+        # input_len with zeros). Use sqrt(dim * eps) so the floor matches
+        # standard RMSNorm's eps semantics (`rsqrt(mean(x²) + eps)`) and is
+        # large enough to survive fp16 (1e-6 alone underflows in fp16).
+        floor_val = torch.sqrt(torch.tensor(self.dim * self.eps, dtype=x.dtype))
+        norm_val = torch.clamp_min(
+            torch.linalg.vector_norm(x, dim=-1, keepdim=True), floor_val
+        )
+        rms_norm_eps0 = (
+            x
+            * torch.sqrt(torch.tensor(self.dim, dtype=x.dtype))
+            * torch.reciprocal(norm_val)
+        )
+        return rms_norm_eps0
+
+    def forward(self, x):
+        output = self._norm(x)
+        return output * self.weight
 
 
 class RMSNormWithInputScale(torch.nn.Module):
@@ -83,3 +128,37 @@ def forward(self, hidden_states: torch.Tensor, gate: torch.Tensor) -> torch.Tens
         hidden_states = self.weight * hidden_states.to(input_dtype)
         hidden_states = hidden_states * F.silu(gate.to(torch.float32))
         return hidden_states.to(input_dtype)
+
+
+def replace_rms_norm_for_coreml_(model: torch.nn.Module) -> torch.nn.Module:
+    """In-place: walk `model` and swap every RMSNorm-family module for RMSNormCoreML.
+
+    Mirrors the post-construction transform pattern used by torchao's
+    `quantize_(model, config)`: instead of threading a `use_coreml_norm` flag
+    through every norm construction site, build the model with the standard
+    norms and then call this once before CoreML export. Trained scale weights
+    are preserved.
+
+    Swaps these classes (everything else is left alone):
+      * `RMSNorm` (this module)
+      * `ScalelessRMSNorm` (this module — no-op weight)
+      * `torch.nn.RMSNorm` (used for affine q_norm/k_norm in StaticAttention)
+    """
+    for name, mod in list(model.named_modules()):
+        if not isinstance(mod, (RMSNorm, ScalelessRMSNorm, torch.nn.RMSNorm)):
+            continue
+        # All three carry the normalized dim either as `dim` or in `normalized_shape[-1]`.
+        dim = getattr(mod, "dim", None) or mod.normalized_shape[-1]
+        eps = getattr(mod, "eps", 1e-6) or 1e-6
+        new = RMSNormCoreML(dim, eps=eps)
+        # Preserve trained scale (no-op for ScalelessRMSNorm).
+        if getattr(mod, "weight", None) is not None:
+            new.weight = mod.weight
+        # Locate parent module via the dotted name and rebind the attribute.
+        if "." in name:
+            parent_name, attr = name.rsplit(".", 1)
+            parent = model.get_submodule(parent_name)
+        else:
+            parent, attr = model, name
+        setattr(parent, attr, new)
+    return model
diff --git a/examples/models/llama/source_transformation/sdpa.py b/examples/models/llama/source_transformation/sdpa.py
index 0285f3562cb..b10f684ccc0 100644
--- a/examples/models/llama/source_transformation/sdpa.py
+++ b/examples/models/llama/source_transformation/sdpa.py
@@ -69,7 +69,7 @@ def forward(
                 0,  # dropout probability. Ignored by the code
                 True,  # is_causal
             )
-        return output.reshape(bsz, seqlen, self.dim).to(dtype=input_dtype)
+        return output.view(bsz, seqlen, self.dim).to(dtype=input_dtype)
 
 
 def _replace_sdpa_with_custom_op(
@@ -198,7 +198,7 @@ def forward(
                 v_scale_fp32,
             )
 
-        return output.reshape(bsz, seqlen, self.dim)
+        return output.view(bsz, seqlen, self.dim)
 
 
 def _update_attention_module_with_quantized_sdpa(
diff --git a/examples/nxp/setup.sh b/examples/nxp/setup.sh
index e6de05b9f47..113b08d24ec 100755
--- a/examples/nxp/setup.sh
+++ b/examples/nxp/setup.sh
@@ -8,7 +8,7 @@ set -u
 EIQ_PYPI_URL="${EIQ_PYPI_URL:-https://eiq.nxp.com/repository}"
 
 # Install eIQ Neutron dependencies - SDK and simulator
-pip install --index-url ${EIQ_PYPI_URL} eiq-neutron-sdk==3.1.0 eiq_nsys
+pip install --index-url ${EIQ_PYPI_URL} eiq-neutron-sdk==3.1.1 eiq_nsys
 
 # Get the directory of the current script
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
diff --git a/examples/qualcomm/executor_runner/qnn_executor_runner.cpp b/examples/qualcomm/executor_runner/qnn_executor_runner.cpp
index df6a8ecbd79..a35a496f22b 100644
--- a/examples/qualcomm/executor_runner/qnn_executor_runner.cpp
+++ b/examples/qualcomm/executor_runner/qnn_executor_runner.cpp
@@ -143,6 +143,13 @@ DEFINE_int32(
     "This is a runtime option and will override the profile level set during AOT. "
     "Refer to QnnExecuTorchProfileLevel under qc_compiler_spec.fbs for more info.");
 
+DEFINE_string(
+    heap_profiling_path,
+    "",
+    "Output path for QNN heap-profiling dump. "
+    "Empty disables heap profiling. "
+    "This is a runtime option and will override the path set during AOT.");
+
 using executorch::aten::Tensor;
 using executorch::aten::TensorImpl;
 using executorch::etdump::ETDumpGen;
@@ -213,7 +220,7 @@ int main(int argc, char** argv) {
   }
 
   // Set runtime options
-  executorch::runtime::BackendOptions<8> backend_options;
+  executorch::runtime::BackendOptions<9> backend_options;
   if (!gflags::GetCommandLineFlagInfoOrDie("log_level").is_default) {
     ET_LOG(Info, "Setting runtime log level: %d", FLAGS_log_level);
     ET_CHECK_MSG(
@@ -291,6 +298,18 @@ int main(int argc, char** argv) {
         "Failed to set backend options: %s",
         QNN_RUNTIME_LPAI_CORE_SELECTION);
   }
+  if (!gflags::GetCommandLineFlagInfoOrDie("heap_profiling_path").is_default) {
+    ET_LOG(
+        Info,
+        "Setting runtime heap_profiling_path: %s",
+        FLAGS_heap_profiling_path.c_str());
+    ET_CHECK_MSG(
+        backend_options.set_option(
+            QNN_RUNTIME_HEAP_PROFILING_PATH,
+            FLAGS_heap_profiling_path.c_str()) == Error::Ok,
+        "Failed to set backend options: %s",
+        QNN_RUNTIME_HEAP_PROFILING_PATH);
+  }
   ET_CHECK_MSG(
       set_option(QNN_BACKEND, backend_options.view()) == Error::Ok,
       "Failed to set runtime options.");
diff --git a/examples/qualcomm/oss_scripts/llama/qnn_multimodal_runner.cpp b/examples/qualcomm/oss_scripts/llama/qnn_multimodal_runner.cpp
index 1a7da5d62d3..29b6b9d7ddc 100644
--- a/examples/qualcomm/oss_scripts/llama/qnn_multimodal_runner.cpp
+++ b/examples/qualcomm/oss_scripts/llama/qnn_multimodal_runner.cpp
@@ -177,14 +177,15 @@ void start_multimodal_runner(
     }
   };
   executorch::extension::llm::GenerationConfig config{
-      true,
-      false,
-      -1,
-      false,
-      FLAGS_seq_len,
-      static_cast(FLAGS_temperature),
-      0,
-      0};
+      .echo = true,
+      .ignore_eos = false,
+      .max_new_tokens = -1,
+      .warming = false,
+      .seq_len = FLAGS_seq_len,
+      .temperature = static_cast(FLAGS_temperature),
+      .num_bos = 0,
+      .num_eos = 0,
+  };
 
   // 1. [Multimodal] Get raw files from input_list.txt
   std::vector audio_raw_files;
diff --git a/examples/qualcomm/util_scripts/qnn_intermediate_debugger_demo.py b/examples/qualcomm/util_scripts/qnn_intermediate_debugger_demo.py
index 07dfc9c9558..e7c7c3985a8 100644
--- a/examples/qualcomm/util_scripts/qnn_intermediate_debugger_demo.py
+++ b/examples/qualcomm/util_scripts/qnn_intermediate_debugger_demo.py
@@ -7,15 +7,14 @@
 import json
 import os
 from multiprocessing.connection import Client
-from typing import Any, Tuple
 
 import numpy as np
 
 import torch
-import torch.nn.functional as F
-from executorch.backends.qualcomm.debugger.metrics_evaluator import (
-    CosineSimilarityEvaluator,
-    MetricEvaluatorBase,
+
+from executorch.backends.qualcomm.debugger.qcom_numerical_comparator_sample import (
+    QcomCosineSimilarityComparator,
+    QcomMSEComparator,
 )
 from executorch.backends.qualcomm.debugger.qnn_intermediate_debugger import (
     OutputFormat,
@@ -29,7 +28,6 @@
     SimpleADB,
 )
 from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype
-from executorch.devtools import Inspector
 from executorch.examples.models.inception_v3.model import InceptionV3Model
 from executorch.examples.qualcomm.utils import (
     get_imagenet_dataset,
@@ -49,12 +47,6 @@ def main(args):
     # ensure the working directory exist.
     os.makedirs(args.artifact, exist_ok=True)
 
-    if not args.compile_only and args.device is None:
-        raise RuntimeError(
-            "device serial is required if not compile only. "
-            "Please specify a device serial by -s/--device argument."
-        )
-
     data_num = 100
 
     inputs, targets = get_imagenet_dataset(
@@ -62,12 +54,13 @@ def main(args):
         data_size=data_num,
         image_shape=(256, 256),
         crop_size=224,
+        shuffle=False,
     )
     pte_filename = "ic3_qnn_debug"
     instance = InceptionV3Model()
     source_model = instance.get_eager_model().eval()
     # Init our QNNIntermediateDebugger and pass it in to build_executorch_binary().
-    qnn_intermediate_debugger = QNNIntermediateDebugger()
+    qnn_intermediate_debugger = QNNIntermediateDebugger(sample_input=inputs[0])
     build_executorch_binary(
         model=source_model,
         qnn_config=qnn_config,
@@ -82,9 +75,6 @@ def main(args):
     inputs = [inputs[0]]
     targets = [targets[0]]
 
-    if args.compile_only:
-        return
-
     # Please ensure that dump_intermediate_outputs are set to true when creating SimpleADB
     adb = SimpleADB(
         qnn_config=qnn_config,
@@ -98,64 +88,52 @@ def main(args):
     output_data_folder = f"{args.artifact}/outputs"
     make_output_dir(output_data_folder)
 
-    class RootMeanSquaredErrorEvaluator(MetricEvaluatorBase):
-        def __init__(self, threshold=0.02):
-            self.threshold = threshold
-
-        def metric_name(self) -> str:
-            return "Root Mean Squared Error"
-
-        def evaluate(
-            self, qnn_output: torch.Tensor, cpu_output: torch.Tensor
-        ) -> Tuple[Any, bool]:
-            mse = F.mse_loss(qnn_output, cpu_output)
-            rmse = torch.sqrt(mse)
-            valid = rmse < self.threshold
-            return rmse, valid
-
     # We will pull the debug output and provide them to the Inspector class.
     # We can then provide our own metrics and output type to generate the intermediate debugging results.
     def validate_intermediate_tensor():
-        inspector = Inspector(
+        qnn_intermediate_debugger.setup_inspector(
             etdump_path=f"{args.artifact}/etdump.etdp",
             debug_buffer_path=f"{args.artifact}/debug_output.bin",
         )
 
-        edge_result = qnn_intermediate_debugger.intermediate_output_module(
-            *(inputs[0])
-        )[0]
+        edge_result = qnn_intermediate_debugger.edge_ep.module()(
+            *(qnn_intermediate_debugger.sample_input)
+        )
 
-        # Optional: Ensures that edge module accuracy aligns with nn.Module
+        # Highly Recommended: Ensures that edge module accuracy aligns with nn.Module
         with torch.no_grad():
-            source_result = source_model(*(inputs[0]))
+            source_result = source_model(*(qnn_intermediate_debugger.sample_input))
             score = torch.nn.functional.cosine_similarity(
                 edge_result.flatten(), source_result.flatten(), dim=0
             ).item()
             print("Cosine Similarity Score between nn.Module and Edge CPU is: ", score)
-
         # Users can generate multiple comparison metrics in a single execution.
-        # Below, we generate 3 metrics.
+
+        cos_comparator = qnn_intermediate_debugger.create_comparator(
+            QcomCosineSimilarityComparator, threshold=0.9
+        )
         qnn_intermediate_debugger.generate_results(
             title="ic3_cos_similarity_debugging_graph",
             path=args.artifact,
-            output_format=OutputFormat.SVG_GRAPHS,
-            inspector=inspector,
-            evaluator=CosineSimilarityEvaluator(0.9),
+            output_format=OutputFormat.SVG_GRAPH,
+            comparator=cos_comparator,
         )
+
         qnn_intermediate_debugger.generate_results(
-            title="ic3_cos_similarity_csv",
+            title="ic3_cos_similarity_debugging_graph",
             path=args.artifact,
-            output_format=OutputFormat.CSV_FILES,
-            inspector=inspector,
-            evaluator=CosineSimilarityEvaluator(0.9),
+            output_format=OutputFormat.CSV_FILE,
+            comparator=cos_comparator,
+        )
+
+        mse_comparator = qnn_intermediate_debugger.create_comparator(
+            QcomMSEComparator, threshold=0.1
         )
-        # Using self defined metrics to print svg graphs
         qnn_intermediate_debugger.generate_results(
-            title="ic3_rmse_debugging_graph",
+            title="ic3_mse_debugging_graph",
             path=args.artifact,
-            output_format=OutputFormat.SVG_GRAPHS,
-            inspector=inspector,
-            evaluator=RootMeanSquaredErrorEvaluator(0.9),
+            output_format=OutputFormat.SVG_GRAPH,
+            comparator=mse_comparator,
         )
 
     adb.pull_debug_output(
@@ -180,8 +158,8 @@ def validate_intermediate_tensor():
             conn.send(
                 json.dumps(
                     {
-                        "svg_path": f"{args.artifact}/ic3_rmse_debugging_graph.svg",
-                        "csv_path": f"{args.artifact}/ic3_cos_similarity_csv.csv",
+                        "svg_path": f"{args.artifact}/ic3_mse_debugging_graph.svg",
+                        "csv_path": f"{args.artifact}/ic3_cos_similarity_debugging_graph.csv",
                     }
                 )
             )
diff --git a/examples/raspberry_pi/setup.sh b/examples/raspberry_pi/setup.sh
index 894a8365e6e..3bf34788d63 100755
--- a/examples/raspberry_pi/setup.sh
+++ b/examples/raspberry_pi/setup.sh
@@ -245,7 +245,7 @@ setup_environment() {
         # Check if conda is available
         if command -v conda &> /dev/null; then
             log_info "Creating conda environment..."
-            conda create -yn executorch python=3.10.0
+            conda create -yn executorch python=3.10
             eval "$(conda shell.bash hook)"
             conda activate executorch
             log_success "Created and activated conda environment: executorch"
diff --git a/examples/riscv/aot_riscv.py b/examples/riscv/aot_riscv.py
index 8076f056ba2..529e2b1e767 100644
--- a/examples/riscv/aot_riscv.py
+++ b/examples/riscv/aot_riscv.py
@@ -3,14 +3,15 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-"""AOT export for the RISC-V Phase 1.0 smoke test.
+"""AOT export for the RISC-V smoke test.
 
-Exports a trivial ``torch.add`` module to a BundledProgram (.bpte) that the
-portable executor_runner can load on a riscv64 target and verify against the
-embedded reference output, emitting ``Test_result: PASS`` on success.
+Exports a small model to a BundledProgram (.bpte) that the portable
+executor_runner can load on a riscv64 target and verify against the embedded
+reference output, emitting ``Test_result: PASS`` on success.
 """
 
 import argparse
+import logging
 from pathlib import Path
 
 import torch
@@ -28,26 +29,186 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         return x + y
 
 
+def build_add():
+    model = AddModule().eval()
+    example_inputs = (torch.ones(1, 4), torch.full((1, 4), 2.0))
+    test_inputs = [
+        (torch.ones(1, 4), torch.full((1, 4), 2.0)),
+        (torch.full((1, 4), 3.0), torch.full((1, 4), 4.0)),
+    ]
+    return model, example_inputs, test_inputs, True
+
+
+def build_mv2():
+    from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
+
+    model = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).eval()
+    torch.manual_seed(0)
+    example_inputs = (torch.randn(1, 3, 224, 224),)
+    test_inputs = [example_inputs]
+    return model, example_inputs, test_inputs, False
+
+
+def build_mobilebert():
+    from transformers import MobileBertConfig, MobileBertModel
+
+    config = MobileBertConfig(
+        vocab_size=1024,
+        hidden_size=128,
+        embedding_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        intermediate_size=128,
+        intra_bottleneck_size=32,
+    )
+
+    class Wrapper(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.model = MobileBertModel(config).eval()
+
+        def forward(self, input_ids):
+            return self.model(input_ids).last_hidden_state
+
+    model = Wrapper().eval()
+    example_inputs = (torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]),)
+    test_inputs = [example_inputs]
+    return model, example_inputs, test_inputs, False
+
+
+def build_llama2():
+    # Use the executorch native Transformer (matches MODEL_NAME_TO_MODEL["llama2"]
+    # in examples/models/__init__.py). Unlike HF LlamaModel, RoPE freqs are
+    # precomputed buffers and just sliced at forward time, so no
+    # torch.arange()/Long causal mask is built per forward — which is what
+    # the PT2E XNNPACK quantizer trips over on HF Llama.
+    from executorch.examples.models.llama.llama_transformer import construct_transformer
+    from executorch.examples.models.llama.model_args import ModelArgs
+
+    seq_len = 8
+    args = ModelArgs(
+        dim=128,
+        n_layers=2,
+        n_heads=4,
+        n_kv_heads=2,  # GQA: kv_heads < n_heads exercises the GQA path
+        vocab_size=1024,
+        hidden_dim=256,  # SwiGLU FFN: gate + up projections at this width
+        max_seq_len=seq_len,
+        max_context_len=seq_len,
+        rope_theta=10000.0,
+    )
+    torch.manual_seed(0)
+    model = construct_transformer(args).eval()
+    example_inputs = (torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=torch.long),)
+    test_inputs = [example_inputs]
+    return model, example_inputs, test_inputs, False
+
+
+def build_resnet18():
+    from torchvision.models import resnet18, ResNet18_Weights
+
+    model = resnet18(weights=ResNet18_Weights.DEFAULT).eval()
+    torch.manual_seed(0)
+    example_inputs = (torch.randn(1, 3, 224, 224),)
+    test_inputs = [example_inputs]
+    return model, example_inputs, test_inputs, False
+
+
+MODELS = {
+    "add": build_add,
+    "mv2": build_mv2,
+    "mobilebert": build_mobilebert,
+    "llama2": build_llama2,
+    "resnet18": build_resnet18,
+}
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--model",
+        choices=sorted(MODELS),
+        default="add",
+        help="Which model to export",
+    )
     parser.add_argument(
         "--output",
         type=Path,
-        default=Path("add_riscv.bpte"),
-        help="Output .bpte path",
+        default=None,
+        help="Output .bpte path (default: _riscv.bpte)",
+    )
+    parser.add_argument(
+        "--xnnpack",
+        action="store_true",
+        help="Lower through the XNNPACK partitioner",
+    )
+    parser.add_argument(
+        "--quantize",
+        action="store_true",
+        help="Produce an 8-bit quantized model",
+    )
+    parser.add_argument(
+        "--debug-xnnpack",
+        action="store_true",
+        help="Enable XNNPACK partitioner DEBUG logging and dump the lowered graph",
     )
     args = parser.parse_args()
 
-    model = AddModule().eval()
-    example_inputs = (torch.ones(1, 4), torch.full((1, 4), 2.0))
+    if args.debug_xnnpack:
+        logging.basicConfig(level=logging.DEBUG)
 
-    exported = export(model, example_inputs)
-    et_program = to_edge_transform_and_lower(exported).to_executorch()
+    if args.output is None:
+        args.output = Path(f"{args.model}_riscv.bpte")
+
+    model, example_inputs, test_inputs, strict = MODELS[args.model]()
+
+    if args.quantize:
+        from executorch.examples.xnnpack import MODEL_NAME_TO_OPTIONS, QuantType
+        from executorch.examples.xnnpack.quantization.utils import quantize
+
+        if args.model not in MODEL_NAME_TO_OPTIONS:
+            parser.error(f"No XNNPACK quantization recipe for model {args.model!r}")
+        quant_type = MODEL_NAME_TO_OPTIONS[args.model].quantization
+        if quant_type == QuantType.NONE:
+            parser.error(f"Quantization recipe for {args.model!r} is NONE")
+        ep = export(model, example_inputs, strict=strict)
+        model = quantize(ep.module(), example_inputs, quant_type)
+
+    exported = export(model, example_inputs, strict=strict)
+    partitioners = []
+    if args.xnnpack:
+        from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
+            XnnpackPartitioner,
+        )
+
+        partitioners.append(XnnpackPartitioner(verbose=args.debug_xnnpack))
+
+    compile_config = None
+    if args.quantize:
+        from executorch.exir import EdgeCompileConfig
+
+        compile_config = EdgeCompileConfig(_check_ir_validity=False)
+
+    edge = to_edge_transform_and_lower(
+        exported, partitioner=partitioners, compile_config=compile_config
+    )
+    delegated = sum(
+        1
+        for n in edge.exported_program().graph.nodes
+        if n.op == "call_function" and "call_delegate" in str(n.target)
+    )
+    print(
+        f"[aot_riscv] model={args.model} xnnpack={args.xnnpack} "
+        f"quantize={args.quantize} delegated_nodes={delegated}"
+    )
+
+    if args.debug_xnnpack:
+        from executorch.exir.backend.utils import print_delegated_graph
+
+        print_delegated_graph(edge.exported_program().graph_module)
+
+    et_program = edge.to_executorch()
 
-    test_inputs = [
-        (torch.ones(1, 4), torch.full((1, 4), 2.0)),
-        (torch.full((1, 4), 3.0), torch.full((1, 4), 4.0)),
-    ]
     test_suite = MethodTestSuite(
         method_name="forward",
         test_cases=[
diff --git a/examples/riscv/etdump_summary.py b/examples/riscv/etdump_summary.py
new file mode 100644
index 00000000000..e4fc5a61d7e
--- /dev/null
+++ b/examples/riscv/etdump_summary.py
@@ -0,0 +1,228 @@
+# Copyright 2026 The ExecuTorch Authors.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Render a per-XNNPACK-op summary from an ETDump file."""
+
+import argparse
+import json
+import re
+import sys
+from collections import defaultdict
+from pathlib import Path
+
+from executorch.devtools import Inspector
+
+
+# "Convolution (NHWC, F32) IGEMM #3" -> ("Convolution (NHWC, F32) IGEMM", 3)
+_SEQ_RE = re.compile(r"^(.*?)\s+#(\d+)$")
+
+# Wrappers around per-op events; kept separate to avoid double-counting children.
+FRAMEWORK_EVENTS = frozenset(
+    {
+        "Method::execute",
+        "Method::init",
+        "Program::load_method",
+        "DELEGATE_CALL",
+        "OPERATOR_CALL",
+    }
+)
+
+_REG_LOG_RE = re.compile(r"Note \(XNNPACK\):.*microkernel '([^']+)'")
+
+
+def parse_run_log(path: Path):
+    syms = set()
+    with open(path, errors="ignore") as f:
+        for line in f:
+            m = _REG_LOG_RE.search(line)
+            if m:
+                syms.add(m.group(1))
+    return sorted(syms)
+
+
+# Two-source mapping from an ETDump op name to a symbol-substring pattern.
+# When the operator type uses xnn_microkernel_type_default, runtime.c does NOT
+# append a category suffix, so we fall back to matching on the base op name.
+_OP_NAME_RE = re.compile(r"^(.*?)\s*\(([^)]*)\)\s*(.*)$")
+_DTYPE_TOKENS = frozenset(
+    {
+        "F32",
+        "F16",
+        "QS8",
+        "QU8",
+        "QC8",
+        "QC4",
+        "QD8",
+        "QC8W",
+        "QC4W",
+        "X8",
+        "X16",
+        "X24",
+        "X32",
+        "X64",
+    }
+)
+# Infix between the kind token and `_ukernel_`: zero or more `_`
+# segments (e.g. `_gemm_ukernel_`, `_gemm_minmax_ukernel_`,
+# `_gemm_minmax_fp32_ukernel_`, ...).
+_INFIX = r"(?:[a-z0-9]+_)*"
+_KIND_PATTERN = {
+    # Microkernel categories appended by runtime.c (xnn_microkernel_type_to_string).
+    "GEMM": r"_gemm_" + _INFIX + r"ukernel_",
+    "IGEMM": r"_igemm_" + _INFIX + r"ukernel_",
+    "DWConv": r"_dwconv_" + _INFIX + r"ukernel_",
+    "Transpose": r"_transposec?_" + _INFIX + r"ukernel_",
+    "Reduce": r"_(?:rsum|rmax|rminmax|rdmax|rdsum)_" + _INFIX + r"ukernel_",
+    "Reduce2": r"_(?:rdmax|rdsum)_" + _INFIX + r"ukernel_",
+    "VMulCAddC": r"_vmulcaddc_" + _INFIX + r"ukernel_",
+    "Average Pooling": r"_(?:avgpool|gavgpool)_" + _INFIX + r"ukernel_",
+    "Pixelwise Average Pooling": r"_pavgpool_" + _INFIX + r"ukernel_",
+    "Conv2D HWC2CHW": r"_conv_hwc2chw_" + _INFIX + r"ukernel_",
+    "SPMM": r"_spmm_" + _INFIX + r"ukernel_",
+    "Subconv2D": r"_subconv2d_" + _INFIX + r"ukernel_",
+    # Base op names (default microkernel type, no category suffix in the ETDump name).
+    "Add": r"_v(?:add|addc)_" + _INFIX + r"ukernel_",
+    "Subtract": r"_v(?:sub|subc|rsubc)_" + _INFIX + r"ukernel_",
+    "Multiply": r"_v(?:mul|mulc)_" + _INFIX + r"ukernel_",
+    "Divide": r"_v(?:div|divc|rdivc)_" + _INFIX + r"ukernel_",
+    "Maximum": r"_v(?:max|maxc)_" + _INFIX + r"ukernel_",
+    "Minimum": r"_v(?:min|minc)_" + _INFIX + r"ukernel_",
+    "Clamp": r"_vclamp_" + _INFIX + r"ukernel_",
+    "Sigmoid": r"_vsigmoid_" + _INFIX + r"ukernel_",
+    "Tanh": r"_vtanh_" + _INFIX + r"ukernel_",
+    "Negate": r"_vneg_" + _INFIX + r"ukernel_",
+    "Abs": r"_vabs_" + _INFIX + r"ukernel_",
+    "Square": r"_vsqr_" + _INFIX + r"ukernel_",
+    "Square Root": r"_vsqrt_" + _INFIX + r"ukernel_",
+    "Reciprocal Square Root": r"_vrsqrt_" + _INFIX + r"ukernel_",
+    "Convert": r"_vcvt_" + _INFIX + r"ukernel_",
+    "Copy": r"_(?:copy|memcpy)_" + _INFIX + r"ukernel_",
+    "Constant Pad": r"_xx_pad_" + _INFIX + r"ukernel_",
+    "Softmax": r"_(?:raddstoreexpminusmax|rmax)_" + _INFIX + r"ukernel_",
+    "Max Pooling": r"_maxpool_" + _INFIX + r"ukernel_",
+}
+
+
+def op_kernels(op_name, kernels):
+    m = _OP_NAME_RE.match(op_name)
+    if not m:
+        return []
+    base, inside, tail = m.group(1).strip(), m.group(2), m.group(3).strip()
+    key = tail if tail in _KIND_PATTERN else (base if base in _KIND_PATTERN else None)
+    if key is None:
+        return []
+    dtype_tokens = [
+        s.strip().lower() for s in inside.split(",") if s.strip() in _DTYPE_TOKENS
+    ]
+    cat_re = re.compile(_KIND_PATTERN[key])
+    return [
+        sym
+        for sym in kernels
+        if cat_re.search(sym) and all(d in sym for d in dtype_tokens)
+    ]
+
+
+def aggregate(etdump_path: Path):
+    insp = Inspector(etdump_path=str(etdump_path))
+    per_op = defaultdict(lambda: {"count": 0, "raw": []})
+    framework = defaultdict(lambda: {"count": 0, "raw": []})
+    for block in insp.event_blocks:
+        for ev in block.events:
+            m = _SEQ_RE.match(ev.name or "")
+            base = m.group(1) if m else (ev.name or "")
+            bucket = framework if base in FRAMEWORK_EVENTS else per_op
+            bucket[base]["count"] += 1
+            bucket[base]["raw"].extend(ev.perf_data.raw if ev.perf_data else [])
+    return per_op, framework
+
+
+def render(per_op, framework, etdump_path, kernels):
+    def rows_of(d):
+        rows = []
+        for name, v in d.items():
+            raw = v["raw"]
+            s = sum(raw)
+            rows.append(
+                {
+                    "op": name,
+                    "count": v["count"],
+                    "sum_ms": s,
+                    "avg_ms": (s / len(raw)) if raw else 0.0,
+                    "max_ms": max(raw) if raw else 0.0,
+                    "kernels": op_kernels(name, kernels) if kernels else [],
+                }
+            )
+        rows.sort(key=lambda r: r["sum_ms"], reverse=True)
+        return rows
+
+    op_rows = rows_of(per_op)
+    fw_rows = rows_of(framework)
+    ops_total = sum(r["sum_ms"] for r in op_rows)
+    fw_total = sum(r["sum_ms"] for r in fw_rows)
+
+    def fmt_table(label, rows, total):
+        print(f"\n[etdump_summary] {label}  total={total:.3f} ms")
+        print(
+            f"{'%':>5}  {'sum_ms':>10}  {'count':>6}  {'avg_ms':>10}  {'max_ms':>10}  op"
+        )
+        for r in rows:
+            pct = (r["sum_ms"] / total * 100.0) if total else 0.0
+            print(
+                f"{pct:5.1f}  {r['sum_ms']:10.3f}  {r['count']:6d}  "
+                f"{r['avg_ms']:10.3f}  {r['max_ms']:10.3f}  {r['op']}"
+            )
+
+    print(f"[etdump_summary] {etdump_path}")
+    fmt_table(f"XNNPACK ops ({len(op_rows)} unique)", op_rows, ops_total)
+    fmt_table(f"Framework wrappers ({len(fw_rows)})", fw_rows, fw_total)
+    if kernels:
+        print(f"\n[etdump_summary] Registered XNNPACK microkernels ({len(kernels)}):")
+        for sym in kernels:
+            print(f"  {sym}")
+
+    return op_rows, fw_rows, ops_total
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("etdump", type=Path)
+    parser.add_argument("--run-log", type=Path, default=None)
+    parser.add_argument("--json", type=Path, default=None)
+    args = parser.parse_args()
+
+    if not args.etdump.exists():
+        print(f"[etdump_summary] missing {args.etdump}", file=sys.stderr)
+        sys.exit(1)
+
+    kernels = []
+    if args.run_log is not None:
+        if not args.run_log.exists():
+            print(f"[etdump_summary] missing run log {args.run_log}", file=sys.stderr)
+            sys.exit(1)
+        kernels = parse_run_log(args.run_log)
+
+    per_op, framework = aggregate(args.etdump)
+    op_rows, fw_rows, ops_total = render(per_op, framework, args.etdump, kernels)
+
+    if args.json is not None:
+        args.json.parent.mkdir(parents=True, exist_ok=True)
+        args.json.write_text(
+            json.dumps(
+                {
+                    "etdump": str(args.etdump),
+                    "run_log": str(args.run_log) if args.run_log else None,
+                    "ops_total_ms": ops_total,
+                    "registered_kernels": kernels,
+                    "ops": op_rows,
+                    "framework": fw_rows,
+                },
+                indent=2,
+            )
+        )
+        print(f"[etdump_summary] wrote {args.json}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/riscv/requirements.txt b/examples/riscv/requirements.txt
new file mode 100644
index 00000000000..273e7156a1d
--- /dev/null
+++ b/examples/riscv/requirements.txt
@@ -0,0 +1,2 @@
+torchvision
+transformers
diff --git a/examples/riscv/run.sh b/examples/riscv/run.sh
index 7c05edcbc8c..2c207816bfc 100755
--- a/examples/riscv/run.sh
+++ b/examples/riscv/run.sh
@@ -20,11 +20,21 @@ build_dir="${et_root_dir}/cmake-out-riscv"
 output_dir="${et_root_dir}/riscv_test"
 qemu="qemu-riscv64-static"
 qemu_timeout="600"
+model="add"
+xnnpack=false
+quantize=false
+debug_xnnpack=false
+verbose_xnnpack=false
 
 usage() {
     cat <          Which model to export and run (default: ${model})
+  --xnnpack               Enable the XNNPACK backend (AOT partitioner + runtime)
+  --quantize              Produce an 8-bit quantized model
+  --verbose-xnnpack       Build XNNPACK with XNN_LOG_LEVEL=4 to log microkernel dispatch at runtime
+  --debug-xnnpack         Enable XNNPACK partitioner DEBUG logging and dump the lowered graph
   --build_only            Only export and cross-compile; do not invoke QEMU
   --build_dir=       CMake build directory (default: ${build_dir})
   --output_dir=      Directory for the exported .bpte (default: ${output_dir})
@@ -36,6 +46,11 @@ EOF
 
 for arg in "$@"; do
     case $arg in
+        --model=*) model="${arg#*=}" ;;
+        --xnnpack) xnnpack=true ;;
+        --quantize) quantize=true ;;
+        --debug-xnnpack) debug_xnnpack=true ;;
+        --verbose-xnnpack) verbose_xnnpack=true ;;
         --build_only) build_only=true ;;
         --build_dir=*) build_dir="${arg#*=}" ;;
         --output_dir=*) output_dir="${arg#*=}" ;;
@@ -47,14 +62,32 @@ for arg in "$@"; do
 done
 
 mkdir -p "${output_dir}"
-bpte_path="${output_dir}/add_riscv.bpte"
+bpte_path="${output_dir}/${model}_riscv.bpte"
 
 echo "[run.sh] Step 1/3: AOT export on host"
-python "${script_dir}/aot_riscv.py" --output "${bpte_path}"
+aot_extra_args=()
+if ${xnnpack}; then
+    aot_extra_args+=(--xnnpack)
+fi
+if ${quantize}; then
+    aot_extra_args+=(--quantize)
+fi
+if ${debug_xnnpack}; then
+    aot_extra_args+=(--debug-xnnpack)
+fi
+python "${script_dir}/aot_riscv.py" --model "${model}" "${aot_extra_args[@]}" --output "${bpte_path}"
 
 echo "[run.sh] Step 2/3: cross-compile executor_runner for riscv64-linux"
+cmake_extra_args=()
+if ${xnnpack}; then
+    cmake_extra_args+=(-DEXECUTORCH_BUILD_XNNPACK=ON)
+fi
+if ${verbose_xnnpack}; then
+    cmake_extra_args+=(-DEXECUTORCH_XNNPACK_LOG_LEVEL=4 -DEXECUTORCH_BUILD_RISCV_ETDUMP=ON)
+fi
 cmake -S "${et_root_dir}" -B "${build_dir}" \
     --preset riscv64-linux \
+    "${cmake_extra_args[@]}" \
     -DCMAKE_BUILD_TYPE=Release
 cmake --build "${build_dir}" -j"$(nproc)" --target executor_runner
 
@@ -84,18 +117,41 @@ hash "${qemu}" 2>/dev/null || {
 # linker (ld-linux-riscv64-lp64d.so.1) referenced in the ELF resolves.
 export QEMU_LD_PREFIX="${QEMU_LD_PREFIX:-/usr/riscv64-linux-gnu}"
 
-log_file=$(mktemp)
-trap 'rm -f "${log_file}"' EXIT
+if [[ -n "${QEMU_CPU+x}" ]]; then
+    echo "[run.sh] QEMU_CPU=${QEMU_CPU}"
+fi
+
+runner_extra_args=()
+if ${quantize}; then
+    runner_extra_args+=(--bundleio_rtol=0.1 --bundleio_atol=0.25)
+fi
+etdump_path=""
+if ${verbose_xnnpack}; then
+    etdump_path="${output_dir}/${model}_riscv.etdump"
+    rm -f "${etdump_path}"
+    runner_extra_args+=(--etdump_path="${etdump_path}")
+fi
+
+# etdump_summary.py reads the XNN_LOG_LEVEL=4 registrations.
+log_file="${output_dir}/${model}_riscv.run.log"
+rm -f "${log_file}"
 
 set +e
 timeout --signal=KILL "${qemu_timeout}" "${qemu}" "${runner}" \
     --model_path="${bpte_path}" \
+    "${runner_extra_args[@]}" \
     2>&1 | tee "${log_file}"
 qemu_status=${PIPESTATUS[0]}
 set -e
 
 echo "[run.sh] qemu exit status: ${qemu_status}"
 
+if [[ -n "${etdump_path}" && -f "${etdump_path}" ]]; then
+    python "${script_dir}/etdump_summary.py" "${etdump_path}" \
+        --run-log "${log_file}" \
+        --json "${etdump_path}.json" || true
+fi
+
 if grep -q "Test_result: PASS" "${log_file}"; then
     echo "[run.sh] Bundled I/O check PASSED"
     exit 0
diff --git a/examples/riscv/setup.sh b/examples/riscv/setup.sh
index c1342c60d5e..955c8ca3386 100755
--- a/examples/riscv/setup.sh
+++ b/examples/riscv/setup.sh
@@ -10,6 +10,8 @@
 
 set -eu
 
+script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd)
+
 if ! command -v apt-get >/dev/null 2>&1; then
     echo "[$(basename "$0")] this setup script targets Debian/Ubuntu (apt-get not found)" >&2
     exit 1
@@ -23,14 +25,23 @@ fi
 ${SUDO} apt-get update
 ${SUDO} apt-get install -y --no-install-recommends \
     build-essential \
-    gcc-riscv64-linux-gnu \
-    g++-riscv64-linux-gnu \
+    gcc${GCC_VERSION:+-${GCC_VERSION}}-riscv64-linux-gnu \
+    g++${GCC_VERSION:+-${GCC_VERSION}}-riscv64-linux-gnu \
     binutils-riscv64-linux-gnu \
     libc6-riscv64-cross \
     libc6-dev-riscv64-cross \
     cmake \
     file \
+    ca-certificates \
     qemu-user-static
 
+if [[ -n "${GCC_VERSION+x}" ]]; then
+    ${SUDO} update-alternatives --install /usr/bin/riscv64-linux-gnu-gcc riscv64-linux-gnu-gcc /usr/bin/riscv64-linux-gnu-gcc${GCC_VERSION:+-${GCC_VERSION}} 100
+    ${SUDO} update-alternatives --install /usr/bin/riscv64-linux-gnu-g++ riscv64-linux-gnu-g++ /usr/bin/riscv64-linux-gnu-g++${GCC_VERSION:+-${GCC_VERSION}} 100
+fi
+
 riscv64-linux-gnu-gcc --version | head -n1
 qemu-riscv64-static --version | head -n1
+
+# Some python packages also need to be installed
+pip install -r "${script_dir}/requirements.txt"
diff --git a/exir/print_program.py b/exir/print_program.py
index c1ec1a0bb8e..18029f4169a 100644
--- a/exir/print_program.py
+++ b/exir/print_program.py
@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -60,6 +61,8 @@ def _scalar_type_str(scalar_type: ScalarType) -> str:
         ScalarType.QUINT8: "qui8",
         ScalarType.QINT32: "qi32",
         ScalarType.BFLOAT16: "bf16",
+        ScalarType.FLOAT8E5M2: "f8e5m2",
+        ScalarType.FLOAT8E4M3FN: "f8e4m3fn",
         ScalarType.QUINT4x2: "qui4x2",
         ScalarType.QUINT2x4: "qui2x4",
     }
diff --git a/exir/serde/export_serialize.py b/exir/serde/export_serialize.py
index 7fd1f9470d4..572d87f2dec 100644
--- a/exir/serde/export_serialize.py
+++ b/exir/serde/export_serialize.py
@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -149,6 +150,8 @@ def _reverse_map(d: Dict[Any, Enum]):
     torch.complex128: ScalarType.COMPLEXDOUBLE,
     torch.bool: ScalarType.BOOL,
     torch.bfloat16: ScalarType.BFLOAT16,
+    torch.float8_e5m2: ScalarType.FLOAT8E5M2,
+    torch.float8_e4m3fn: ScalarType.FLOAT8E4M3FN,
     torch.uint16: ScalarType.UINT16
 }
 
diff --git a/exir/serde/schema.py b/exir/serde/schema.py
index f91526c385f..bd1904aa34f 100644
--- a/exir/serde/schema.py
+++ b/exir/serde/schema.py
@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -16,7 +17,7 @@
 from executorch.exir.serde.union import _Union
 
 # NOTE: Please update this value if any modifications are made to the schema
-SCHEMA_VERSION = (5, 3)
+SCHEMA_VERSION = (5, 4)
 TREESPEC_VERSION = 1
 
 
@@ -36,6 +37,8 @@ class ScalarType(IntEnum):
     BOOL = 12
     BFLOAT16 = 13
     UINT16 = 14
+    FLOAT8E5M2 = 15
+    FLOAT8E4M3FN = 16
 
 class Layout(IntEnum):
     Unknown = 0
diff --git a/exir/tensor.py b/exir/tensor.py
index b1619d16bdf..e40d6f10168 100644
--- a/exir/tensor.py
+++ b/exir/tensor.py
@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -292,6 +293,8 @@ def memory_format_enum(memory_format: torch.memory_format) -> int:
     torch.qint32: ScalarType.QINT32,
     torch.bfloat16: ScalarType.BFLOAT16,
     torch.quint4x2: ScalarType.QUINT4x2,
+    torch.float8_e5m2: ScalarType.FLOAT8E5M2,
+    torch.float8_e4m3fn: ScalarType.FLOAT8E4M3FN,
     torch.uint16: ScalarType.UINT16,
     torch.uint32: ScalarType.UINT32,
 }
diff --git a/exir/tests/test_tensor.py b/exir/tests/test_tensor.py
index c5383b0dac2..1a73e81319c 100644
--- a/exir/tests/test_tensor.py
+++ b/exir/tests/test_tensor.py
@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -90,6 +91,14 @@ def test_normal_tensor_conversion(self) -> None:
         # whereas strides for torch.memory_format = torch.channels_last is
         # (3*4*5, 1, 5*3, 3))
 
+    def test_fp8_tensor_conversion(self) -> None:
+        for dtype in (torch.float8_e5m2, torch.float8_e4m3fn):
+            normal_tensor = torch.randn(2, 2, 3, dtype=torch.float32).to(dtype)
+            flatbuffer_tensor = make_tensor_value(
+                1, 0, TensorSpec.from_tensor(normal_tensor)
+            )
+            self.compare_tensors(normal_tensor, flatbuffer_tensor)
+
     def test_allocation_info_succeeds(self) -> None:
         test_cases = (
             (
diff --git a/extension/android/executorch_android/build.gradle b/extension/android/executorch_android/build.gradle
index f84aafe138c..3ee5b5877b3 100644
--- a/extension/android/executorch_android/build.gradle
+++ b/extension/android/executorch_android/build.gradle
@@ -114,3 +114,17 @@ repositories {
         url "https://central.sonatype.com/repository/maven-snapshots/"
     }
 }
+
+android.libraryVariants.all { variant ->
+    task("generate${variant.name.substring(0, 1).toUpperCase()}${variant.name.substring(1)}Javadoc", type: Javadoc) {
+        source = variant.javaCompileProvider.get().source
+        classpath += project.files(android.getBootClasspath().join(File.pathSeparator))
+        classpath += variant.javaCompileProvider.get().classpath
+        options {
+            overview = "src/main/javadoc/overview.html"
+            windowTitle = "ExecuTorch Android Java API"
+            docTitle = "ExecuTorch Android Java API"
+            links("https://docs.oracle.com/en/java/javase/11/docs/api/")
+        }
+    }
+}
diff --git a/extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/LlmModuleConversationHistoryTest.kt b/extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/LlmModuleConversationHistoryTest.kt
new file mode 100644
index 00000000000..c75e2ee498d
--- /dev/null
+++ b/extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/LlmModuleConversationHistoryTest.kt
@@ -0,0 +1,196 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+package org.pytorch.executorch
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import java.io.File
+import java.io.IOException
+import org.apache.commons.io.FileUtils
+import org.junit.After
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertNotEquals
+import org.junit.Assert.assertTrue
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.pytorch.executorch.TestFileUtils.getTestFilePath
+import org.pytorch.executorch.extension.llm.LlmCallback
+import org.pytorch.executorch.extension.llm.LlmModule
+
+/**
+ * Behavioral tests for multi-turn / conversation-history semantics on [LlmModule].
+ *
+ * These tests run on the TinyStories-110M fixture pulled by `android_test_setup.sh`, which is too
+ * small and not instruction-tuned, so we cannot assert anything about the *content* of generated
+ * text (e.g. "did the model recall the user's name"). Instead, we assert structural invariants of
+ * the KV-cache + reset plumbing that any conversation-history feature depends on:
+ * 1. Determinism after [LlmModule.resetContext] at temperature=0 (greedy decode).
+ * 2. State preservation across successive [LlmModule.generate] calls (no reset → output diverges).
+ * 3. [LlmModule.prefillPrompt] influences the next [LlmModule.generate] call.
+ * 4. [LlmModule.resetContext] fully clears prefilled state.
+ *
+ * All tests run on both internal (fbsource Sandcastle) and OSS (GitHub Actions) Android CI because
+ * the fixture is fetched from the public `ossci-android` S3 bucket by `android_test_setup.sh` and
+ * the test only depends on the public `LlmModule` API.
+ */
+@RunWith(AndroidJUnit4::class)
+class LlmModuleConversationHistoryTest {
+
+  private lateinit var llmModule: LlmModule
+
+  @Before
+  @Throws(IOException::class)
+  fun setUp() {
+    val pteFile = File(getTestFilePath(TEST_FILE_NAME))
+    requireNotNull(javaClass.getResourceAsStream(TEST_FILE_NAME)) {
+          "Test resource $TEST_FILE_NAME not found; did android_test_setup.sh run?"
+        }
+        .use { pteStream -> FileUtils.copyInputStreamToFile(pteStream, pteFile) }
+
+    val tokenizerFile = File(getTestFilePath(TOKENIZER_FILE_NAME))
+    requireNotNull(javaClass.getResourceAsStream(TOKENIZER_FILE_NAME)) {
+          "Test resource $TOKENIZER_FILE_NAME not found; did android_test_setup.sh run?"
+        }
+        .use { tokenizerStream -> FileUtils.copyInputStreamToFile(tokenizerStream, tokenizerFile) }
+
+    llmModule =
+        LlmModule(getTestFilePath(TEST_FILE_NAME), getTestFilePath(TOKENIZER_FILE_NAME), 0.0f)
+    llmModule.load()
+  }
+
+  @After
+  fun tearDown() {
+    if (::llmModule.isInitialized) {
+      llmModule.close()
+    }
+  }
+
+  /**
+   * resetContext() + greedy decode (temperature=0) must produce identical output across two runs
+   * with the same prompt. This is the foundational invariant any conversation-history feature
+   * relies on: clearing the KV cache truly returns the model to a clean state.
+   */
+  @Test
+  @Throws(IOException::class)
+  fun testResetContextProducesDeterministicOutput() {
+    val firstRun = generateAndCollect(PROMPT_A)
+    llmModule.resetContext()
+    val secondRun = generateAndCollect(PROMPT_A)
+
+    assertTrue("Expected non-empty generation on first run", firstRun.isNotEmpty())
+    assertTrue("Expected non-empty generation on second run", secondRun.isNotEmpty())
+    assertEquals(
+        "Greedy generation after resetContext() must be deterministic for the same prompt.",
+        firstRun,
+        secondRun,
+    )
+  }
+
+  /**
+   * Without resetContext() between calls, KV-cache state persists and influences subsequent
+   * generation. Generating the same prompt twice in a row should produce different output the
+   * second time (because the KV cache is no longer empty and start position is non-zero), or the
+   * second call may throw because the runtime detects the stale KV state.
+   *
+   * Either outcome proves state persistence. If this test ever starts failing (i.e. both calls
+   * succeed with equal output), the runtime is silently dropping state between generate() calls —
+   * that would break multi-turn conversations.
+   */
+  @Test
+  @Throws(IOException::class)
+  fun testKvCacheStatePersistsAcrossGenerateCalls() {
+    val firstRun = generateAndCollect(PROMPT_A)
+    assertTrue("Expected non-empty generation on first run", firstRun.isNotEmpty())
+
+    try {
+      val secondRun = generateAndCollect(PROMPT_A)
+      assertNotEquals(
+          "Without resetContext(), repeated generate() calls must reflect persisted KV state.",
+          firstRun,
+          secondRun,
+      )
+    } catch (_: ExecutorchRuntimeException) {
+      // The second generate() threw because KV-cache state from the first call
+      // affected execution — this also proves state persistence.
+    }
+  }
+
+  /**
+   * prefillPrompt() must influence the next generate() — i.e. prefilled tokens are part of the
+   * conversation history. If prefilling has no effect, multi-turn flows that rely on injecting
+   * prior turns via prefill are broken.
+   */
+  @Test
+  @Throws(IOException::class)
+  fun testPrefillPromptInfluencesNextGeneration() {
+    val baselineRun = generateAndCollect(PROMPT_A)
+
+    llmModule.resetContext()
+    llmModule.prefillPrompt(PREFILL_HISTORY)
+    val withHistoryRun = generateAndCollect(PROMPT_A)
+
+    assertTrue("Expected non-empty baseline generation", baselineRun.isNotEmpty())
+    assertTrue("Expected non-empty post-prefill generation", withHistoryRun.isNotEmpty())
+    assertNotEquals(
+        "prefillPrompt() must alter the KV state seen by the next generate() call.",
+        baselineRun,
+        withHistoryRun,
+    )
+  }
+
+  /**
+   * resetContext() must fully clear prefilled state — running prefill then resetting then
+   * generating should match a clean-slate generation of the same prompt.
+   */
+  @Test
+  @Throws(IOException::class)
+  fun testResetContextClearsPrefilledHistory() {
+    val cleanRun = generateAndCollect(PROMPT_A)
+
+    llmModule.resetContext()
+    llmModule.prefillPrompt(PREFILL_HISTORY)
+    llmModule.resetContext()
+    val postResetRun = generateAndCollect(PROMPT_A)
+
+    assertTrue("Expected non-empty clean run", cleanRun.isNotEmpty())
+    assertTrue("Expected non-empty post-reset run", postResetRun.isNotEmpty())
+    assertEquals(
+        "resetContext() after a prefillPrompt() must fully clear KV state.",
+        cleanRun,
+        postResetRun,
+    )
+  }
+
+  private fun generateAndCollect(prompt: String): List {
+    val collector = CollectingCallback()
+    llmModule.generate(prompt, SEQ_LEN, collector)
+    return collector.tokens()
+  }
+
+  private class CollectingCallback : LlmCallback {
+    private val tokens: MutableList = ArrayList()
+
+    override fun onResult(result: String) {
+      tokens.add(result)
+    }
+
+    override fun onStats(stats: String) = Unit
+
+    fun tokens(): List = tokens.toList()
+  }
+
+  companion object {
+    private const val TEST_FILE_NAME = "/stories.pte"
+    private const val TOKENIZER_FILE_NAME = "/tokenizer.bin"
+
+    /** Short prompt; SEQ_LEN kept small to keep the test fast on CI emulators/devices. */
+    private const val PROMPT_A = "Once"
+    private const val PREFILL_HISTORY = "Long ago, in a small village by the sea, "
+    private const val SEQ_LEN = 24
+  }
+}
diff --git a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/Tensor.java b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/Tensor.java
index b86c464960d..f810ee6070f 100644
--- a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/Tensor.java
+++ b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/Tensor.java
@@ -554,6 +554,122 @@ public float[] getDataAsFloatArray() {
         "Tensor of type " + getClass().getSimpleName() + " cannot return data as float array.");
   }
 
+  /**
+   * Copies the tensor's data into a caller-provided {@link FloatBuffer}, avoiding the per-call
+   * {@code float[]} allocation that {@link #getDataAsFloatArray()} performs. The destination
+   * buffer's position is advanced by the number of elements written; its content from the starting
+   * position must have at least {@link #numel()} elements of remaining capacity.
+   *
+   * 
Useful in steady-state inference loops where the same output tensor shape is read every
+   * frame: pre-allocate a {@code FloatBuffer} once (e.g. via {@link #allocateFloatBuffer(int)}) and
+   * reuse it across calls.
+   *
+   * 
Supported by float32 (zero-copy bulk put) and float16 (per-element half→float widening,
+   * matching {@link #getDataAsFloatArray()} on that subclass). For raw fp16 bits without widening,
+   * use {@link #copyDataInto(ShortBuffer)}.
+   *
+   * @param dst the destination buffer; must have remaining capacity {@code >=} {@link #numel()}.
+   * @throws IllegalStateException if it is called for a tensor type that does not support a float
+   *     view.
+   * @throws java.nio.BufferOverflowException if {@code dst} does not have enough remaining
+   *     capacity.
+   */
+  public void copyDataInto(FloatBuffer dst) {
+    throw new IllegalStateException(
+        "Tensor of type " + getClass().getSimpleName() + " cannot copy data into FloatBuffer.");
+  }
+
+  /**
+   * Copies the tensor's data into a caller-provided {@link ByteBuffer}, avoiding the per-call
+   * {@code byte[]} allocation that {@link #getDataAsByteArray()} performs.
+   *
+   * @param dst the destination buffer; must have remaining capacity {@code >=} {@link #numel()}.
+   * @throws IllegalStateException if it is called for a non-int8 tensor.
+   * @throws java.nio.BufferOverflowException if {@code dst} does not have enough remaining
+   *     capacity.
+   */
+  public void copyDataInto(ByteBuffer dst) {
+    throw new IllegalStateException(
+        "Tensor of type " + getClass().getSimpleName() + " cannot copy data into ByteBuffer.");
+  }
+
+  /**
+   * Copies the tensor's data into a caller-provided {@link ByteBuffer}, avoiding the per-call
+   * {@code byte[]} allocation that {@link #getDataAsUnsignedByteArray()} performs. The bytes carry
+   * the raw uint8 bits — Java's signed {@code byte} representation, with values {@code >127}
+   * appearing negative; reinterpret with {@code & 0xFF} when reading.
+   *
+   * @param dst the destination buffer; must have remaining capacity {@code >=} {@link #numel()}.
+   * @throws IllegalStateException if it is called for a non-uint8 tensor.
+   * @throws java.nio.BufferOverflowException if {@code dst} does not have enough remaining
+   *     capacity.
+   */
+  public void copyDataIntoUnsigned(ByteBuffer dst) {
+    throw new IllegalStateException(
+        "Tensor of type "
+            + getClass().getSimpleName()
+            + " cannot copy data into ByteBuffer (unsigned).");
+  }
+
+  /**
+   * Copies the tensor's data into a caller-provided {@link IntBuffer}, avoiding the per-call {@code
+   * int[]} allocation that {@link #getDataAsIntArray()} performs.
+   *
+   * @param dst the destination buffer; must have remaining capacity {@code >=} {@link #numel()}.
+   * @throws IllegalStateException if it is called for a non-int32 tensor.
+   * @throws java.nio.BufferOverflowException if {@code dst} does not have enough remaining
+   *     capacity.
+   */
+  public void copyDataInto(IntBuffer dst) {
+    throw new IllegalStateException(
+        "Tensor of type " + getClass().getSimpleName() + " cannot copy data into IntBuffer.");
+  }
+
+  /**
+   * Copies the tensor's data into a caller-provided {@link LongBuffer}, avoiding the per-call
+   * {@code long[]} allocation that {@link #getDataAsLongArray()} performs.
+   *
+   * @param dst the destination buffer; must have remaining capacity {@code >=} {@link #numel()}.
+   * @throws IllegalStateException if it is called for a non-int64 tensor.
+   * @throws java.nio.BufferOverflowException if {@code dst} does not have enough remaining
+   *     capacity.
+   */
+  public void copyDataInto(LongBuffer dst) {
+    throw new IllegalStateException(
+        "Tensor of type " + getClass().getSimpleName() + " cannot copy data into LongBuffer.");
+  }
+
+  /**
+   * Copies the tensor's data into a caller-provided {@link DoubleBuffer}, avoiding the per-call
+   * {@code double[]} allocation that {@link #getDataAsDoubleArray()} performs.
+   *
+   * @param dst the destination buffer; must have remaining capacity {@code >=} {@link #numel()}.
+   * @throws IllegalStateException if it is called for a non-float64 tensor.
+   * @throws java.nio.BufferOverflowException if {@code dst} does not have enough remaining
+   *     capacity.
+   */
+  public void copyDataInto(DoubleBuffer dst) {
+    throw new IllegalStateException(
+        "Tensor of type " + getClass().getSimpleName() + " cannot copy data into DoubleBuffer.");
+  }
+
+  /**
+   * Copies the tensor's data into a caller-provided {@link ShortBuffer}, avoiding the per-call
+   * {@code short[]} allocation that {@link #getDataAsShortArray()} performs. For float16 tensors
+   * this writes the raw 16-bit half-precision bits with no widening; use {@link
+   * #copyDataInto(FloatBuffer)} if you want the values widened to fp32.
+   *
+   * @param dst the destination buffer; must have remaining capacity {@code >=} {@link #numel()}.
+   * @throws IllegalStateException if it is called for a tensor type whose backing storage is not a
+   *     {@code ShortBuffer}.
+   * @throws java.nio.BufferOverflowException if {@code dst} does not have enough remaining
+   *     capacity.
+   */
+  public void copyDataInto(ShortBuffer dst) {
+    throw new IllegalStateException(
+        "Tensor of type " + getClass().getSimpleName() + " cannot copy data into ShortBuffer.");
+  }
+
   /**
    * @return a Java long array that contains the tensor data. This may be a copy or reference.
    * @throws IllegalStateException if it is called for a non-int64 tensor.
@@ -604,6 +720,12 @@ public byte[] getDataAsUnsignedByteArray() {
       return arr;
     }
 
+    @Override
+    public void copyDataIntoUnsigned(ByteBuffer dst) {
+      data.rewind();
+      dst.put(data);
+    }
+
     @Override
     public String toString() {
       return String.format("Tensor(%s, dtype=torch.uint8)", Arrays.toString(shape));
@@ -636,6 +758,12 @@ public byte[] getDataAsByteArray() {
       return arr;
     }
 
+    @Override
+    public void copyDataInto(ByteBuffer dst) {
+      data.rewind();
+      dst.put(data);
+    }
+
     @Override
     public String toString() {
       return String.format("Tensor(%s, dtype=torch.int8)", Arrays.toString(shape));
@@ -668,6 +796,12 @@ public int[] getDataAsIntArray() {
       return arr;
     }
 
+    @Override
+    public void copyDataInto(IntBuffer dst) {
+      data.rewind();
+      dst.put(data);
+    }
+
     @Override
     public String toString() {
       return String.format("Tensor(%s, dtype=torch.int32)", Arrays.toString(shape));
@@ -690,6 +824,12 @@ public float[] getDataAsFloatArray() {
       return arr;
     }
 
+    @Override
+    public void copyDataInto(FloatBuffer dst) {
+      data.rewind();
+      dst.put(data);
+    }
+
     @Override
     public DType dtype() {
       return DType.FLOAT;
@@ -732,6 +872,12 @@ public short[] getDataAsShortArray() {
       return arr;
     }
 
+    @Override
+    public void copyDataInto(ShortBuffer dst) {
+      data.rewind();
+      dst.put(data);
+    }
+
     @Override
     public float[] getDataAsFloatArray() {
       data.rewind();
@@ -743,6 +889,21 @@ public float[] getDataAsFloatArray() {
       return arr;
     }
 
+    @Override
+    public void copyDataInto(FloatBuffer dst) {
+      data.rewind();
+      int remaining = data.remaining();
+      // Match the all-or-nothing semantics of bulk FloatBuffer.put(FloatBuffer):
+      // verify capacity up front so an undersized destination throws before any
+      // partial widening is observed in dst.
+      if (dst.remaining() < remaining) {
+        throw new java.nio.BufferOverflowException();
+      }
+      for (int i = 0; i < remaining; i++) {
+        dst.put(halfBitsToFloat(data.get()));
+      }
+    }
+
     @Override
     public String toString() {
       return String.format("Tensor(%s, dtype=torch.float16)", Arrays.toString(shape));
@@ -800,6 +961,12 @@ public long[] getDataAsLongArray() {
       return arr;
     }
 
+    @Override
+    public void copyDataInto(LongBuffer dst) {
+      data.rewind();
+      dst.put(data);
+    }
+
     @Override
     public String toString() {
       return String.format("Tensor(%s, dtype=torch.int64)", Arrays.toString(shape));
@@ -832,6 +999,12 @@ public double[] getDataAsDoubleArray() {
       return arr;
     }
 
+    @Override
+    public void copyDataInto(DoubleBuffer dst) {
+      data.rewind();
+      dst.put(data);
+    }
+
     @Override
     public String toString() {
       return String.format("Tensor(%s, dtype=torch.float64)", Arrays.toString(shape));
diff --git a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/package-info.java b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/package-info.java
index 2fcc8c9ec6b..86e19d09133 100644
--- a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/package-info.java
+++ b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/package-info.java
@@ -1,2 +1,51 @@
-/** Extension for LLM related use cases for ExecuTorch Android Java/JNI package. */
+/**
+ * ExecuTorch LLM extension for Android.
+ *
+ * 
This package provides Java bindings for running large language models (LLMs) on Android using
+ * ExecuTorch. It supports text generation, tokenization, and streaming token callbacks.
+ *
+ * 
Quick Start

+ *
+ * 
{@code
+ * import org.pytorch.executorch.extension.llm.LlmModule;
+ *
+ * // Load a Llama model
+ * LlmModule llm = new LlmModule(
+ *     "/data/local/tmp/llama.pte",
+ *     "/data/local/tmp/tokenizer.bin",
+ *     0.8f
+ * );
+ * llm.load();
+ *
+ * // Generate text token by token
+ * llm.generate("Hello, my name is", 200, new LlmCallback() {
+ *     public void onResult(String token) {
+ *         System.out.print(token);
+ *     }
+ *     public void onStats(String stats) {
+ *         System.out.println("\nStats: " + stats);
+ *     }
+ * });
+ * }

+ *
+ * 
Key Classes

+ *
+ * 
    
+ *   
  • {@link org.pytorch.executorch.extension.llm.LlmModule} — load and run an LLM
+ *   
  • {@link org.pytorch.executorch.extension.llm.LlmModuleConfig} — configure model paths and
+ *       settings
+ *   
  • {@link org.pytorch.executorch.extension.llm.LlmGenerationConfig} — control generation
+ *       (temperature, seq length)
+ * 

+ *
+ * 
More Resources

+ *
+ * 
+ */
 package org.pytorch.executorch.extension.llm;
diff --git a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/package-info.java b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/package-info.java
index 01d55ebc72b..7a5ed0bb5a5 100644
--- a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/package-info.java
+++ b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/package-info.java
@@ -1,2 +1,57 @@
-/** ExecuTorch Android Java/JNI package. This is the main package for generic use cases. */
+/**
+ * ExecuTorch Android Java API.
+ *
+ * 
This package provides Java bindings for running ExecuTorch models on Android. Use these
+ * classes to load a {@code .pte} model file and run inference directly from your Java or Kotlin
+ * Android app — no C++ required.
+ *
+ * 
Quick Start

+ *
+ * 
Step 1. Add the dependency to your {@code app/build.gradle.kts}:
+ *
+ * 
{@code
+ * dependencies {
+ *     implementation("org.pytorch:executorch-android:${executorch_version}")
+ * }
+ * }

+ *
+ * 
Step 2. Load your model and run inference:
+ *
+ * 
{@code
+ * import org.pytorch.executorch.EValue;
+ * import org.pytorch.executorch.Module;
+ * import org.pytorch.executorch.Tensor;
+ *
+ * // Load your exported .pte model file
+ * Module module = Module.load("/data/local/tmp/model.pte");
+ *
+ * // Build an input tensor  e.g. a 1x3x224x224 image
+ * float[] inputData = new float[1 * 3 * 224 * 224];
+ * Tensor inputTensor = Tensor.fromBlob(inputData, new long[]{1, 3, 224, 224});
+ *
+ * // Run inference
+ * EValue[] output = module.forward(EValue.from(inputTensor));
+ *
+ * // Read the result
+ * float[] scores = output[0].toTensor().getDataAsFloatArray();
+ * }

+ *
+ * 
Key Classes

+ *
+ * 
    
+ *   
  • {@link org.pytorch.executorch.Module} — load and run a {@code .pte} model
+ *   
  • {@link org.pytorch.executorch.Tensor} — create input tensors and read outputs
+ *   
  • {@link org.pytorch.executorch.EValue} — wrap inputs and unwrap outputs
+ *   
  • {@link org.pytorch.executorch.DType} — supported data types (FLOAT, INT32, etc.)
+ * 

+ *
+ * 
More Resources

+ *
+ * 
+ */
 package org.pytorch.executorch;
diff --git a/extension/android/executorch_android/src/main/javadoc/overview.html b/extension/android/executorch_android/src/main/javadoc/overview.html
new file mode 100644
index 00000000000..bd02ef2fcf9
--- /dev/null
+++ b/extension/android/executorch_android/src/main/javadoc/overview.html
@@ -0,0 +1,84 @@
+
+
+
+  ExecuTorch Android Java API
+
+
+
+

+  The ExecuTorch Android Java API lets you run PyTorch models on Android
+  devices using a simple Java or Kotlin interface.
+

+
+

+  ExecuTorch is PyTorch's solution for on-device AI — from smartphones to
+  microcontrollers. The Java API wraps the native ExecuTorch runtime and gives
+  you clean Java classes to load models, build tensors, and run inference.
+

+
+
Quick Start

+
+
Add the library to your app:

+
+
+// app/build.gradle.kts
+dependencies {
+    implementation("org.pytorch:executorch-android:${executorch_version}")
+}
+

+
+
Load a model and run inference:

+
+
+import org.pytorch.executorch.EValue;
+import org.pytorch.executorch.Module;
+import org.pytorch.executorch.Tensor;
+
+// Load your exported .pte model
+Module module = Module.load("/data/local/tmp/model.pte");
+
+// Create an input tensor (1x3x224x224 image)
+float[] data = new float[1 * 3 * 224 * 224];
+Tensor input = Tensor.fromBlob(data, new long[]{1, 3, 224, 224});
+
+// Run inference
+EValue[] output = module.forward(EValue.from(input));
+float[] scores = output[0].toTensor().getDataAsFloatArray();
+

+
+
Packages

+
+
+  
+    
+    
+  
+  
+    
+    
+  
+  
+    
+    
+  
+
org.pytorch.executorchCore API. Contains Module to load and run models, Tensor for tensor operations,
+        and EValue to wrap inputs and outputs.
org.pytorch.executorch.extension.llmLLM extension. Contains LlmModule for running large language models like Llama
+        with streaming token generation.
org.pytorch.executorch.annotationsAPI annotations. Experimental marks APIs that may change in future releases.

+
+
Resources

+
+
+
+
+
diff --git a/extension/android/executorch_android/src/test/java/org/pytorch/executorch/TensorTest.kt b/extension/android/executorch_android/src/test/java/org/pytorch/executorch/TensorTest.kt
index 24e19c1ded1..7d4cea59803 100644
--- a/extension/android/executorch_android/src/test/java/org/pytorch/executorch/TensorTest.kt
+++ b/extension/android/executorch_android/src/test/java/org/pytorch/executorch/TensorTest.kt
@@ -274,6 +274,353 @@ class TensorTest {
         .hasMessage("Tensor of type Tensor_float32 cannot return data as long array.")
   }
 
+  @Test
+  fun testCopyDataIntoFloat32() {
+    val data = floatArrayOf(Float.MIN_VALUE, 0f, 0.1f, Float.MAX_VALUE)
+    val shape = longArrayOf(2, 2)
+    val tensor = Tensor.fromBlob(data, shape)
+
+    val dst = Tensor.allocateFloatBuffer(4)
+    tensor.copyDataInto(dst)
+    assertEquals(4, dst.position())
+    dst.rewind()
+    for (i in data.indices) {
+      assertEquals(data[i].toDouble(), dst.get().toDouble(), 1e-5)
+    }
+
+    // Verify reuse: a second call refills the same buffer in place.
+    dst.rewind()
+    tensor.copyDataInto(dst)
+    assertEquals(4, dst.position())
+    dst.rewind()
+    for (i in data.indices) {
+      assertEquals(data[i].toDouble(), dst.get().toDouble(), 1e-5)
+    }
+  }
+
+  @Test
+  fun testCopyDataIntoFloat32_writesAtDstPosition() {
+    val data = floatArrayOf(1f, 2f, 3f, 4f)
+    val shape = longArrayOf(4)
+    val tensor = Tensor.fromBlob(data, shape)
+
+    // Pre-fill a larger buffer; copyDataInto should write at the current
+    // position and advance it, not overwrite from index 0.
+    val dst = Tensor.allocateFloatBuffer(8)
+    dst.put(floatArrayOf(-1f, -1f))
+    assertEquals(2, dst.position())
+    tensor.copyDataInto(dst)
+    assertEquals(6, dst.position())
+    dst.rewind()
+    assertEquals(-1f.toDouble(), dst.get().toDouble(), 0.0)
+    assertEquals(-1f.toDouble(), dst.get().toDouble(), 0.0)
+    for (i in data.indices) {
+      assertEquals(data[i].toDouble(), dst.get().toDouble(), 1e-5)
+    }
+  }
+
+  @Test
+  fun testCopyDataIntoFloat32_overflow() {
+    val data = floatArrayOf(1f, 2f, 3f, 4f)
+    val tensor = Tensor.fromBlob(data, longArrayOf(4))
+    val dst = Tensor.allocateFloatBuffer(2)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(java.nio.BufferOverflowException::class.java)
+  }
+
+  @Test
+  fun testCopyDataIntoFloat16() {
+    // 0x0000=+0, 0x3C00=1.0, 0x4000=2.0, 0xC000=-2.0
+    val halfBits =
+        shortArrayOf(0x0000.toShort(), 0x3C00.toShort(), 0x4000.toShort(), 0xC000.toShort())
+    val tensor = Tensor.fromBlob(halfBits, longArrayOf(4))
+    assertEquals(DType.HALF, tensor.dtype())
+
+    val dst = Tensor.allocateFloatBuffer(4)
+    tensor.copyDataInto(dst)
+    assertEquals(4, dst.position())
+    dst.rewind()
+    assertEquals(0.0, dst.get().toDouble(), 0.0)
+    assertEquals(1.0, dst.get().toDouble(), 0.0)
+    assertEquals(2.0, dst.get().toDouble(), 0.0)
+    assertEquals(-2.0, dst.get().toDouble(), 0.0)
+  }
+
+  @Test
+  fun testCopyDataIntoFloat16_overflowIsAtomic() {
+    // The fp16 path widens element-by-element rather than via bulk put. Verify
+    // that an undersized destination throws BufferOverflowException up front
+    // and leaves dst unmodified, matching the all-or-nothing semantics of the
+    // float32 / int / etc. paths.
+    val halfBits =
+        shortArrayOf(0x0000.toShort(), 0x3C00.toShort(), 0x4000.toShort(), 0xC000.toShort())
+    val tensor = Tensor.fromBlob(halfBits, longArrayOf(4))
+    val dst = Tensor.allocateFloatBuffer(2)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(java.nio.BufferOverflowException::class.java)
+    assertEquals(0, dst.position())
+  }
+
+  @Test
+  fun testCopyDataIntoFloat_unsupportedDtype() {
+    val tensor = Tensor.fromBlob(intArrayOf(1, 2, 3, 4), longArrayOf(4))
+    val dst = Tensor.allocateFloatBuffer(4)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(IllegalStateException::class.java)
+        .hasMessage("Tensor of type Tensor_int32 cannot copy data into FloatBuffer.")
+  }
+
+  @Test
+  fun testCopyDataIntoInt32() {
+    val data = intArrayOf(Int.MIN_VALUE, 0, 1, Int.MAX_VALUE)
+    val tensor = Tensor.fromBlob(data, longArrayOf(4))
+    val dst = Tensor.allocateIntBuffer(4)
+    tensor.copyDataInto(dst)
+    assertEquals(4, dst.position())
+    dst.rewind()
+    for (i in data.indices) {
+      assertEquals(data[i], dst.get())
+    }
+  }
+
+  @Test
+  fun testCopyDataIntoInt32_unsupportedDtype() {
+    val tensor = Tensor.fromBlob(floatArrayOf(1f, 2f), longArrayOf(2))
+    val dst = Tensor.allocateIntBuffer(2)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(IllegalStateException::class.java)
+        .hasMessage("Tensor of type Tensor_float32 cannot copy data into IntBuffer.")
+  }
+
+  @Test
+  fun testCopyDataIntoInt64() {
+    val data = longArrayOf(Long.MIN_VALUE, 0, 1, Long.MAX_VALUE)
+    val tensor = Tensor.fromBlob(data, longArrayOf(4))
+    val dst = Tensor.allocateLongBuffer(4)
+    tensor.copyDataInto(dst)
+    assertEquals(4, dst.position())
+    dst.rewind()
+    for (i in data.indices) {
+      assertEquals(data[i], dst.get())
+    }
+  }
+
+  @Test
+  fun testCopyDataIntoInt64_unsupportedDtype() {
+    val tensor = Tensor.fromBlob(floatArrayOf(1f, 2f), longArrayOf(2))
+    val dst = Tensor.allocateLongBuffer(2)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(IllegalStateException::class.java)
+        .hasMessage("Tensor of type Tensor_float32 cannot copy data into LongBuffer.")
+  }
+
+  @Test
+  fun testCopyDataIntoFloat64() {
+    val data = doubleArrayOf(Double.MIN_VALUE, 0.0, 0.1, Double.MAX_VALUE)
+    val tensor = Tensor.fromBlob(data, longArrayOf(4))
+    val dst = Tensor.allocateDoubleBuffer(4)
+    tensor.copyDataInto(dst)
+    assertEquals(4, dst.position())
+    dst.rewind()
+    for (i in data.indices) {
+      assertEquals(data[i], dst.get(), 1e-12)
+    }
+  }
+
+  @Test
+  fun testCopyDataIntoFloat64_unsupportedDtype() {
+    val tensor = Tensor.fromBlob(floatArrayOf(1f, 2f), longArrayOf(2))
+    val dst = Tensor.allocateDoubleBuffer(2)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(IllegalStateException::class.java)
+        .hasMessage("Tensor of type Tensor_float32 cannot copy data into DoubleBuffer.")
+  }
+
+  @Test
+  fun testCopyDataIntoInt8() {
+    val data = byteArrayOf(Byte.MIN_VALUE, 0, 1, Byte.MAX_VALUE)
+    val tensor = Tensor.fromBlob(data, longArrayOf(4))
+    val dst = Tensor.allocateByteBuffer(4)
+    tensor.copyDataInto(dst)
+    assertEquals(4, dst.position())
+    dst.rewind()
+    for (i in data.indices) {
+      assertEquals(data[i], dst.get())
+    }
+  }
+
+  @Test
+  fun testCopyDataIntoInt8_rejectsUInt8() {
+    val tensor = Tensor.fromBlobUnsigned(byteArrayOf(0, 1, 2, 3), longArrayOf(4))
+    val dst = Tensor.allocateByteBuffer(4)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(IllegalStateException::class.java)
+        .hasMessage("Tensor of type Tensor_uint8 cannot copy data into ByteBuffer.")
+  }
+
+  @Test
+  fun testCopyDataIntoUnsignedUInt8() {
+    val data = byteArrayOf(0, 1, 127, -1) // -1 == 255 unsigned
+    val tensor = Tensor.fromBlobUnsigned(data, longArrayOf(4))
+    val dst = Tensor.allocateByteBuffer(4)
+    tensor.copyDataIntoUnsigned(dst)
+    assertEquals(4, dst.position())
+    dst.rewind()
+    for (i in data.indices) {
+      assertEquals(data[i], dst.get())
+    }
+  }
+
+  @Test
+  fun testCopyDataIntoUnsigned_rejectsInt8() {
+    val tensor = Tensor.fromBlob(byteArrayOf(0, 1, 2, 3), longArrayOf(4))
+    val dst = Tensor.allocateByteBuffer(4)
+    assertThatThrownBy { tensor.copyDataIntoUnsigned(dst) }
+        .isInstanceOf(IllegalStateException::class.java)
+        .hasMessage("Tensor of type Tensor_int8 cannot copy data into ByteBuffer (unsigned).")
+  }
+
+  @Test
+  fun testCopyDataIntoByte_unsupportedDtype() {
+    val tensor = Tensor.fromBlob(floatArrayOf(1f, 2f), longArrayOf(2))
+    val dst = Tensor.allocateByteBuffer(2 * java.lang.Float.BYTES)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(IllegalStateException::class.java)
+        .hasMessage("Tensor of type Tensor_float32 cannot copy data into ByteBuffer.")
+    assertThatThrownBy { tensor.copyDataIntoUnsigned(dst) }
+        .isInstanceOf(IllegalStateException::class.java)
+        .hasMessage("Tensor of type Tensor_float32 cannot copy data into ByteBuffer (unsigned).")
+  }
+
+  @Test
+  fun testCopyDataIntoFloat16RawBits() {
+    val halfBits =
+        shortArrayOf(0x0000.toShort(), 0x3C00.toShort(), 0x4000.toShort(), 0xC000.toShort())
+    val tensor = Tensor.fromBlob(halfBits, longArrayOf(4))
+    assertEquals(DType.HALF, tensor.dtype())
+    val dst = Tensor.allocateHalfBuffer(4)
+    tensor.copyDataInto(dst)
+    assertEquals(4, dst.position())
+    dst.rewind()
+    for (i in halfBits.indices) {
+      assertEquals(halfBits[i], dst.get())
+    }
+  }
+
+  @Test
+  fun testCopyDataIntoShort_unsupportedDtype() {
+    val tensor = Tensor.fromBlob(floatArrayOf(1f, 2f), longArrayOf(2))
+    val dst = Tensor.allocateHalfBuffer(2)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(IllegalStateException::class.java)
+        .hasMessage("Tensor of type Tensor_float32 cannot copy data into ShortBuffer.")
+  }
+
+  @Test
+  fun testCopyDataIntoInt32_overflow() {
+    val tensor = Tensor.fromBlob(intArrayOf(1, 2, 3, 4), longArrayOf(4))
+    val dst = Tensor.allocateIntBuffer(2)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(java.nio.BufferOverflowException::class.java)
+  }
+
+  @Test
+  fun testCopyDataIntoInt64_overflow() {
+    val tensor = Tensor.fromBlob(longArrayOf(1, 2, 3, 4), longArrayOf(4))
+    val dst = Tensor.allocateLongBuffer(2)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(java.nio.BufferOverflowException::class.java)
+  }
+
+  @Test
+  fun testCopyDataIntoFloat64_overflow() {
+    val tensor = Tensor.fromBlob(doubleArrayOf(1.0, 2.0, 3.0, 4.0), longArrayOf(4))
+    val dst = Tensor.allocateDoubleBuffer(2)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(java.nio.BufferOverflowException::class.java)
+  }
+
+  @Test
+  fun testCopyDataIntoInt8_overflow() {
+    val tensor = Tensor.fromBlob(byteArrayOf(1, 2, 3, 4), longArrayOf(4))
+    val dst = Tensor.allocateByteBuffer(2)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(java.nio.BufferOverflowException::class.java)
+  }
+
+  @Test
+  fun testCopyDataIntoUnsignedUInt8_overflow() {
+    val tensor = Tensor.fromBlobUnsigned(byteArrayOf(1, 2, 3, 4), longArrayOf(4))
+    val dst = Tensor.allocateByteBuffer(2)
+    assertThatThrownBy { tensor.copyDataIntoUnsigned(dst) }
+        .isInstanceOf(java.nio.BufferOverflowException::class.java)
+  }
+
+  @Test
+  fun testCopyDataIntoFloat16RawBits_overflow() {
+    val halfBits =
+        shortArrayOf(0x0000.toShort(), 0x3C00.toShort(), 0x4000.toShort(), 0xC000.toShort())
+    val tensor = Tensor.fromBlob(halfBits, longArrayOf(4))
+    val dst = Tensor.allocateHalfBuffer(2)
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(java.nio.BufferOverflowException::class.java)
+  }
+
+  @Test
+  fun testCopyDataIntoFloat16_writesAtDstPosition() {
+    // 0x3C00=1.0, 0x4000=2.0, 0x4200=3.0, 0x4400=4.0
+    val halfBits =
+        shortArrayOf(0x3C00.toShort(), 0x4000.toShort(), 0x4200.toShort(), 0x4400.toShort())
+    val tensor = Tensor.fromBlob(halfBits, longArrayOf(4))
+
+    // Pre-fill a larger buffer; the fp16 widening path should write at the
+    // current position and advance it, not overwrite from index 0.
+    val dst = Tensor.allocateFloatBuffer(8)
+    dst.put(floatArrayOf(-1f, -1f))
+    assertEquals(2, dst.position())
+    tensor.copyDataInto(dst)
+    assertEquals(6, dst.position())
+    dst.rewind()
+    assertEquals(-1f.toDouble(), dst.get().toDouble(), 0.0)
+    assertEquals(-1f.toDouble(), dst.get().toDouble(), 0.0)
+    assertEquals(1.0, dst.get().toDouble(), 0.0)
+    assertEquals(2.0, dst.get().toDouble(), 0.0)
+    assertEquals(3.0, dst.get().toDouble(), 0.0)
+    assertEquals(4.0, dst.get().toDouble(), 0.0)
+  }
+
+  @Test
+  fun testCopyDataInto_emptyTensor() {
+    val floatTensor = Tensor.fromBlob(floatArrayOf(), longArrayOf(0))
+    val floatDst = Tensor.allocateFloatBuffer(4)
+    floatDst.put(floatArrayOf(7f, 8f))
+    assertEquals(2, floatDst.position())
+    floatTensor.copyDataInto(floatDst)
+    assertEquals(2, floatDst.position())
+
+    // Same for the fp16 widening path, whose explicit remaining() check is
+    // worth exercising at zero.
+    val halfTensor = Tensor.fromBlob(shortArrayOf(), longArrayOf(0))
+    val halfWidenDst = Tensor.allocateFloatBuffer(2)
+    halfTensor.copyDataInto(halfWidenDst)
+    assertEquals(0, halfWidenDst.position())
+  }
+
+  @Test
+  fun testCopyDataInto_rejectsReadOnlyBuffer() {
+    val tensor = Tensor.fromBlob(floatArrayOf(1f, 2f, 3f, 4f), longArrayOf(4))
+    val dst = Tensor.allocateFloatBuffer(4).asReadOnlyBuffer()
+    assertThatThrownBy { tensor.copyDataInto(dst) }
+        .isInstanceOf(java.nio.ReadOnlyBufferException::class.java)
+
+    // fp16 widening path uses element-by-element put, which also rejects.
+    val halfTensor =
+        Tensor.fromBlob(shortArrayOf(0x3C00.toShort(), 0x4000.toShort()), longArrayOf(2))
+    val halfWidenDst = Tensor.allocateFloatBuffer(2).asReadOnlyBuffer()
+    assertThatThrownBy { halfTensor.copyDataInto(halfWidenDst) }
+        .isInstanceOf(java.nio.ReadOnlyBufferException::class.java)
+  }
+
   @Test
   fun testIllegalArguments() {
     val data = floatArrayOf(Float.MIN_VALUE, 0f, 0.1f, Float.MAX_VALUE)
diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm
index 53b7d997c6a..7dda8898e25 100644
--- a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm
+++ b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm
@@ -252,27 +252,6 @@ @implementation ExecuTorchModule {
   std::unique_ptr _module;
   NSMutableDictionary *> *_inputs;
   NSMutableDictionary *> *_outputs;
-  // Strong reference to the most recently passed BackendOptionsMap. The
-  // C++ Module borrows a pointer into the map's underlying C++ storage and
-  // dereferences it during lazy load_method calls (triggered by forward),
-  // so the ObjC wrapper must keep it alive. ARC handles the lifetime.
-  //
-  // INVARIANT: this ivar is only ever overwritten with another non-nil
-  // BackendOptionsMap, and never reset to nil while `_module` is alive.
-  // Resetting to nil would release the C++ map while `_module` still holds
-  // a borrowed pointer into it.
-  //
-  // THREAD SAFETY: like the rest of `ExecuTorchModule`, write access here
-  // is not thread-safe. The ARC retain/release on assignment is non-atomic
-  // for direct ivars; serialize `loadWithOptions:` calls externally if you
-  // share a `Module` across threads.
-  //
-  // TODO: remove this ivar once the C++ Module owns its LoadBackendOptionsMap
-  // by value (today it borrows a raw pointer). With owned options the ObjC
-  // wrapper has nothing to retain, the thread-safety caveat above goes
-  // away, and -loadMethod:options: / -loadWithOptions: stop needing a
-  // custom lifetime contract between the bindings and the C++ layer.
-  ExecuTorchBackendOptionsMap *_loadedBackendOptions;
 }
 
 - (instancetype)initWithFilePath:(NSString *)filePath
@@ -354,25 +333,10 @@ - (BOOL)loadWithOptions:(ExecuTorchBackendOptionsMap *)options
            verification:(ExecuTorchVerification)verification
                   error:(NSError **)error {
   NSParameterAssert(options);
-  // Retain the options object so the C++ borrowed pointer it contains stays
-  // valid for the lifetime of any methods loaded with these options.
-  // (Methods load lazily during forward(), so the borrow may outlive this
-  // call.) See ExecuTorchBackendOptionsMap.h for the lifetime contract.
-  //
-  // No rollback on failure: Module::load updates its backend_options_ raw
-  // pointer BEFORE attempting load_internal, so after a failed call the
-  // C++ side already references `options`. The ObjC retain therefore
-  // always matches what C++ points at, even on the failure path — a
-  // two-phase commit here would instead leave C++ pointing at a map the
-  // wrapper no longer retains. See:
-  // https://github.com/pytorch/executorch/blob/6412f55a54dd3ce1f4ed220a3e96ee19b8f37967/extension/module/module.cpp#L192-L197
-  //
-  // TODO: once Module::load is made transactional (i.e. it only commits
-  // `backend_options_` after load_internal succeeds), replace the
-  // unconditional assignment below with a proper two-phase commit that
-  // only overwrites _loadedBackendOptions on success. This removes the
-  // "match C++'s unconditional write" workaround documented above.
-  _loadedBackendOptions = options;
+  // Module deep-copies the LoadBackendOptionsMap on the C++ side, so we
+  // do not need to retain `options` past this call. ARC releases the
+  // wrapper when the parameter goes out of scope and the Module's owned
+  // copy keeps lazy load_method paths working.
   const auto errorCode = _module->load(*[options cppMap],
                                         static_cast(verification));
   if (errorCode != Error::Ok) {
@@ -395,22 +359,10 @@ - (BOOL)loadMethod:(NSString *)methodName
            options:(ExecuTorchBackendOptionsMap *)options
              error:(NSError **)error {
   NSParameterAssert(options);
-  // Do NOT assign to _loadedBackendOptions here. Module::load_method
-  // consumes `backend_options` synchronously within this call — it is
-  // passed through to program_->load_method and is not cached on the
-  // Module. Only Module::load(backend_options, ...) stores the pointer
-  // (via backend_options_). ARC keeps `options` alive for the call
-  // duration via the parameter, so no ivar retention is needed here.
-  // See:
-  //   load_method: https://github.com/pytorch/executorch/blob/6412f55a54dd3ce1f4ed220a3e96ee19b8f37967/extension/module/module.cpp#L353-L409
-  //   load (stores backend_options_): https://github.com/pytorch/executorch/blob/6412f55a54dd3ce1f4ed220a3e96ee19b8f37967/extension/module/module.cpp#L195
-  //
-  // Overwriting _loadedBackendOptions would release any map previously
-  // installed by -loadWithOptions:, but the C++ Module's backend_options_
-  // raw pointer would still reference that released map's storage — a
-  // use-after-free on the next lazy load_method. The XCTest
-  // testMixedLoadWithOptionsAndLoadMethodWithOptionsOnMultiMethodModel
-  // pins this invariant via a weak reference.
+  // load_method consumes `options` synchronously: the cppMap pointer is
+  // passed through to program_->load_method and is not cached on the C++
+  // Module. ARC keeps `options` alive for the duration of this call via
+  // the parameter, so no extra retention is needed here.
   const auto errorCode = _module->load_method(methodName.UTF8String,
                                                /*planned_memory=*/nullptr,
                                                /*event_tracer=*/nullptr,
diff --git a/extension/apple/ExecuTorch/__tests__/ModuleTest.swift b/extension/apple/ExecuTorch/__tests__/ModuleTest.swift
index db75345f071..deea7d57189 100644
--- a/extension/apple/ExecuTorch/__tests__/ModuleTest.swift
+++ b/extension/apple/ExecuTorch/__tests__/ModuleTest.swift
@@ -525,22 +525,23 @@ class ModuleTest: XCTestCase {
   }
 
   // Mixed sequence on a multi-method delegated model:
-  //   1. load(optionsA)            — installs optionsA; the C++ Module
-  //      stores a raw pointer into optionsA's storage and the ObjC
-  //      wrapper retains optionsA via _loadedBackendOptions.
+  //   1. load(optionsA)              — installs optionsA. The C++ Module
+  //      deep-copies optionsA at load time; the Apple binding does NOT
+  //      retain the input map. After the autoreleasepool drains, the
+  //      caller's optionsA is dropped — the C++ Module's internal copy
+  //      keeps the configuration alive.
   //   2. load("mul", options: optionsB) — loads "mul" explicitly with
-  //      optionsB, synchronously. Must NOT release optionsA (doing so
-  //      would leave _module->backend_options_ dangling).
-  //   3. forward(inputs)           — triggers a lazy load_method on
-  //      "forward" which falls back to the stored pointer (into optionsA).
+  //      optionsB; the C++ Module deep-copies optionsB.
+  //   3. forward(inputs)             — triggers a lazy load_method on
+  //      "forward" which consults the deep-copied optionsA stored
+  //      inside the C++ Module.
   //
-  // The XCTAssertNotNil(weakA) after step 2 is the deterministic check:
-  // a buggy loadMethod:options: that assigns `_loadedBackendOptions =
-  // options` releases optionsA's last strong ref there, weakA becomes
-  // nil, and the assertion fails independent of heap layout. With the
-  // correct implementation weakA stays non-nil. The forward/execute
-  // assertions additionally verify the positive path end-to-end.
-  func testMixedLoadWithOptionsAndLoadMethodWithOptionsOnMultiMethodModel() throws {
+  // The XCTAssertNil(weakA) after step 1 is the deterministic check:
+  // it proves the Apple binding does not silently retain the input map
+  // (a regression would re-introduce a strong reference). The forward
+  // assertions verify the positive path: even though the caller has
+  // dropped optionsA, the deep-copied configuration is still applied.
+  func testMultiMethodLoadAndForwardWithBackendOptions() throws {
     let modelPath = try requireFixture("add_mul_coreml", ofType: "pte")
     let module = Module(filePath: modelPath)
 
@@ -552,7 +553,9 @@ class ModuleTest: XCTestCase {
       weakA = optionsA
       try module.load(options: optionsA)
     }
-    XCTAssertNotNil(weakA, "Module must retain optionsA after load(optionsA)")
+    XCTAssertNil(weakA,
+      "Apple binding must not retain the input BackendOptionsMap — " +
+      "the C++ Module deep-copies it at load time")
 
     try autoreleasepool {
       let optionsB = try BackendOptionsMap(options: [
@@ -561,11 +564,9 @@ class ModuleTest: XCTestCase {
       try module.load("mul", options: optionsB)
     }
     XCTAssertTrue(module.isLoaded("mul"))
-    XCTAssertNotNil(weakA,
-      "load(\"mul\", options: optionsB) must not release optionsA — " +
-      "_module->backend_options_ still points into its storage")
 
-    // Lazy load_method("forward") must still see valid optionsA storage.
+    // Lazy load_method("forward") must still apply the configuration
+    // from the now-released optionsA, via the C++ Module's deep copy.
     let inputs: [Tensor] = [Tensor([2]), Tensor([3])]
     let addOuts: [Value] = try module.forward(inputs)
     XCTAssertEqual(addOuts.first?.tensor(), Tensor([Float(5)]))
diff --git a/extension/apple/ExecuTorch/__tests__/ObjC/ModuleTestObjC.m b/extension/apple/ExecuTorch/__tests__/ObjC/ModuleTestObjC.m
index 69efaa25304..491f1e2d25f 100644
--- a/extension/apple/ExecuTorch/__tests__/ObjC/ModuleTestObjC.m
+++ b/extension/apple/ExecuTorch/__tests__/ObjC/ModuleTestObjC.m
@@ -15,6 +15,55 @@ @interface ModuleTestObjC : XCTestCase
 
 @implementation ModuleTestObjC
 
+// Pins the deep-copy contract: -loadWithOptions: must not retain the
+// options wrapper. The C++ Module deep-copies the LoadBackendOptionsMap
+// (and its BackendOption arrays) into Module-owned storage, so dropping
+// the ObjC wrapper after loadWithOptions: returns must be safe and a
+// subsequent forward() must continue to work using the Module's owned
+// copy. If a future refactor reverts to the borrowed-pointer design and
+// reintroduces an ivar to retain the wrapper, this test fails (the
+// weak reference stays non-nil).
+- (void)testLoadWithOptionsDoesNotRetainOptions {
+  NSString *modelPath = [self requireFixture:@"add_coreml" ofType:@"pte"];
+  if (!modelPath) return;
+  NSError *error = nil;
+  ExecuTorchModule *module = [[ExecuTorchModule alloc] initWithFilePath:modelPath];
+
+  __weak ExecuTorchBackendOptionsMap *weakOptions = nil;
+  @autoreleasepool {
+    ExecuTorchBackendOptionsMap *options = [ExecuTorchBackendOptionsMap mapWithOptions:@{
+      @"CoreMLBackend": @[
+        [ExecuTorchBackendOption optionWithKey:@"compute_unit" stringValue:@"cpu_only"],
+      ],
+    } error:&error];
+    XCTAssertNotNil(options, @"%@", error);
+    weakOptions = options;
+    XCTAssertTrue([module loadWithOptions:options error:&error], @"%@", error);
+  }
+  // Local + autoreleased refs have drained. With deep-copy, the wrapper
+  // can be reclaimed and the C++ Module keeps its own copy of the
+  // backend options for any future lazy load_method calls.
+  XCTAssertNil(weakOptions,
+      @"loadWithOptions: must not retain the map (Module deep-copies). "
+      @"See module.cpp Module::load(LoadBackendOptionsMap, ...) deep-copy.");
+
+  // Forward must still execute against the Module-owned options copy.
+  ExecuTorchTensor *one =
+      [[ExecuTorchTensor alloc] initWithScalars:@[@1.0f] dataType:ExecuTorchDataTypeFloat];
+  NSArray *outputs =
+      [module forwardWithTensors:@[one, one] error:&error];
+  XCTAssertNotNil(outputs, @"%@", error);
+
+  __block float result = NAN;
+  [outputs.firstObject.tensorValue
+      bytesWithHandler:^(const void *bytes, NSInteger count, ExecuTorchDataType dt) {
+    if (dt == ExecuTorchDataTypeFloat && count >= 1) {
+      result = ((const float *)bytes)[0];
+    }
+  }];
+  XCTAssertEqual(result, 2.0f);
+}
+
 - (NSBundle *)resourceBundle {
 #if SWIFT_PACKAGE
   return SWIFTPM_MODULE_BUNDLE;
diff --git a/extension/data_loader/mman.h b/extension/data_loader/mman.h
index 26a9ee08067..a7a335961c8 100644
--- a/extension/data_loader/mman.h
+++ b/extension/data_loader/mman.h
@@ -17,6 +17,7 @@
 
 #ifndef _WIN32
 
+#include 
 #include 
 #include 
 
@@ -43,6 +44,34 @@ ET_INLINE off_t get_mmap_offset(size_t offset) {
   return static_cast(offset);
 }
 
+/**
+ * Hint the kernel to prefetch pages eagerly and to optimize for sequential
+ * reads. Intended to reduce page-fault stutter during model initialization
+ * when the caller does not want to mlock the pages into RAM.
+ */
+ET_INLINE void madvise_pages_willneed_sequential(void* addr, size_t len) {
+  ::madvise(addr, len, MADV_WILLNEED);
+  ::madvise(addr, len, MADV_SEQUENTIAL);
+}
+
+/**
+ * On Apple platforms, schedule kernel read-ahead on the file descriptor itself
+ * via fcntl(F_RDADVISE). This is more aggressive than madvise for cold starts:
+ * it brings pages into the unified buffer cache so first-touch faults are
+ * serviced from RAM instead of storage. No-op on non-Apple POSIX platforms.
+ */
+ET_INLINE void fcntl_rdadvise_apple(int fd, size_t file_size) {
+#if defined(__APPLE__)
+  struct radvisory advice;
+  advice.ra_offset = 0;
+  advice.ra_count = static_cast(file_size);
+  ::fcntl(fd, F_RDADVISE, &advice);
+#else
+  (void)fd;
+  (void)file_size;
+#endif
+}
+
 #else
 
 #define NOMINMAX
@@ -80,4 +109,21 @@ ET_INLINE uint64_t get_mmap_offset(size_t offset) {
   return static_cast(offset);
 }
 
+/**
+ * No-op on Windows: there is no direct equivalent to madvise(MADV_WILLNEED |
+ * MADV_SEQUENTIAL) and the existing mman_windows shim does not implement one.
+ */
+ET_INLINE void madvise_pages_willneed_sequential(void* addr, size_t len) {
+  (void)addr;
+  (void)len;
+}
+
+/**
+ * No-op on Windows: F_RDADVISE is an Apple-specific fcntl command.
+ */
+ET_INLINE void fcntl_rdadvise_apple(int fd, size_t file_size) {
+  (void)fd;
+  (void)file_size;
+}
+
 #endif
diff --git a/extension/data_loader/mmap_data_loader.cpp b/extension/data_loader/mmap_data_loader.cpp
index 5d77b67cc59..dc9e1a615bf 100644
--- a/extension/data_loader/mmap_data_loader.cpp
+++ b/extension/data_loader/mmap_data_loader.cpp
@@ -249,6 +249,11 @@ Result MmapDataLoader::load(
     // No need to keep track of this. munmap() will unlock as a side effect.
   }
 
+  if (mlock_config_ == MlockConfig::UseMadvise) {
+    madvise_pages_willneed_sequential(pages, map_size);
+    fcntl_rdadvise_apple(fd_, file_size_);
+  }
+
   // The requested data is at an offset into the mapped pages.
   const void* data = static_cast(pages) + offset - range.start;
 
diff --git a/extension/data_loader/mmap_data_loader.h b/extension/data_loader/mmap_data_loader.h
index c0496a39d4b..2bbdd96013b 100644
--- a/extension/data_loader/mmap_data_loader.h
+++ b/extension/data_loader/mmap_data_loader.h
@@ -38,6 +38,10 @@ class MmapDataLoader final : public executorch::runtime::DataLoader {
     UseMlock,
     /// Call `mlock()` on loaded pages, ignoring errors if it fails.
     UseMlockIgnoreErrors,
+    /// Use madvise(MADV_WILLNEED | MADV_SEQUENTIAL) instead of mlock.
+    /// Tells the kernel to prefetch pages eagerly and optimize for
+    /// sequential reads, without pinning them in RAM.
+    UseMadvise,
   };
 
   /**
diff --git a/extension/data_loader/test/mmap_data_loader_test.cpp b/extension/data_loader/test/mmap_data_loader_test.cpp
index df071fd7474..e08001af245 100644
--- a/extension/data_loader/test/mmap_data_loader_test.cpp
+++ b/extension/data_loader/test/mmap_data_loader_test.cpp
@@ -244,6 +244,12 @@ TEST_F(MmapDataLoaderTest, InBoundsLoadsSucceedUseMlockIgnoreErrors) {
       MmapDataLoader::MlockConfig::UseMlockIgnoreErrors);
 }
 
+TEST_F(MmapDataLoaderTest, InBoundsLoadsSucceedUseMadvise) {
+  // There's no portable way to verify madvise() is called, but exercise the
+  // path to make sure the code still behaves correctly.
+  test_in_bounds_loads_succeed(MmapDataLoader::MlockConfig::UseMadvise);
+}
+
 TEST_F(MmapDataLoaderTest, FinalPageOfUnevenFileSucceeds) {
   // Create a file whose length is not an even multiple of a page.
   // Each 4-byte word in the file has a different value.
diff --git a/extension/llm/runner/irunner.h b/extension/llm/runner/irunner.h
index c9584dda879..bb7dd767fea 100644
--- a/extension/llm/runner/irunner.h
+++ b/extension/llm/runner/irunner.h
@@ -33,14 +33,6 @@ struct GenerationConfig {
   // Whether to echo the input prompt in the output
   bool echo = true;
 
-  // Grammar definition for constrained decoding (e.g. a JSON schema, regex,
-  // Lark CFG, or GBNF grammar). Empty string means no constraint.
-  std::string grammar;
-
-  // Grammar format: "json_schema", "regex", "lark", or "gbnf".
-  // Only used when grammar is non-empty.
-  std::string grammar_type;
-
   // Whether to ignore EOS token and continue generating until max_new_tokens
   bool ignore_eos = false;
 
diff --git a/extension/module/module.cpp b/extension/module/module.cpp
index ec7236276f5..5422fb15b71 100644
--- a/extension/module/module.cpp
+++ b/extension/module/module.cpp
@@ -20,6 +20,7 @@ namespace extension {
 namespace ET_MODULE_NAMESPACE {
 
 using ET_MERGED_DATA_MAP_NAMESPACE::MergedDataMap;
+using ET_RUNTIME_NAMESPACE::Kernel;
 using ET_RUNTIME_NAMESPACE::MethodMeta;
 using ET_RUNTIME_NAMESPACE::Program;
 
@@ -70,6 +71,17 @@ runtime::Result> make_data_loader(
           std::move(*res_mlock_ignore));
       break;
     }
+    case Module::LoadMode::MmapUseMadvise: {
+      auto res_madvise = MmapDataLoader::from(
+          file_path.c_str(), MmapDataLoader::MlockConfig::UseMadvise);
+      if (!res_madvise.ok()) {
+        return res_madvise.error();
+      }
+      data_loader =
+          std::make_unique>(
+              std::move(*res_madvise));
+      break;
+    }
   }
   return data_loader;
 }
@@ -192,8 +204,49 @@ runtime::Error Module::load(const Program::Verification verification) {
 runtime::Error Module::load(
     const LoadBackendOptionsMap& backend_options,
     const Program::Verification verification) {
-  backend_options_ = &backend_options;
-  return load_internal(verification);
+  // load_internal does not read backend options, so run it first; on
+  // failure we skip the deep-copy work entirely and leave the prior
+  // installed options (if any) in place.
+  ET_CHECK_OK_OR_RETURN_ERROR(load_internal(verification));
+
+  // Deep-copy the input into local storage so the Module owns the
+  // BackendOption arrays for the lifetime of any methods loaded with
+  // these options. Build BOTH the storage and the map in locals so any
+  // mid-loop failure (or exception from emplace) leaves the prior
+  // installed state untouched -- the two members are only committed
+  // together at the end on full success.
+  //
+  // local_storage is reserve()'d up front, so emplace_back() never
+  // reallocates the outer buffer and the inner vectors keep stable
+  // addresses while we build local_map. The final move of
+  // local_storage into backend_options_storage_ uses std::vector's
+  // O(1) buffer transfer, so the heap buffers that local_map's spans
+  // point into remain valid after the move; the static_assert documents
+  // the inner-vector property we rely on for that span stability.
+  static_assert(
+      std::is_nothrow_move_constructible_v>,
+      "Moving local_storage must not move-construct the inner vectors; "
+      "local_map's spans reference their heap buffers.");
+
+  std::vector> local_storage;
+  local_storage.reserve(backend_options.size());
+  LoadBackendOptionsMap local_map;
+  for (size_t i = 0; i < backend_options.size(); ++i) {
+    const auto entry = backend_options.entry_at(i);
+    local_storage.emplace_back(entry.options.begin(), entry.options.end());
+    auto& owned = local_storage.back();
+    // The input map was already valid, so set_options should not fail
+    // here; assert it loudly rather than leaving partial state behind.
+    ET_CHECK_OK_OR_RETURN_ERROR(local_map.set_options(
+        entry.backend_id,
+        runtime::Span(owned.data(), owned.size())));
+  }
+
+  // Single commit point: both members updated together.
+  backend_options_storage_ = std::move(local_storage);
+  backend_options_map_ = std::move(local_map);
+
+  return runtime::Error::Ok;
 }
 
 runtime::Error Module::load_internal(const Program::Verification verification) {
@@ -354,13 +407,17 @@ runtime::Error Module::load_method(
     const std::string& method_name,
     runtime::HierarchicalAllocator* planned_memory,
     torch::executor::EventTracer* event_tracer,
-    const LoadBackendOptionsMap* backend_options) {
+    const LoadBackendOptionsMap* backend_options,
+    std::vector kernel_registry) {
   if (!is_method_loaded(method_name)) {
     ET_CHECK_OK_OR_RETURN_ERROR(load());
 
-    // Use passed backend_options, or fall back to stored one from load()
-    const LoadBackendOptionsMap* effective_backend_options =
-        backend_options ? backend_options : backend_options_;
+    // Use passed backend_options, or fall back to stored ones from load().
+    // An empty stored map behaves identically to nullptr downstream, so we
+    // only forward the stored map when it actually has entries.
+    const LoadBackendOptionsMap* effective_backend_options = backend_options
+        ? backend_options
+        : (backend_options_map_.size() > 0 ? &backend_options_map_ : nullptr);
 
     MethodHolder method_holder;
 
@@ -391,12 +448,16 @@ runtime::Error Module::load_method(
 
     method_holder.memory_manager = std::make_unique(
         memory_allocator_.get(), planned_memory, temp_allocator_.get());
+    method_holder.kernel_registry = std::move(kernel_registry);
     auto res_method = program_->load_method(
         method_name.c_str(),
         method_holder.memory_manager.get(),
         event_tracer ? event_tracer : this->event_tracer(),
         merged_data_map_.get(),
-        effective_backend_options);
+        effective_backend_options,
+        runtime::Span(
+            method_holder.kernel_registry.data(),
+            method_holder.kernel_registry.size()));
     if (!res_method.ok()) {
       return res_method.error();
     }
diff --git a/extension/module/module.h b/extension/module/module.h
index 08a68b2676b..47ead23032e 100644
--- a/extension/module/module.h
+++ b/extension/module/module.h
@@ -14,6 +14,8 @@
 #include 
 #include 
 
+#include 
+#include 
 #include 
 
 #ifdef USE_ATEN_LIB
@@ -25,6 +27,7 @@
 namespace executorch {
 namespace extension {
 
+using ET_RUNTIME_NAMESPACE::Kernel;
 using ET_RUNTIME_NAMESPACE::Method;
 using ET_RUNTIME_NAMESPACE::MethodMeta;
 using ET_RUNTIME_NAMESPACE::NamedDataMap;
@@ -51,6 +54,8 @@ class Module {
     MmapUseMlock,
     /// Use memory locking and ignore errors.
     MmapUseMlockIgnoreErrors,
+    /// Use mmap with madvise(MADV_WILLNEED | MADV_SEQUENTIAL) hints.
+    MmapUseMadvise,
   };
 
   /**
@@ -185,9 +190,18 @@ class Module {
   /**
    * Loads the program with per-delegate runtime options.
    *
-   * @param[in] backend_options A LoadBackendOptionsMap containing per-delegate
-   * load-time configuration options. The caller must ensure this object
-   * outlives any methods loaded with these options.
+   * The Module deep-copies `backend_options` into internal storage, so the
+   * caller may release the input (and any backing BackendOption arrays its
+   * Spans referenced) immediately after this call returns. Future lazy
+   * `load_method` calls (e.g. triggered by `forward`) consume the
+   * Module-owned copy.
+   *
+   * Transactional: on failure, the previously-installed backend options
+   * (if any) are left in place; the input is not committed.
+   *
+   * @param[in] backend_options A LoadBackendOptionsMap containing
+   * per-delegate load-time configuration options. Deep-copied into the
+   * Module on success; not retained on failure.
    * @param[in] verification The type of verification to do before returning
    * success.
    *
@@ -198,6 +212,21 @@ class Module {
       const Program::Verification verification =
           Program::Verification::Minimal);
 
+  /**
+   * Returns the deep-copied LoadBackendOptionsMap most recently installed
+   * via `load(LoadBackendOptionsMap, ...)`. The returned reference is owned
+   * by the Module and remains valid until the next call to
+   * `load(LoadBackendOptionsMap, ...)` or until the Module is destroyed.
+   *
+   * If `load(LoadBackendOptionsMap, ...)` has never been called, returns a
+   * default-constructed (empty, `size() == 0`) map.
+   *
+   * @returns Const reference to the Module-owned LoadBackendOptionsMap.
+   */
+  inline const LoadBackendOptionsMap& backend_options() const {
+    return backend_options_map_;
+  }
+
   /**
    * Checks if the program is loaded.
    *
@@ -253,7 +282,8 @@ class Module {
       const std::string& method_name,
       runtime::HierarchicalAllocator* planned_memory = nullptr,
       torch::executor::EventTracer* event_tracer = nullptr,
-      const LoadBackendOptionsMap* backend_options = nullptr);
+      const LoadBackendOptionsMap* backend_options = nullptr,
+      std::vector kernel_registry = {});
 
   ET_DEPRECATED ET_NODISCARD runtime::Error inline load_method(
       const std::string& method_name,
@@ -301,9 +331,14 @@ class Module {
   ET_NODISCARD inline runtime::Error load_forward(
       runtime::HierarchicalAllocator* planned_memory = nullptr,
       torch::executor::EventTracer* event_tracer = nullptr,
-      const LoadBackendOptionsMap* backend_options = nullptr) {
+      const LoadBackendOptionsMap* backend_options = nullptr,
+      std::vector kernel_registry = {}) {
     return load_method(
-        "forward", planned_memory, event_tracer, backend_options);
+        "forward",
+        planned_memory,
+        event_tracer,
+        backend_options,
+        std::move(kernel_registry));
   }
 
   ET_DEPRECATED ET_NODISCARD inline runtime::Error load_forward(
@@ -696,6 +731,7 @@ class Module {
     std::unique_ptr planned_memory;
     std::unique_ptr memory_manager;
     std::unique_ptr method;
+    std::vector kernel_registry;
   };
 
   std::string file_path_;
@@ -711,7 +747,14 @@ class Module {
   std::unique_ptr merged_data_map_;
   std::vector> shared_arenas_;
   ET_DEPRECATED std::vector debug_buffer_;
-  const LoadBackendOptionsMap* backend_options_ = nullptr;
+  // Module-owned deep-copy of the backend options most recently installed
+  // via load(LoadBackendOptionsMap, ...). `backend_options_storage_` owns
+  // the per-backend BackendOption arrays; `backend_options_map_` is a
+  // LoadBackendOptionsMap whose Spans reference those owned arrays. An
+  // empty map (`size() == 0`) is observationally indistinguishable from
+  // "never set" by downstream consumers, so we don't track that bit.
+  std::vector> backend_options_storage_;
+  LoadBackendOptionsMap backend_options_map_;
   bool share_memory_arenas_;
 
   ET_NODISCARD runtime::Error load_internal(
diff --git a/extension/module/targets.bzl b/extension/module/targets.bzl
index 6d60429bc51..fa80203831a 100644
--- a/extension/module/targets.bzl
+++ b/extension/module/targets.bzl
@@ -27,6 +27,8 @@ def define_common_targets():
                 "//executorch/extension/named_data_map:merged_data_map" + aten_suffix,
             ],
             exported_deps = [
+                "//executorch/runtime/backend:backend_options",
+                "//executorch/runtime/backend:backend_options_map",
                 "//executorch/runtime/executor:program_no_prim_ops" + aten_suffix,
             ],
         )
diff --git a/extension/module/test/module_test.cpp b/extension/module/test/module_test.cpp
index 7e1d657094c..71c55ae8751 100644
--- a/extension/module/test/module_test.cpp
+++ b/extension/module/test/module_test.cpp
@@ -9,7 +9,9 @@
 #include 
 
 #include 
+#include 
 #include 
+#include 
 
 #include 
 
@@ -50,6 +52,23 @@ TEST_F(ModuleTest, TestLoad) {
   EXPECT_TRUE(module.is_loaded());
 }
 
+TEST_F(ModuleTest, TestLoadMmapUseMadvise) {
+  Module module(model_path_, Module::LoadMode::MmapUseMadvise);
+
+  EXPECT_FALSE(module.is_loaded());
+  const auto error = module.load();
+  EXPECT_EQ(error, Error::Ok);
+  EXPECT_TRUE(module.is_loaded());
+
+  auto tensor = make_tensor_ptr({2, 2}, {1.f, 2.f, 3.f, 4.f});
+
+  const auto result = module.execute("forward", {tensor, tensor, 1.0});
+  EXPECT_EQ(result.error(), Error::Ok);
+
+  const auto expected = make_tensor_ptr({2, 2}, {2.f, 4.f, 6.f, 8.f});
+  EXPECT_TENSOR_CLOSE(result->at(0).toTensor(), *expected.get());
+}
+
 TEST_F(ModuleTest, TestLoadNonExistent) {
   Module module("/path/to/nonexistent/file.pte");
   const auto error = module.load();
@@ -646,6 +665,115 @@ TEST_F(ModuleTest, TestLoadWithEmptyLoadBackendOptionsMap) {
   EXPECT_TRUE(module.is_loaded());
 }
 
+TEST_F(ModuleTest, TestLoadWithBackendOptionsRollbackOnFailure) {
+  // Module pointed at a non-existent file so `load_internal` will fail.
+  Module module("/this/path/should/not/exist.pte");
+
+  {
+    // `bo1` lives only in this scope. The Module deep-copies the input,
+    // so dropping `bo1` is always safe regardless of whether the load
+    // succeeded, but on the failure path we additionally verify the
+    // Module did NOT install the input options (transactional contract).
+    LoadBackendOptionsMap bo1;
+    BackendOptions<2> opts;
+    opts.set_option("rollback_test", true);
+    ASSERT_EQ(bo1.set_options("RollbackBackend", opts.view()), Error::Ok);
+
+    const auto load_error = module.load(bo1);
+    EXPECT_NE(load_error, Error::Ok);
+    EXPECT_FALSE(module.is_loaded());
+  }
+  // `bo1` is destroyed. Module must remain in a usable state and a
+  // subsequent `load_method` should fail with the same load-time error
+  // (file not found) rather than crashing.
+  EXPECT_FALSE(module.is_loaded());
+  const auto method_error = module.load_method("forward");
+  EXPECT_NE(method_error, Error::Ok);
+  EXPECT_FALSE(module.is_method_loaded("forward"));
+}
+
+TEST_F(ModuleTest, TestLoadDeepCopiesBackendOptionsInputCanBeReleased) {
+  // Pin the deep-copy contract: the caller may release the input
+  // LoadBackendOptionsMap (and the BackendOption arrays its Spans
+  // referenced) immediately after `load()` returns. A subsequent
+  // `load_method` must use the Module-owned copy via the fallback path,
+  // not dereference the released input.
+  Module module(model_path_);
+
+  {
+    LoadBackendOptionsMap bo;
+    BackendOptions<2> opts;
+    opts.set_option("persist_test", true);
+    ASSERT_EQ(bo.set_options("PersistBackend", opts.view()), Error::Ok);
+
+    ASSERT_EQ(module.load(bo), Error::Ok);
+    // `bo` and `opts` go out of scope here; their storage is freed.
+  }
+
+  // load_method without explicit backend_options falls back to the
+  // Module's stored copy. With the old borrowed-pointer design this
+  // would have been a use-after-free; with deep-copy it is safe.
+  EXPECT_EQ(module.load_method("forward"), Error::Ok);
+  EXPECT_TRUE(module.is_method_loaded("forward"));
+
+  // Forward should still execute correctly using the Module-owned
+  // backend options.
+  auto tensor = make_tensor_ptr({2, 2}, {1.f, 2.f, 3.f, 4.f});
+  const auto result = module.execute("forward", {tensor, tensor, 1.0});
+  EXPECT_EQ(result.error(), Error::Ok);
+}
+
+TEST_F(ModuleTest, TestLoadStoresBackendOptionsForReadback) {
+  // Verify that Module deep-copies the input LoadBackendOptionsMap into
+  // its own storage so callers can both (a) release the input
+  // immediately and (b) read back exactly what was stored via the
+  // public `backend_options()` accessor.
+  Module module(model_path_);
+
+  // Default-constructed: no options stored yet.
+  EXPECT_EQ(module.backend_options().size(), 0u);
+
+  {
+    LoadBackendOptionsMap bo;
+    BackendOptions<2> opts;
+    opts.set_option("num_threads", 8);
+    opts.set_option("enable_profiling", true);
+    ASSERT_EQ(bo.set_options("MyBackend", opts.view()), Error::Ok);
+
+    ASSERT_EQ(module.load(bo), Error::Ok);
+    // `bo` and `opts` go out of scope here; their backing storage is
+    // freed. Anything we read back from `module.backend_options()` must
+    // therefore live in Module-owned storage.
+  }
+
+  const auto& stored = module.backend_options();
+  ASSERT_EQ(stored.size(), 1u);
+
+  const auto entry = stored.entry_at(0);
+  EXPECT_STREQ(entry.backend_id, "MyBackend");
+  ASSERT_EQ(entry.options.size(), 2u);
+
+  // Look up each option by key so the value assertions are direct and
+  // independent of insertion order.
+  const BackendOption* num_threads_opt = nullptr;
+  const BackendOption* enable_profiling_opt = nullptr;
+  for (const auto& opt : entry.options) {
+    if (std::strcmp(opt.key, "num_threads") == 0) {
+      num_threads_opt = &opt;
+    } else if (std::strcmp(opt.key, "enable_profiling") == 0) {
+      enable_profiling_opt = &opt;
+    }
+  }
+
+  ASSERT_NE(num_threads_opt, nullptr);
+  ASSERT_TRUE(std::holds_alternative(num_threads_opt->value));
+  EXPECT_EQ(std::get(num_threads_opt->value), 8);
+
+  ASSERT_NE(enable_profiling_opt, nullptr);
+  ASSERT_TRUE(std::holds_alternative(enable_profiling_opt->value));
+  EXPECT_TRUE(std::get(enable_profiling_opt->value));
+}
+
 TEST_F(ModuleTest, TestLoadBackendOptionsMapPersistedAcrossLoadMethod) {
   Module module(model_path_);
 
diff --git a/kernels/test/UnaryUfuncRealHBBF16ToFloatHBF16Test.cpp b/kernels/test/UnaryUfuncRealHBBF16ToFloatHBF16Test.cpp
index 74ad17df789..243e69b8c8d 100644
--- a/kernels/test/UnaryUfuncRealHBBF16ToFloatHBF16Test.cpp
+++ b/kernels/test/UnaryUfuncRealHBBF16ToFloatHBF16Test.cpp
@@ -7,6 +7,7 @@
  */
 
 #include 
+#include 
 
 namespace torch::executor::testing {
 void UnaryUfuncRealHBBF16ToFloatHBF16Test::test_bool_input() {
@@ -25,9 +26,9 @@ void UnaryUfuncRealHBBF16ToFloatHBF16Test::test_bool_input() {
 }
 
 void UnaryUfuncRealHBBF16ToFloatHBF16Test::test_mismatched_input_shapes_dies() {
-  if (get_supported_features()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched input shapes";
-  }
+  ET_SKIP_IF(
+      get_supported_features()->is_aten,
+      "ATen kernel can handle mismatched input shapes");
   TensorFactory tf;
 
   executorch::aten::Tensor a = tf.ones(/*sizes=*/{4});
@@ -122,9 +123,9 @@ void UnaryUfuncRealHBBF16ToFloatHBF16Test::
 
 void UnaryUfuncRealHBBF16ToFloatHBF16Test::
     test_all_real_input_float_output_unbound_dynamism_support() {
-  if (!get_supported_features()->is_aten) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !get_supported_features()->is_aten,
+      "Dynamic shape unbound not supported");
 #define TEST_ENTRY(ctype, dtype)            \
   test_floating_point_op_out<               \
       executorch::aten::ScalarType::dtype,  \
@@ -136,9 +137,9 @@ void UnaryUfuncRealHBBF16ToFloatHBF16Test::
 
 void UnaryUfuncRealHBBF16ToFloatHBF16Test::
     test_all_real_input_double_output_unbound_dynamism_support() {
-  if (!get_supported_features()->is_aten) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !get_supported_features()->is_aten,
+      "Dynamic shape unbound not supported");
 #define TEST_ENTRY(ctype, dtype)             \
   test_floating_point_op_out<                \
       executorch::aten::ScalarType::dtype,   \
diff --git a/kernels/test/UnaryUfuncRealHBBF16ToFloatHBF16Test.h b/kernels/test/UnaryUfuncRealHBBF16ToFloatHBF16Test.h
index d1e812ec2c2..fb2eada0a0f 100644
--- a/kernels/test/UnaryUfuncRealHBBF16ToFloatHBF16Test.h
+++ b/kernels/test/UnaryUfuncRealHBBF16ToFloatHBF16Test.h
@@ -10,6 +10,7 @@
 
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
diff --git a/kernels/test/op__clone_dim_order_test.cpp b/kernels/test/op__clone_dim_order_test.cpp
index f009ce1b195..ad4c7479185 100644
--- a/kernels/test/op__clone_dim_order_test.cpp
+++ b/kernels/test/op__clone_dim_order_test.cpp
@@ -11,6 +11,7 @@
 #include  // Declares the operator.
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -155,9 +156,9 @@ TEST_F(OpDimOrderCloneTest, AllDtypesSupported) {
 
 // Cloning with mismatched input and output tensor shapes should fail.
 TEST_F(OpDimOrderCloneTest, MismatchedSizesDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "Skipping: ATen kernel supports mismatched sizes.";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "Skipping: ATen kernel supports mismatched sizes.");
   TensorFactory tf;
   Tensor input = tf.make(/*sizes=*/{3, 1, 1, 2}, /*data=*/{1, 2, 3, 4, 5, 6});
   Tensor out = tf.zeros({3, 2, 1, 1});
@@ -178,10 +179,9 @@ TEST_F(OpDimOrderCloneTest, MismatchedSizesDie) {
 
 // Cloning with an unsupported memory format should fail.
 TEST_F(OpDimOrderCloneTest, MismatchedMemoryFormatDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP()
-        << "Skipping: ATen kernel supports non-contiguous memory formats.";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "Skipping: ATen kernel supports non-contiguous memory formats.");
   TensorFactory tf_in;
   TensorFactory tf_out;
   Tensor input =
@@ -210,10 +210,9 @@ TEST_F(OpDimOrderCloneTest, MismatchedMemoryFormatDies) {
 // Cloning with non‑blocking=true should fail because portable kernels only
 // support blocking.
 TEST_F(OpDimOrderCloneTest, MismatchedBlockingDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP()
-        << "Skipping: ATen kernel supports non-blocking data transfer.";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "Skipping: ATen kernel supports non-blocking data transfer.");
   TensorFactory tf;
   Tensor input = tf.make(/*sizes=*/{3, 1, 1, 2}, /*data=*/{1, 2, 3, 4, 5, 6});
   Tensor out = tf.zeros(/*sizes=*/{3, 1, 1, 2});
@@ -244,9 +243,9 @@ TEST_F(OpDimOrderCloneTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpDimOrderCloneTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Skipping: Dynamic shape unbound not supported.";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Skipping: Dynamic shape unbound not supported.");
   test_dynamic_shape(
       {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op__empty_dim_order_test.cpp b/kernels/test/op__empty_dim_order_test.cpp
index b3534948c8d..9a1b2ba62e1 100644
--- a/kernels/test/op__empty_dim_order_test.cpp
+++ b/kernels/test/op__empty_dim_order_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -147,9 +148,9 @@ TEST_F(OpEmptyDimOrderOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpEmptyDimOrderOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   TensorFactory tf;
 
   int64_t sizes[2] = {3, 2};
diff --git a/kernels/test/op__to_dim_order_copy_test.cpp b/kernels/test/op__to_dim_order_copy_test.cpp
index c2458ae6540..2d562defce7 100644
--- a/kernels/test/op__to_dim_order_copy_test.cpp
+++ b/kernels/test/op__to_dim_order_copy_test.cpp
@@ -14,6 +14,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -459,9 +460,9 @@ TEST_F(OpToDimOrderCopyTest, HardcodeFloatConvertInt) {
 }
 
 TEST_F(OpToDimOrderCopyTest, MismatchedSizesDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched sizes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched sizes");
   TensorFactory tf;
   Tensor input = tf.make(/*sizes=*/{3, 1, 1, 2}, /*data=*/{1, 2, 3, 4, 5, 6});
   Tensor out = tf.zeros({3, 2, 1, 1});
@@ -484,9 +485,9 @@ TEST_F(OpToDimOrderCopyTest, MismatchedSizesDie) {
 // should not be allowed. The function is expected death if using the illegal
 // memory format.
 TEST_F(OpToDimOrderCopyTest, MismatchedMemoryFormatDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle non contiguous memory formats";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle non contiguous memory formats");
   TensorFactory tf_in;
   TensorFactory tf_out;
   Tensor input =
@@ -514,9 +515,9 @@ TEST_F(OpToDimOrderCopyTest, MismatchedMemoryFormatDies) {
 
 // Only blocking data transfer supported
 TEST_F(OpToDimOrderCopyTest, MismatchedBlockingDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle non blocking data transfer";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle non blocking data transfer");
   TensorFactory tf;
   Tensor input = tf.make(/*sizes=*/{3, 1, 1, 2}, /*data=*/{1, 2, 3, 4, 5, 6});
   Tensor out = tf.zeros(/*sizes=*/{3, 1, 1, 2});
@@ -547,9 +548,9 @@ TEST_F(OpToDimOrderCopyTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpToDimOrderCopyTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   test_dynamic_shape(
       {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_add_test.cpp b/kernels/test/op_add_test.cpp
index 60faa4efb47..bf795d3c47a 100644
--- a/kernels/test/op_add_test.cpp
+++ b/kernels/test/op_add_test.cpp
@@ -10,6 +10,7 @@
 #include 
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -686,10 +687,9 @@ TEST_F(OpAddOutKernelTest, MismatchedNonBroadcastableInputShapesDies) {
 }
 
 TEST_F(OpAddOutKernelTest, MismatchedOutputShapesDies) {
-  if (SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP()
-        << "The current kernel supports implicitly resizing output tensor";
-  }
+  ET_SKIP_IF(
+      SupportedFeatures::get()->output_resize,
+      "The current kernel supports implicitly resizing output tensor");
 
   TensorFactory tf;
 
diff --git a/kernels/test/op_addmm_test.cpp b/kernels/test/op_addmm_test.cpp
index ff02d9c0a79..858c46323c2 100644
--- a/kernels/test/op_addmm_test.cpp
+++ b/kernels/test/op_addmm_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -42,11 +43,9 @@ class OpAddmmOutTest : public OperatorTest {
     TensorFactory tf;
 
     if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-      if (DTYPE == ScalarType::Half) {
-        GTEST_SKIP()
-            << "skip Half because torch::executor::aten::mm_out does not support Half";
-        return;
-      }
+      ET_SKIP_IF(
+          DTYPE == ScalarType::Half,
+          "skip Half because torch::executor::aten::mm_out does not support Half");
     }
 
     // matmul gives 4 * 2 * 3 = 24, α * 24 = 48, 48 + β * self = 51
@@ -205,9 +204,9 @@ TEST_F(OpAddmmOutTest, MismatchedDimensionSizeDies) {
   Tensor right_out = tf.ones({2, 2});
   Tensor wrong_out = tf.ones({2, 2, 3});
 
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched dimensions");
 
   ET_EXPECT_KERNEL_FAILURE(
       context_,
@@ -228,9 +227,9 @@ TEST_F(OpAddmmOutTest, WrongOutShapeDies) {
   Tensor right_out = tf.ones({10, 4});
   Tensor wrong_out = tf.ones({7, 5});
 
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle wrong out shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle wrong out shape");
 
   ET_EXPECT_KERNEL_FAILURE(
       context_, op_addmm_out(self, x, y, Scalar(1), Scalar(1), wrong_out));
diff --git a/kernels/test/op_amax_test.cpp b/kernels/test/op_amax_test.cpp
index 703495584ce..5c6fbc54987 100644
--- a/kernels/test/op_amax_test.cpp
+++ b/kernels/test/op_amax_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -260,9 +261,9 @@ void OpAmaxOutTest::test_amax_out_dtype() {
 }
 
 TEST_F(OpAmaxOutTest, InvalidDimensionListDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
 #define TEST_ENTRY(ctype, dtype) \
   test_amax_out_invalid_dimensions();
   ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
@@ -270,9 +271,9 @@ TEST_F(OpAmaxOutTest, InvalidDimensionListDies) {
 }
 
 TEST_F(OpAmaxOutTest, InvalidShapeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
 #define TEST_ENTRY(ctype, dtype) \
   test_amax_out_invalid_shape();
   ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
@@ -280,9 +281,9 @@ TEST_F(OpAmaxOutTest, InvalidShapeDies) {
 }
 
 TEST_F(OpAmaxOutTest, MismatchedDTypesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   TensorFactory tf_float;
   TensorFactory tf_int;
 
diff --git a/kernels/test/op_amin_test.cpp b/kernels/test/op_amin_test.cpp
index dc41676c03a..564f444e268 100644
--- a/kernels/test/op_amin_test.cpp
+++ b/kernels/test/op_amin_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -259,9 +260,9 @@ void OpAminOutTest::test_amin_out_dtype() {
 }
 
 TEST_F(OpAminOutTest, InvalidDimensionListDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
 #define TEST_ENTRY(ctype, dtype) \
   test_amin_out_invalid_dimensions();
   ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
@@ -269,9 +270,9 @@ TEST_F(OpAminOutTest, InvalidDimensionListDies) {
 }
 
 TEST_F(OpAminOutTest, InvalidShapeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
 #define TEST_ENTRY(ctype, dtype) \
   test_amin_out_invalid_shape();
   ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
@@ -279,9 +280,9 @@ TEST_F(OpAminOutTest, InvalidShapeDies) {
 }
 
 TEST_F(OpAminOutTest, MismatchedDTypesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   TensorFactory tf_float;
   TensorFactory tf_int;
 
diff --git a/kernels/test/op_any_test.cpp b/kernels/test/op_any_test.cpp
index 1853b96ab7d..cb831bffa8a 100644
--- a/kernels/test/op_any_test.cpp
+++ b/kernels/test/op_any_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -99,9 +100,9 @@ class OpAnyOutTest : public OperatorTest {
 };
 
 TEST_F(OpAnyOutTest, MismatchedDimensionsDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched dimensions");
   TensorFactory tff;
   const std::vector size{2, 2};
 
diff --git a/kernels/test/op_arange_test.cpp b/kernels/test/op_arange_test.cpp
index e7b9ae7c9ea..82bc46a1858 100644
--- a/kernels/test/op_arange_test.cpp
+++ b/kernels/test/op_arange_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -114,9 +115,9 @@ TEST_F(OpArangeOutTest, FloatNumberNotEqualIntSupport) {
 }
 
 TEST_F(OpArangeOutTest, OutDimUnsupportedDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched out dim";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched out dim");
   TensorFactory tf;
 
   Scalar end = Scalar(5);
@@ -150,9 +151,9 @@ TEST_F(OpArangeOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpArangeOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "Dynamic Unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "Dynamic Unbound not supported");
   TensorFactory tf;
 
   Tensor expected_result = tf.make({5}, {0, 1, 2, 3, 4});
@@ -196,9 +197,9 @@ TEST_F(OpArangeStartOutTest, FloatNumberNotEqualIntSupport) {
 }
 
 TEST_F(OpArangeStartOutTest, OutDimUnsupportedDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched out dim";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched out dim");
   TensorFactory tf;
 
   Scalar start = Scalar(0);
@@ -235,9 +236,9 @@ TEST_F(OpArangeStartOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpArangeStartOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "Dynamic Unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "Dynamic Unbound not supported");
   TensorFactory tf;
 
   Tensor expected_result = tf.make({5}, {0, 1, 2, 3, 4});
diff --git a/kernels/test/op_as_strided_copy_test.cpp b/kernels/test/op_as_strided_copy_test.cpp
index cb0191c69a8..e025c7afaf6 100644
--- a/kernels/test/op_as_strided_copy_test.cpp
+++ b/kernels/test/op_as_strided_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -197,9 +198,9 @@ TEST_F(OpAsStridedCopyOutTest, AllScalarInputOutputSupport) {
 }
 
 TEST_F(OpAsStridedCopyOutTest, InvalidParametersDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle invalid parameter";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle invalid parameter");
 #define TEST_ENTRY(ctype, dtype) \
   test_as_strided_copy_out_invalid_parameters();
   ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
@@ -326,9 +327,9 @@ TEST_F(OpAsStridedCopyOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpAsStridedCopyOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   /* %python
   out_args = "{1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op) */
diff --git a/kernels/test/op_bitwise_not_test.cpp b/kernels/test/op_bitwise_not_test.cpp
index 702486f0d2a..ddfb5b9c0d1 100644
--- a/kernels/test/op_bitwise_not_test.cpp
+++ b/kernels/test/op_bitwise_not_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -96,9 +97,9 @@ TEST_F(OpBitwiseNotOutTest, BoolInputOutputSupport) {
 
 // Mismatched shape tests.
 TEST_F(OpBitwiseNotOutTest, MismatchedShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf;
 
   Tensor a = tf.ones(/*sizes=*/{4});
diff --git a/kernels/test/op_bmm_test.cpp b/kernels/test/op_bmm_test.cpp
index c870c412035..eeb5d8d8458 100644
--- a/kernels/test/op_bmm_test.cpp
+++ b/kernels/test/op_bmm_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -239,9 +240,9 @@ TEST_F(OpBmmOutTest, MismatchedDimensionsDies) {
 }
 
 TEST_F(OpBmmOutTest, MismatchedDimensionSizeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched dimension size";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched dimension size");
   TensorFactory tf;
 
   Tensor x = tf.ones({2, 10, 3});
@@ -261,9 +262,9 @@ TEST_F(OpBmmOutTest, MismatchedDimensionSizeDies) {
 }
 
 TEST_F(OpBmmOutTest, WrongOutShapeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle wrong out shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle wrong out shape");
   TensorFactory tf;
 
   Tensor x = tf.ones({2, 10, 3});
diff --git a/kernels/test/op_cat_test.cpp b/kernels/test/op_cat_test.cpp
index 4ea131452c7..d3bda1e8abd 100644
--- a/kernels/test/op_cat_test.cpp
+++ b/kernels/test/op_cat_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -178,9 +179,9 @@ TEST_F(OpCatOutTest, SmokeDim1) {
 }
 
 TEST_F(OpCatOutTest, SixteenBitFloatSupport) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "Test Half/BF16 support only for ExecuTorch mode";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "Test Half/BF16 support only for ExecuTorch mode");
   test_16bit_dtype();
   test_16bit_dtype();
 }
@@ -246,9 +247,9 @@ TEST_F(OpCatOutTest, AllDtypesSupported) {
 }
 
 TEST_F(OpCatOutTest, EmptyInputTensorShapeIgnored) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel doesn't ignore empty input tensor shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel doesn't ignore empty input tensor shape");
   TensorFactory tf;
 
   // An empty tensor with a shape totally different from the non-empty inputs.
@@ -330,9 +331,9 @@ TEST_F(OpCatOutTest, MismatchedDtypesDies) {
 }
 
 TEST_F(OpCatOutTest, MismatchedDimensionsDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched dimensions");
   TensorFactory tf;
   Tensor out = tf.zeros({2, 2});
 
@@ -346,9 +347,9 @@ TEST_F(OpCatOutTest, MismatchedDimensionsDies) {
 }
 
 TEST_F(OpCatOutTest, MismatchedDimensionSizeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched dimension size";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched dimension size");
   TensorFactory tf;
   Tensor out = tf.zeros({2, 2});
 
@@ -363,9 +364,9 @@ TEST_F(OpCatOutTest, MismatchedDimensionSizeDies) {
 }
 
 TEST_F(OpCatOutTest, WrongOutShapeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle wrong out shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle wrong out shape");
   TensorFactory tf;
 
   // Should be {4, 3} to match the inputs when calling cat() with dim 0.
@@ -440,9 +441,9 @@ TEST_F(OpCatOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpCatOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   /* %python
   out_args = "{1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op_tensor_list_in) */
diff --git a/kernels/test/op_clamp_test.cpp b/kernels/test/op_clamp_test.cpp
index aeb44f1d7ab..2962cb98963 100644
--- a/kernels/test/op_clamp_test.cpp
+++ b/kernels/test/op_clamp_test.cpp
@@ -12,6 +12,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -360,9 +361,9 @@ TEST_F(OpClampOutTest, DoubleTensors) {
 //
 
 TEST_F(OpClampOutTest, ByteTensorNegativeClampDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle negative clamp on byte tensor";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle negative clamp on byte tensor");
   // Cannot be represented by a uint8_t.
   expect_bad_clamp_value_dies(-1);
 }
diff --git a/kernels/test/op_clone_test.cpp b/kernels/test/op_clone_test.cpp
index 57a8aed2d6c..aec1c473f18 100644
--- a/kernels/test/op_clone_test.cpp
+++ b/kernels/test/op_clone_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -73,9 +74,9 @@ class OpCloneTest : public OperatorTest {
 
 // regular test for clone.out
 TEST_F(OpCloneTest, AllDtypesSupported) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
 #define TEST_ENTRY(ctype, dtype) test_dtype();
   ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
 #undef TEST_ENTRY
@@ -88,9 +89,9 @@ TEST_F(OpCloneTest, EmptyInputSupported) {
 }
 
 TEST_F(OpCloneTest, MismatchedSizesDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched sizes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched sizes");
   TensorFactory tf;
   Tensor input = tf.make(/*sizes=*/{3, 1, 1, 2}, /*data=*/{1, 2, 3, 4, 5, 6});
   Tensor out = tf.zeros({3, 2, 1, 1});
@@ -114,9 +115,9 @@ TEST_F(OpCloneTest, MismatchedTypesDie) {
 // MemoryFormat::Contiguous should not be allowed. The function is expected
 // depth if using the illegal memory format.
 TEST_F(OpCloneTest, MismatchedMemoryFormatDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle non contiguous memory formats";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle non contiguous memory formats");
   TensorFactory tf_in;
   TensorFactory tf_out;
   Tensor input =
diff --git a/kernels/test/op_constant_pad_nd_test.cpp b/kernels/test/op_constant_pad_nd_test.cpp
index 7f44068d9cb..7bd908e0ecb 100644
--- a/kernels/test/op_constant_pad_nd_test.cpp
+++ b/kernels/test/op_constant_pad_nd_test.cpp
@@ -10,6 +10,7 @@
 #include 
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -463,9 +464,9 @@ TEST_F(OpConstantPadNDOutTest, TooManyPaddingElementsFail) {
 }
 
 TEST_F(OpConstantPadNDOutTest, IncorrectOutputShapeFail) {
-  if (SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle reshape output";
-  }
+  ET_SKIP_IF(
+      SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle reshape output");
 
   TensorFactory tf;
 
diff --git a/kernels/test/op_convolution_test.cpp b/kernels/test/op_convolution_test.cpp
index 1e0e406af44..89a46038fae 100644
--- a/kernels/test/op_convolution_test.cpp
+++ b/kernels/test/op_convolution_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -466,9 +467,9 @@ TEST_F(OpConvOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpConvOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   test_dynamic_shape(
       {1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_copy_test.cpp b/kernels/test/op_copy_test.cpp
index 97fd7e7e6c0..b8c15a147f3 100644
--- a/kernels/test/op_copy_test.cpp
+++ b/kernels/test/op_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -205,9 +206,9 @@ TEST_F(OpCopyTest, ResizeOutDie) {
 #endif
 
 TEST_F(OpCopyTest, MismatchedSizesDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched sizes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched sizes");
   TensorFactory tf;
   Tensor self = tf.make(/*sizes=*/{3, 1, 1, 2}, /*data=*/{1, 2, 3, 4, 5, 6});
   Tensor src = tf.make(/*sizes=*/{3, 1, 1, 2}, /*data=*/{1, 2, 3, 4, 5, 6});
@@ -230,9 +231,9 @@ TEST_F(OpCopyTest, MismatchedSrcOutTypesDie) {
 // MemoryFormat::Contiguous should not be allowed. The function is expected
 // depth if using the illegal memory format.
 TEST_F(OpCopyTest, BlockingDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle non-contiguous memory formats";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle non-contiguous memory formats");
   TensorFactory tf_in;
   TensorFactory tf_out;
   Tensor self = tf_in.make(/*sizes=*/{3, 1, 1, 2}, /*data=*/{1, 2, 3, 4, 5, 6});
@@ -248,17 +249,17 @@ TEST_F(OpCopyTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpCopyTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
 }
 
 TEST_F(OpCopyTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_cumsum_test.cpp b/kernels/test/op_cumsum_test.cpp
index 720f7bd98e9..48605024ff1 100644
--- a/kernels/test/op_cumsum_test.cpp
+++ b/kernels/test/op_cumsum_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -104,9 +105,9 @@ class OpCumSumOutTest : public OperatorTest {
 };
 
 TEST_F(OpCumSumOutTest, MismatchedDimensionsDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched dimensions");
   TensorFactory tff;
 
   Tensor in = tff.make({1, 3}, {0, 1, 2});
diff --git a/kernels/test/op_detach_copy_test.cpp b/kernels/test/op_detach_copy_test.cpp
index fba497c75ab..42d287ec4c1 100644
--- a/kernels/test/op_detach_copy_test.cpp
+++ b/kernels/test/op_detach_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -88,9 +89,9 @@ TEST_F(OpDetachCopyOutTest, AllScalarInputOutputSupport) {
 
 // Mismatched shape tests.
 TEST_F(OpDetachCopyOutTest, MismatchedShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
 #define TEST_ENTRY(ctype, dtype) \
   test_detach_copy_out_invalid_shape();
   ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
diff --git a/kernels/test/op_div_test.cpp b/kernels/test/op_div_test.cpp
index 84d33fa2757..d313e85a338 100644
--- a/kernels/test/op_div_test.cpp
+++ b/kernels/test/op_div_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -427,9 +428,9 @@ TEST_F(OpDivOutTest, BroadcastDimSizeMissingBA) {
 //
 
 TEST_F(OpDivOutTest, MismatchedShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_float;
 
diff --git a/kernels/test/op_embedding_test.cpp b/kernels/test/op_embedding_test.cpp
index abee1be43e0..1f0e172a0e8 100644
--- a/kernels/test/op_embedding_test.cpp
+++ b/kernels/test/op_embedding_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -218,9 +219,9 @@ TEST_F(OpEmbeddingOutTest, WeightWrongDimensionsDies) {
 }
 
 TEST_F(OpEmbeddingOutTest, WrongOutShapeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle wrong out shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle wrong out shape");
   TensorFactory tff;
   // clang-format off
   Tensor weight = tff.make(
@@ -455,9 +456,9 @@ TEST_F(OpEmbeddingOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpEmbeddingOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   test_dynamic_shape(
       {1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_empty_test.cpp b/kernels/test/op_empty_test.cpp
index 23173b1feae..e7bbf1f9dfb 100644
--- a/kernels/test/op_empty_test.cpp
+++ b/kernels/test/op_empty_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -78,9 +79,9 @@ TEST_F(OpEmptyOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpEmptyOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   TensorFactory tf;
 
   int64_t sizes[2] = {3, 2};
diff --git a/kernels/test/op_eq_test.cpp b/kernels/test/op_eq_test.cpp
index 539fb172f85..24cb00772a8 100644
--- a/kernels/test/op_eq_test.cpp
+++ b/kernels/test/op_eq_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -82,9 +83,9 @@ TEST_F(OpEqScalarOutTest, BoolInputDtype) {
 
 // Mismatched shape tests.
 TEST_F(OpEqScalarOutTest, MismatchedShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
@@ -96,9 +97,9 @@ TEST_F(OpEqScalarOutTest, MismatchedShapesDies) {
 }
 
 TEST_F(OpEqScalarOutTest, AllRealOutputDTypes) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle non-bool output dtype";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle non-bool output dtype");
 #define TEST_ENTRY(ctype, dtype) test_eq_all_output_dtypes();
   ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
@@ -159,9 +160,9 @@ TEST_F(OpEqScalarOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpEqScalarOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   /* %python
   out_args = "{1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op_out_dtype) */
diff --git a/kernels/test/op_expand_copy_test.cpp b/kernels/test/op_expand_copy_test.cpp
index 86d3858c830..790ed8e82d5 100644
--- a/kernels/test/op_expand_copy_test.cpp
+++ b/kernels/test/op_expand_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -216,9 +217,9 @@ TEST_F(OpExpandOutTest, BadOutDataTypeGoodShapeDeath) {
 }
 
 TEST_F(OpExpandOutTest, BadOutShapeGoodDataTypeDeath) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle this";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle this");
   TensorFactory tf;
   Tensor a = tf.make(/*sizes*/ {1, 2}, /*data=*/{42, 96});
   Tensor out = tf.ones({2, 6, 4});
@@ -314,9 +315,9 @@ TEST_F(OpExpandOutTest, ResizedOutput) {
 #endif
 
 TEST_F(OpExpandOutTest, ImplicitTrue) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle this";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle this");
   TensorFactory tf;
   Tensor a = tf.ones({2, 2});
   Tensor out = tf.zeros({2, 2});
@@ -378,9 +379,9 @@ TEST_F(OpExpandOutTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpExpandOutTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{10, 10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
   %rewrite(unary_op) */
@@ -418,9 +419,9 @@ TEST_F(OpExpandOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpExpandOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op) */
diff --git a/kernels/test/op_fft_c2r_test.cpp b/kernels/test/op_fft_c2r_test.cpp
index 58c8a216e42..dd9ae7b3e2c 100644
--- a/kernels/test/op_fft_c2r_test.cpp
+++ b/kernels/test/op_fft_c2r_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -160,10 +161,9 @@ TEST_F(OpFftC2rOutTest, MultipleDims) {
 }
 
 TEST_F(OpFftC2rOutTest, InvalidNorm) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen MKL path does not validate norm";
-    return;
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen MKL path does not validate norm");
   auto invalid_norm = [this](int64_t norm) {
     test_dtype(norm);
   };
@@ -174,10 +174,9 @@ TEST_F(OpFftC2rOutTest, InvalidNorm) {
 }
 
 TEST_F(OpFftC2rOutTest, InvalidDim) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen fails UBSAN";
-    return;
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen fails UBSAN");
   auto negative_dim = [this]() {
     test_dtype(0, -1);
     test_dtype(0, 3);
diff --git a/kernels/test/op_fft_r2c_test.cpp b/kernels/test/op_fft_r2c_test.cpp
index 8730053bdc0..6db4ea617c4 100644
--- a/kernels/test/op_fft_r2c_test.cpp
+++ b/kernels/test/op_fft_r2c_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -131,10 +132,9 @@ TEST_F(OpFftR2cOutTest, MultipleDims) {
 }
 
 TEST_F(OpFftR2cOutTest, InvalidNorm) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen MKL path does not validate norm";
-    return;
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen MKL path does not validate norm");
   auto invalid_norm = [this](int64_t norm) {
     test_dtype(norm);
   };
@@ -145,10 +145,9 @@ TEST_F(OpFftR2cOutTest, InvalidNorm) {
 }
 
 TEST_F(OpFftR2cOutTest, InvalidDim) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen fails UBSAN";
-    return;
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen fails UBSAN");
   auto negative_dim = [this]() {
     test_dtype(0, -1);
     test_dtype(0, 3);
@@ -159,10 +158,9 @@ TEST_F(OpFftR2cOutTest, InvalidDim) {
 
 // TODO: support this and patch test accordingly!
 TEST_F(OpFftR2cOutTest, TwoSidedIsNotSupported) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen supports two-sided";
-    return;
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen supports two-sided");
   auto twosided = [this]() {
     test_dtype(
         0, 1, false);
diff --git a/kernels/test/op_fill_test.cpp b/kernels/test/op_fill_test.cpp
index c1c50206152..4ad5af749b8 100644
--- a/kernels/test/op_fill_test.cpp
+++ b/kernels/test/op_fill_test.cpp
@@ -10,6 +10,7 @@
 #include 
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -142,9 +143,9 @@ TEST_F(OpFillTest, MismatchedOtherPropertiesDies) {
 
 TEST_F(OpFillTest, MismatchedOutputShapesDies) {
   // Skip ATen test since it supports `self` and `out` having different shapes.
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched output shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched output shape");
 
   TensorFactory tf;
 
diff --git a/kernels/test/op_floor_divide_test.cpp b/kernels/test/op_floor_divide_test.cpp
index 166f7fdd4f9..7939ea775af 100644
--- a/kernels/test/op_floor_divide_test.cpp
+++ b/kernels/test/op_floor_divide_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -156,9 +157,9 @@ TEST_F(OpFloorDivideTest, MismatchedInputShapesDies) {
 }
 
 TEST_F(OpFloorDivideTest, MismatchedOutputShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched output shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched output shape");
   TensorFactory tf;
 
   const std::vector sizes = {2, 2};
diff --git a/kernels/test/op_full_like_test.cpp b/kernels/test/op_full_like_test.cpp
index 23ac4e685f9..4f9ab255708 100644
--- a/kernels/test/op_full_like_test.cpp
+++ b/kernels/test/op_full_like_test.cpp
@@ -10,6 +10,7 @@
 #include 
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -105,9 +106,9 @@ TEST_F(OpFullLikeTest, AllDtypeOutputPasses) {
 }
 
 TEST_F(OpFullLikeTest, MismatchedShapeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
 #define TEST_ENTRY(ctype, dtype) \
   test_full_like_out_mismatched_shape();
   ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
diff --git a/kernels/test/op_gather_test.cpp b/kernels/test/op_gather_test.cpp
index ff67d9b8fd9..768bc422ad8 100644
--- a/kernels/test/op_gather_test.cpp
+++ b/kernels/test/op_gather_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -268,9 +269,9 @@ TEST_F(OpGatherOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpGatherOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_ge_test.cpp b/kernels/test/op_ge_test.cpp
index 4fd0aa515b3..07812e99993 100644
--- a/kernels/test/op_ge_test.cpp
+++ b/kernels/test/op_ge_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -92,9 +93,9 @@ TEST_F(OpGeScalarOutTest, BoolInputDtype) {
 
 // Mismatched shape tests.
 TEST_F(OpGeScalarOutTest, MismatchedInOutShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
@@ -135,9 +136,9 @@ TEST_F(OpGeTensorOutTest, AllDtypesSupported) {
 }
 
 TEST_F(OpGeTensorOutTest, MismatchedInShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
@@ -149,9 +150,9 @@ TEST_F(OpGeTensorOutTest, MismatchedInShapesDies) {
 }
 
 TEST_F(OpGeTensorOutTest, MismatchedInOutShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
diff --git a/kernels/test/op_gelu_test.cpp b/kernels/test/op_gelu_test.cpp
index 9303b034ca2..b2e8902f4e3 100644
--- a/kernels/test/op_gelu_test.cpp
+++ b/kernels/test/op_gelu_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -79,9 +80,7 @@ TEST_F(OpGeluTest, BFloat16Tensors) {
 }
 
 TEST_F(OpGeluTest, DoubleTensors) {
-  if (!SupportedFeatures::get()->op_gelu_dtype_double) {
-    GTEST_SKIP();
-  }
+  ET_SKIP_IF(!SupportedFeatures::get()->op_gelu_dtype_double, "");
 
   test_gelu_execution();
 }
diff --git a/kernels/test/op_glu_test.cpp b/kernels/test/op_glu_test.cpp
index 9bee3a6a5a2..f5aaa54a49f 100644
--- a/kernels/test/op_glu_test.cpp
+++ b/kernels/test/op_glu_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -171,9 +172,9 @@ TEST_F(OpGluOutTest, InfinityAndNANTest) {
 }
 
 TEST_F(OpGluOutTest, MismatchedShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
 #define TEST_ENTRY(ctype, dtype) \
   test_glu_out_mismatched_shape();
   ET_FORALL_FLOAT_TYPES(TEST_ENTRY);
diff --git a/kernels/test/op_gt_test.cpp b/kernels/test/op_gt_test.cpp
index 028e7d16878..649fc17722e 100644
--- a/kernels/test/op_gt_test.cpp
+++ b/kernels/test/op_gt_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -92,9 +93,9 @@ TEST_F(OpGtScalarOutTest, BoolInputDtype) {
 
 // Mismatched shape tests.
 TEST_F(OpGtScalarOutTest, MismatchedInOutShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
@@ -135,9 +136,9 @@ TEST_F(OpGtTensorOutTest, AllDtypesSupported) {
 }
 
 TEST_F(OpGtTensorOutTest, MismatchedInShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
@@ -149,9 +150,9 @@ TEST_F(OpGtTensorOutTest, MismatchedInShapesDies) {
 }
 
 TEST_F(OpGtTensorOutTest, MismatchedInOutShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
diff --git a/kernels/test/op_index_put_test.cpp b/kernels/test/op_index_put_test.cpp
index f1021d9ad61..1c3ef58eb6b 100644
--- a/kernels/test/op_index_put_test.cpp
+++ b/kernels/test/op_index_put_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -997,17 +998,17 @@ TEST_F(OpIndexPutOutTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpIndexPutOutTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {10, 10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
 }
 
 TEST_F(OpIndexPutOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_index_select_test.cpp b/kernels/test/op_index_select_test.cpp
index 33adf166682..5b4ccf33da6 100644
--- a/kernels/test/op_index_select_test.cpp
+++ b/kernels/test/op_index_select_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -282,9 +283,9 @@ TEST_F(OpIndexSelectOutTest, AllDtypesSupported) {
 // In this test we are gonnna find if our select function support non-empty
 // tensor input and empty-size tensor output.
 TEST_F(OpIndexSelectOutTest, NonEmptyInputEmptyOutputWithMismatchDimDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle out with mismatched dimensions");
   TensorFactory tf;
   TensorFactory tfl;
   Tensor x = tf.make({10}, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9});
@@ -356,9 +357,9 @@ TEST_F(OpIndexSelectOutTest, MismatchedDtypesDies) {
 }
 
 TEST_F(OpIndexSelectOutTest, OutMatchNumelLackDimAtEndDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle out with mismatched dimensions");
   TensorFactory tf;
   TensorFactory tfl;
 
@@ -374,9 +375,9 @@ TEST_F(OpIndexSelectOutTest, OutMatchNumelLackDimAtEndDies) {
 }
 
 TEST_F(OpIndexSelectOutTest, OutMatchNumelExtraDimAtFrontDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle out with mismatched dimensions");
   TensorFactory tf;
   TensorFactory tfl;
 
@@ -392,9 +393,9 @@ TEST_F(OpIndexSelectOutTest, OutMatchNumelExtraDimAtFrontDies) {
 }
 
 TEST_F(OpIndexSelectOutTest, OutSizeMismatchDimDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle out with mismatched dimensions");
   TensorFactory tf;
   TensorFactory tfl;
 
@@ -505,17 +506,17 @@ TEST_F(OpIndexSelectOutTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpIndexSelectOutTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {10, 10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
 }
 
 TEST_F(OpIndexSelectOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_index_test.cpp b/kernels/test/op_index_test.cpp
index 8816d0a8d3f..48d0b0b80b1 100644
--- a/kernels/test/op_index_test.cpp
+++ b/kernels/test/op_index_test.cpp
@@ -10,6 +10,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -433,9 +434,9 @@ TEST_F(OpIndexTensorOutTest, IndicesWithNullTensorsSupported) {
 }
 
 TEST_F(OpIndexTensorOutTest, IndicesWithOnlyNullTensorsSupported) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   test_indices_with_only_null_tensors_enumerate_in_types();
 }
 
@@ -452,9 +453,9 @@ TEST_F(OpIndexTensorOutTest, TooManyNullIndices) {
 }
 
 TEST_F(OpIndexTensorOutTest, EmptyIndicesSupported) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   TensorFactory tf;
 
   // Using empty tensors as input.
@@ -622,9 +623,7 @@ TEST_F(OpIndexTensorOutTest, InvalidIndicesShapesDies) {
 }
 
 TEST_F(OpIndexTensorOutTest, InvalidIndicesShapeDies2) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "";
-  }
+  ET_SKIP_IF(torch::executor::testing::SupportedFeatures::get()->is_aten, "");
   TensorFactory tf;
   TensorFactory tfl;
 
diff --git a/kernels/test/op_le_test.cpp b/kernels/test/op_le_test.cpp
index 1baf098f9dd..3bc14a56548 100644
--- a/kernels/test/op_le_test.cpp
+++ b/kernels/test/op_le_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -92,9 +93,9 @@ TEST_F(OpLeScalarOutTest, BoolInputDtype) {
 
 // Mismatched shape tests.
 TEST_F(OpLeScalarOutTest, MismatchedInOutShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
@@ -135,9 +136,9 @@ TEST_F(OpLeTensorOutTest, AllDtypesSupported) {
 }
 
 TEST_F(OpLeTensorOutTest, MismatchedInShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
@@ -149,9 +150,9 @@ TEST_F(OpLeTensorOutTest, MismatchedInShapesDies) {
 }
 
 TEST_F(OpLeTensorOutTest, MismatchedInOutShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
diff --git a/kernels/test/op_lift_fresh_copy_test.cpp b/kernels/test/op_lift_fresh_copy_test.cpp
index 215ad4e05c6..31fe6d3402c 100644
--- a/kernels/test/op_lift_fresh_copy_test.cpp
+++ b/kernels/test/op_lift_fresh_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -68,9 +69,9 @@ TEST_F(OpLiftFreshCopyTest, EmptyInputSupported) {
 }
 
 TEST_F(OpLiftFreshCopyTest, MismatchedSizesDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched sizes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched sizes");
   TensorFactory tf;
   Tensor self = tf.make(/*sizes=*/{3, 1, 1, 2}, /*data=*/{1, 2, 3, 4, 5, 6});
   Tensor out = tf.zeros({3, 2, 1, 1});
diff --git a/kernels/test/op_linear_test.cpp b/kernels/test/op_linear_test.cpp
index 9b0ba782271..406b3b32c11 100644
--- a/kernels/test/op_linear_test.cpp
+++ b/kernels/test/op_linear_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -45,11 +46,9 @@ class OpLinearOutTest : public OperatorTest {
     TensorFactory tf;
 
     if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-      if (DTYPE == ScalarType::Half) {
-        GTEST_SKIP()
-            << "skip Half because torch::executor::aten::mm_out does not support Half";
-        return;
-      }
+      ET_SKIP_IF(
+          DTYPE == ScalarType::Half,
+          "skip Half because torch::executor::aten::mm_out does not support Half");
     }
 
     // matmul gives 19 * 2 * 3 = 114
@@ -218,9 +217,9 @@ TEST_F(OpLinearOutTest, MismatchedDimensionsDies) {
 }
 
 TEST_F(OpLinearOutTest, MismatchedDimensionSizeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched dimension size";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched dimension size");
   TensorFactory tf;
   Tensor x = tf.full({2, 2}, 3);
 
@@ -237,9 +236,9 @@ TEST_F(OpLinearOutTest, MismatchedDimensionSizeDies) {
 }
 
 TEST_F(OpLinearOutTest, WrongOutShapeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle wrong out shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle wrong out shape");
   TensorFactory tf;
   Tensor x = tf.ones({10, 3});
 
diff --git a/kernels/test/op_log_softmax_test.cpp b/kernels/test/op_log_softmax_test.cpp
index 84255b8a29c..88a8660faf6 100644
--- a/kernels/test/op_log_softmax_test.cpp
+++ b/kernels/test/op_log_softmax_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -145,9 +146,9 @@ TEST_F(OpLogSoftmaxOutTest, Smoke) {
 }
 
 TEST_F(OpLogSoftmaxOutTest, AllDtypesSupported) {
-  if (!SupportedFeatures::get()->op_log_softmax_dtype_double) {
-    GTEST_SKIP() << "This kernel does not support dtype double";
-  }
+  ET_SKIP_IF(
+      !SupportedFeatures::get()->op_log_softmax_dtype_double,
+      "This kernel does not support dtype double");
 
 #define TEST_ENTRY(ctype, dtype) test_dtype();
   ET_FORALL_FLOATHBF16_TYPES(TEST_ENTRY)
@@ -159,9 +160,9 @@ TEST_F(OpLogSoftmaxOutTest, NonContiguous) {
 }
 
 TEST_F(OpLogSoftmaxOutTest, MismatchedDimensionsDies) {
-  if (SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen currently supports mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      SupportedFeatures::get()->is_aten,
+      "ATen currently supports mismatched dimensions");
 
   TensorFactory tff;
 
@@ -177,9 +178,9 @@ TEST_F(OpLogSoftmaxOutTest, MismatchedDimensionsDies) {
 }
 
 TEST_F(OpLogSoftmaxOutTest, MismatchedDimensionSizeDies) {
-  if (SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen currently supports mismatched dimension size";
-  }
+  ET_SKIP_IF(
+      SupportedFeatures::get()->is_aten,
+      "ATen currently supports mismatched dimension size");
 
   TensorFactory tf;
 
@@ -194,13 +195,13 @@ TEST_F(OpLogSoftmaxOutTest, MismatchedDimensionSizeDies) {
 }
 
 TEST_F(OpLogSoftmaxOutTest, TestWithLargeNumber) {
-  if (!SupportedFeatures::get()->op_log_softmax_dtype_double) {
-    GTEST_SKIP() << "This kernel does not support dtype double";
-  }
+  ET_SKIP_IF(
+      !SupportedFeatures::get()->op_log_softmax_dtype_double,
+      "This kernel does not support dtype double");
 
-  if (SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen does not support mixing float and double";
-  }
+  ET_SKIP_IF(
+      SupportedFeatures::get()->is_aten,
+      "ATen does not support mixing float and double");
 
   TensorFactory tf;
 
@@ -229,13 +230,12 @@ TEST_F(OpLogSoftmaxOutTest, TestWithLargeNumber) {
 }
 
 TEST_F(OpLogSoftmaxOutTest, NegativeDim) {
-  if (!SupportedFeatures::get()->op_log_softmax_dtype_double) {
-    GTEST_SKIP() << "This kernel does not support dtype double";
-  }
+  ET_SKIP_IF(
+      !SupportedFeatures::get()->op_log_softmax_dtype_double,
+      "This kernel does not support dtype double");
 
-  if (SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen does not support negative dim";
-  }
+  ET_SKIP_IF(
+      SupportedFeatures::get()->is_aten, "ATen does not support negative dim");
 
   TensorFactory tf;
 
@@ -472,7 +472,7 @@ TEST_F(OpLogSoftmaxOutTest, DoubleCase) {
   if (!SupportedFeatures::get()->op_log_softmax_dtype_double) {
     // For optimized kernels, we expect the call above to fail gracefully
     expect_failure();
-    GTEST_SKIP() << "This kernel does not support dtype double";
+    ET_SKIP_IF(true, "This kernel does not support dtype double");
   }
 
   // Verify output dimensions
diff --git a/kernels/test/op_logical_not_test.cpp b/kernels/test/op_logical_not_test.cpp
index d06a3dcefea..c9afd6d39ae 100644
--- a/kernels/test/op_logical_not_test.cpp
+++ b/kernels/test/op_logical_not_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -104,9 +105,9 @@ class OpLogicalNotOutTest : public OperatorTest {
 };
 
 TEST_F(OpLogicalNotOutTest, MismatchedDimensionsDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched dimensions");
   TensorFactory tff;
   const std::vector size{2, 2};
 
diff --git a/kernels/test/op_logit_test.cpp b/kernels/test/op_logit_test.cpp
index 0056e984bb7..effdc0d0377 100644
--- a/kernels/test/op_logit_test.cpp
+++ b/kernels/test/op_logit_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -134,9 +135,9 @@ TEST_F(OpLogitOutTest, AllRealInputDoubleOutputSupportEpsSet) {
 
 // Mismatched shape tests.
 TEST_F(OpLogitOutTest, MismatchedShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf;
   TensorFactory tf_out;
 
diff --git a/kernels/test/op_lt_test.cpp b/kernels/test/op_lt_test.cpp
index c17d69c37da..5b92f7556bd 100644
--- a/kernels/test/op_lt_test.cpp
+++ b/kernels/test/op_lt_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -92,9 +93,9 @@ TEST_F(OpLtScalarOutTest, BoolInputDtype) {
 
 // Mismatched shape tests.
 TEST_F(OpLtScalarOutTest, MismatchedInOutShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
@@ -135,9 +136,9 @@ TEST_F(OpLtTensorOutTest, AllDtypesSupported) {
 }
 
 TEST_F(OpLtTensorOutTest, MismatchedInShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
@@ -149,9 +150,9 @@ TEST_F(OpLtTensorOutTest, MismatchedInShapesDies) {
 }
 
 TEST_F(OpLtTensorOutTest, MismatchedInOutShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
diff --git a/kernels/test/op_masked_fill_test.cpp b/kernels/test/op_masked_fill_test.cpp
index b36b54c2b81..443f8909b8a 100644
--- a/kernels/test/op_masked_fill_test.cpp
+++ b/kernels/test/op_masked_fill_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -246,9 +247,9 @@ TEST_F(OpMaskedFillTest, BroadcastTest) {
 }
 
 TEST_F(OpMaskedFillTest, MismatchedOutputShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched output shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched output shape");
   TensorFactory tf;
   TensorFactory tf_bool;
 
diff --git a/kernels/test/op_max_test.cpp b/kernels/test/op_max_test.cpp
index 53c90ae909c..d49dcbc9ed0 100644
--- a/kernels/test/op_max_test.cpp
+++ b/kernels/test/op_max_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -281,9 +282,9 @@ TEST_F(OpMaxUnaryOutTest, EmptyFloatingInput) {
 }
 
 TEST_F(OpMaxOutTest, MismatchedDimensionsDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
 #define TEST_ENTRY(ctype, dtype) \
   test_max_out_invalid_dimensions();
   ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
@@ -291,9 +292,9 @@ TEST_F(OpMaxOutTest, MismatchedDimensionsDies) {
 }
 
 TEST_F(OpMaxOutTest, MismatchedDTypesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   TensorFactory tf_float;
   TensorFactory tf_long;
 
@@ -392,9 +393,9 @@ TEST_F(OpMaxOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpMaxOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   test_dynamic_shape(
       {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_mean_test.cpp b/kernels/test/op_mean_test.cpp
index 23f4b675d68..6633ab9c3c7 100644
--- a/kernels/test/op_mean_test.cpp
+++ b/kernels/test/op_mean_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -263,9 +264,9 @@ void OpMeanOutTest::
 }
 
 TEST_F(OpMeanOutTest, InvalidDimensionListDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   // Use a two layer switch to hanldle each possible data pair
 #define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE, OUTPUT_CTYPE, OUTPUT_DTYPE) \
   test_mean_dim_out_invalid_dimensions<                                   \
@@ -281,9 +282,9 @@ TEST_F(OpMeanOutTest, InvalidDimensionListDies) {
 }
 
 TEST_F(OpMeanOutTest, InvalidShapeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   // Use a two layer switch to hanldle each possible data pair
 #define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE, OUTPUT_CTYPE, OUTPUT_DTYPE) \
   test_mean_dim_out_invalid_shape<                                        \
@@ -299,9 +300,9 @@ TEST_F(OpMeanOutTest, InvalidShapeDies) {
 }
 
 TEST_F(OpMeanOutTest, MismatchedDTypesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   TensorFactory tf_float;
   TensorFactory tf_int;
 
@@ -351,9 +352,9 @@ TEST_F(OpMeanOutTest, AllRealInputFloatOutputPasses) {
 }
 
 TEST_F(OpMeanOutTest, HalfSupport) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "Test Half support only for ExecuTorch mode";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "Test Half support only for ExecuTorch mode");
 #define TEST_ENTRY(ctype, dtype) \
   test_mean_dim_out_dtype();
   ET_FORALL_REALH_TYPES(TEST_ENTRY);
diff --git a/kernels/test/op_min_test.cpp b/kernels/test/op_min_test.cpp
index ebbca989051..684782a8d87 100644
--- a/kernels/test/op_min_test.cpp
+++ b/kernels/test/op_min_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -277,9 +278,9 @@ TEST_F(OpMinUnaryOutTest, EmptyFloatingInput) {
 }
 
 TEST_F(OpMinOutTest, MismatchedDimensionsDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
 #define TEST_ENTRY(ctype, dtype) \
   test_min_out_invalid_dimensions();
   ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
@@ -287,9 +288,9 @@ TEST_F(OpMinOutTest, MismatchedDimensionsDies) {
 }
 
 TEST_F(OpMinOutTest, MismatchedDTypesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   TensorFactory tf_float;
   TensorFactory tf_long;
 
@@ -366,9 +367,9 @@ TEST_F(OpMinOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpMinOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   test_dynamic_shape(
       {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_minimum_test.cpp b/kernels/test/op_minimum_test.cpp
index 9c256963943..1e8fc65a385 100644
--- a/kernels/test/op_minimum_test.cpp
+++ b/kernels/test/op_minimum_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -126,9 +127,9 @@ TEST_F(OpMinimumOutTest, MismatchedOutputShapesDies) {
 }
 
 TEST_F(OpMinimumOutTest, MismatchedOutputShapeWithSingletonDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched output shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched output shape");
   // First argument is singleton but second and output has different shape.
   TensorFactory tf;
   Tensor out = tf.zeros({4, 4});
@@ -186,9 +187,9 @@ TEST_F(OpMinimumOutTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpMinimumOutTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
   %rewrite(binary_op) */
@@ -227,9 +228,9 @@ TEST_F(OpMinimumOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpMinimumOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(binary_op) */
diff --git a/kernels/test/op_mm_test.cpp b/kernels/test/op_mm_test.cpp
index 62d5ed29e26..ff3dfc14071 100644
--- a/kernels/test/op_mm_test.cpp
+++ b/kernels/test/op_mm_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -36,11 +37,9 @@ class OpMmOutTest : public OperatorTest {
     TensorFactory tf;
 
     if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-      if (DTYPE == ScalarType::Half) {
-        GTEST_SKIP()
-            << "skip Half because torch::executor::aten::mm_out does not support Half";
-        return;
-      }
+      ET_SKIP_IF(
+          DTYPE == ScalarType::Half,
+          "skip Half because torch::executor::aten::mm_out does not support Half");
     }
 
     // matmul gives 4 * 2 * 3 = 24
@@ -135,9 +134,9 @@ TEST_F(OpMmOutTest, MismatchedDimensionsDies) {
 }
 
 TEST_F(OpMmOutTest, MismatchedDimensionSizeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched dimension size";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched dimension size");
   TensorFactory tf;
   Tensor x = tf.full({2, 2}, 3);
 
@@ -154,9 +153,9 @@ TEST_F(OpMmOutTest, MismatchedDimensionSizeDies) {
 }
 
 TEST_F(OpMmOutTest, WrongOutShapeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle wrong out shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle wrong out shape");
   TensorFactory tf;
   Tensor x = tf.ones({10, 3});
 
diff --git a/kernels/test/op_mul_test.cpp b/kernels/test/op_mul_test.cpp
index 4553f8a53b6..e4121651b5a 100644
--- a/kernels/test/op_mul_test.cpp
+++ b/kernels/test/op_mul_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -430,9 +431,9 @@ TEST_F(OpMulOutTest, OptimizedPathIgnoresLeading1Dimensions) {
 
 // Mismatched shape tests.
 TEST_F(OpMulOutTest, MismatchedNonBroadcastableInputShapesDies) {
-  if (SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen currently supports mismatched shapes";
-  }
+  ET_SKIP_IF(
+      SupportedFeatures::get()->is_aten,
+      "ATen currently supports mismatched shapes");
 
   TensorFactory tf;
 
@@ -519,9 +520,9 @@ TEST_F(OpMulOutTest, AllComplexDtypesSupported) {
 }
 
 TEST_F(OpMulOutTest, MismatchedOutputShapesDies) {
-  if (SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen currently supports mismatched shapes";
-  }
+  ET_SKIP_IF(
+      SupportedFeatures::get()->is_aten,
+      "ATen currently supports mismatched shapes");
 
   TensorFactory tf;
 
diff --git a/kernels/test/op_native_layer_norm_test.cpp b/kernels/test/op_native_layer_norm_test.cpp
index 930214d238c..b190a3e89ad 100644
--- a/kernels/test/op_native_layer_norm_test.cpp
+++ b/kernels/test/op_native_layer_norm_test.cpp
@@ -12,6 +12,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -447,9 +448,9 @@ TEST_F(OpNativeLayerNormTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpNativeLayerNormTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   test_dynamic_shape(
       {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_ne_test.cpp b/kernels/test/op_ne_test.cpp
index 46681b02725..a1107233887 100644
--- a/kernels/test/op_ne_test.cpp
+++ b/kernels/test/op_ne_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -102,9 +103,9 @@ TEST_F(OpNeScalarOutTest, BoolInputDtype) {
 
 // Mismatched shape tests.
 TEST_F(OpNeScalarOutTest, MismatchedShapesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf_int;
   TensorFactory tf_bool;
 
@@ -180,9 +181,9 @@ TEST_F(OpNeScalarOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpNeScalarOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   /* %python
   out_args = "{1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op_out_dtype) */
diff --git a/kernels/test/op_permute_copy_test.cpp b/kernels/test/op_permute_copy_test.cpp
index bb7b186a672..3273ea35481 100644
--- a/kernels/test/op_permute_copy_test.cpp
+++ b/kernels/test/op_permute_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -452,9 +453,9 @@ TEST_F(OpPermuteCopyTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpPermuteCopyTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   /* %python
   out_args = "{1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op) */
diff --git a/kernels/test/op_repeat_test.cpp b/kernels/test/op_repeat_test.cpp
index 5a04044d34c..d6e290b1493 100644
--- a/kernels/test/op_repeat_test.cpp
+++ b/kernels/test/op_repeat_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -87,9 +88,9 @@ class OpRepeatOutTest : public OperatorTest {
 };
 
 TEST_F(OpRepeatOutTest, AllDtypesSupported) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
 #define TEST_ENTRY(ctype, dtype) run_dtype_tests();
   ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
 #undef TEST_ENTRY
@@ -209,9 +210,9 @@ TEST_F(OpRepeatOutTest, NegativeRepeatDie) {
 }
 
 TEST_F(OpRepeatOutTest, WrongOutputShapeDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle wrong output shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle wrong output shape");
   TensorFactory tf;
 
   Tensor x = tf.ones(
@@ -246,9 +247,9 @@ TEST_F(OpRepeatOutTest, OutputDtypeMismatchedDie) {
 // Right now we only support the dimension of input and output no larger
 // than 16.
 TEST_F(OpRepeatOutTest, TooManyDimensionsDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle larger number of dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle larger number of dimensions");
   TensorFactory tf;
 
   Tensor x = tf.ones(
@@ -361,9 +362,9 @@ TEST_F(OpRepeatOutTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpRepeatOutTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
   %rewrite(unary_op) */
@@ -384,9 +385,9 @@ TEST_F(OpRepeatOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpRepeatOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op) */
diff --git a/kernels/test/op_rsub_test.cpp b/kernels/test/op_rsub_test.cpp
index e2bcbd78dcc..fff4e604e29 100644
--- a/kernels/test/op_rsub_test.cpp
+++ b/kernels/test/op_rsub_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -223,9 +224,9 @@ TEST_F(OpRSubScalarOutTest, MismatchedOutputDtypeDies) {
 // Mismatched shape tests.
 
 TEST_F(OpRSubScalarOutTest, MismatchedOutputShapesDies) {
-  if (SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle output shapes";
-  }
+  ET_SKIP_IF(
+      SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle output shapes");
 
   TensorFactory tf;
 
@@ -248,17 +249,17 @@ TEST_F(OpRSubScalarOutTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpRSubScalarOutTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
 }
 
 TEST_F(OpRSubScalarOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_scalar_tensor_test.cpp b/kernels/test/op_scalar_tensor_test.cpp
index 0be6f395eb0..b47ecd19e20 100644
--- a/kernels/test/op_scalar_tensor_test.cpp
+++ b/kernels/test/op_scalar_tensor_test.cpp
@@ -10,6 +10,7 @@
 #include 
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -91,28 +92,28 @@ class OpScalarTensorOutTest : public OperatorTest {
 
 ET_FORALL_REAL_TYPES_AND3(Half, Bool, BFloat16, GENERATE_TEST_0D)
 
-#define GENERATE_TEST(ctype, dtype)                                    \
-  TEST_F(OpScalarTensorOutTest, dtype##Tensors) {                      \
-    if (torch::executor::testing::SupportedFeatures::get()->is_aten) { \
-      GTEST_SKIP() << "ATen kernel resizes output to shape {}";        \
-    }                                                                  \
-    test_scalar_tensor_out_1d(2);            \
-    test_scalar_tensor_out_2d(2);            \
-    test_scalar_tensor_out_3d(2);            \
-    test_scalar_tensor_out_1d(4);            \
-    test_scalar_tensor_out_2d(4);            \
-    test_scalar_tensor_out_3d(4);            \
-    test_scalar_tensor_out_1d(7);            \
-    test_scalar_tensor_out_2d(7);            \
-    test_scalar_tensor_out_3d(7);            \
+#define GENERATE_TEST(ctype, dtype)                                  \
+  TEST_F(OpScalarTensorOutTest, dtype##Tensors) {                    \
+    ET_SKIP_IF(                                                      \
+        torch::executor::testing::SupportedFeatures::get()->is_aten, \
+        "ATen kernel resizes output to shape {}");                   \
+    test_scalar_tensor_out_1d(2);          \
+    test_scalar_tensor_out_2d(2);          \
+    test_scalar_tensor_out_3d(2);          \
+    test_scalar_tensor_out_1d(4);          \
+    test_scalar_tensor_out_2d(4);          \
+    test_scalar_tensor_out_3d(4);          \
+    test_scalar_tensor_out_1d(7);          \
+    test_scalar_tensor_out_2d(7);          \
+    test_scalar_tensor_out_3d(7);          \
   }
 
 ET_FORALL_REAL_TYPES_AND3(Half, Bool, BFloat16, GENERATE_TEST)
 
 TEST_F(OpScalarTensorOutTest, InvalidOutShapeFails) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel will reshape output";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel will reshape output");
 
   TensorFactory tf;
   std::vector sizes{1, 2, 1};
diff --git a/kernels/test/op_scatter_add_test.cpp b/kernels/test/op_scatter_add_test.cpp
index d5511b72683..89770898f9b 100644
--- a/kernels/test/op_scatter_add_test.cpp
+++ b/kernels/test/op_scatter_add_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -324,9 +325,9 @@ TEST_F(OpScatterAddOutTest, InvalidDimensionsDies) {
 }
 
 TEST_F(OpScatterAddOutTest, MismatchedShapeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shape");
 #define TEST_ENTRY(CTYPE, DTYPE) \
   test_scatter_add_out_mismatched_shape();
   ET_FORALL_REAL_TYPES(TEST_ENTRY);
@@ -388,17 +389,17 @@ TEST_F(OpScatterAddOutTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpScatterAddOutTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {10, 10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
 }
 
 TEST_F(OpScatterAddOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_scatter_test.cpp b/kernels/test/op_scatter_test.cpp
index dac9017d188..79f106e3578 100644
--- a/kernels/test/op_scatter_test.cpp
+++ b/kernels/test/op_scatter_test.cpp
@@ -10,6 +10,7 @@
 #include 
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -468,9 +469,9 @@ TEST_F(OpScatterValueOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpScatterValueOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_select_copy_test.cpp b/kernels/test/op_select_copy_test.cpp
index c78c09628c8..470e8cae717 100644
--- a/kernels/test/op_select_copy_test.cpp
+++ b/kernels/test/op_select_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -395,9 +396,9 @@ TEST_F(OpSelectCopyIntOutTest, MismatchedDtypesDies) {
 }
 
 TEST_F(OpSelectCopyIntOutTest, OutMatchNumelLackDimAtEndDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle out with mismatched dimensions");
   TensorFactory tf;
   Tensor x = tf.zeros({1, 2, 2, 1});
 
@@ -410,9 +411,9 @@ TEST_F(OpSelectCopyIntOutTest, OutMatchNumelLackDimAtEndDies) {
 }
 
 TEST_F(OpSelectCopyIntOutTest, OutMatchNumelExtraDimAtFrontDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle out with mismatched dimensions");
   TensorFactory tf;
   Tensor x = tf.zeros({2, 2});
 
@@ -425,9 +426,9 @@ TEST_F(OpSelectCopyIntOutTest, OutMatchNumelExtraDimAtFrontDies) {
 }
 
 TEST_F(OpSelectCopyIntOutTest, OutSizeMismatchDimDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle out with mismatched dimensions");
   TensorFactory tf;
 
   Tensor x = tf.zeros({2, 4, 7, 5});
@@ -484,9 +485,9 @@ TEST_F(OpSelectCopyIntOutTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpSelectCopyIntOutTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{5, 5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
   %rewrite(unary_op) */
@@ -521,9 +522,9 @@ TEST_F(OpSelectCopyIntOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpSelectCopyIntOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op) */
diff --git a/kernels/test/op_select_scatter_test.cpp b/kernels/test/op_select_scatter_test.cpp
index 185331d399c..268c186004a 100644
--- a/kernels/test/op_select_scatter_test.cpp
+++ b/kernels/test/op_select_scatter_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -651,17 +652,17 @@ TEST_F(OpSelectScatterOutTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpSelectScatterOutTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {10, 10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
 }
 
 TEST_F(OpSelectScatterOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_sigmoid_test.cpp b/kernels/test/op_sigmoid_test.cpp
index 57771cc3c40..2d75d5c02e4 100644
--- a/kernels/test/op_sigmoid_test.cpp
+++ b/kernels/test/op_sigmoid_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -89,9 +90,9 @@ class OpSigmoidOutTest : public OperatorTest {
 };
 
 TEST_F(OpSigmoidOutTest, AllRealInputHalfOutputSupport) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "Test Half support only for ExecuTorch mode";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "Test Half support only for ExecuTorch mode");
 #define TEST_ENTRY(ctype, dtype) \
   test_integer_sigmoid_out();
   ET_FORALL_REALH_TYPES(TEST_ENTRY);
@@ -124,9 +125,9 @@ TEST_F(OpSigmoidOutTest, BooleanInputDoubleOutputSupport) {
 
 // Mismatched shape tests.
 TEST_F(OpSigmoidOutTest, MismatchedShapesDies) {
-  if (SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
 
   TensorFactory tf;
   TensorFactory tf_out;
diff --git a/kernels/test/op_sign_test.cpp b/kernels/test/op_sign_test.cpp
index 2754e784928..950419240e8 100644
--- a/kernels/test/op_sign_test.cpp
+++ b/kernels/test/op_sign_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -44,18 +45,18 @@ class OpSignTest : public OperatorTest {
 };
 
 TEST_F(OpSignTest, ETSanityCheckFloat) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen returns 0 on NAN input";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen returns 0 on NAN input");
 #define TEST_ENTRY(ctype, dtype) test_et_dtype();
   ET_FORALL_FLOATHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 
 TEST_F(OpSignTest, ATenSanityCheckFloat) {
-  if (!torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ET returns NAN on NAN input";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ET returns NAN on NAN input");
   TensorFactory tf;
 
   Tensor in = tf.make({1, 7}, {-INFINITY, -3., -1.5, 0., 1.5, NAN, INFINITY});
diff --git a/kernels/test/op_slice_copy_test.cpp b/kernels/test/op_slice_copy_test.cpp
index c7e8a0acf66..e2a48322463 100644
--- a/kernels/test/op_slice_copy_test.cpp
+++ b/kernels/test/op_slice_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -48,7 +49,7 @@ class OpSliceCopyTensorOutTest : public OperatorTest {
         5,   6,   7,   8, // [1, :]
         9,  10,  11,  12, // [2, :]
       });
-  
+
     // op_slice_copy_tensor_out(input, /*dim=*/0, /*start=*/0, /*end=*/2, /*step=*/1, out),
     // The result should equal to input[0:2:1, :]
     Tensor expect_ret = tf.make(
@@ -428,9 +429,9 @@ TEST_F(OpSliceCopyTensorOutTest, LegalStepsSupported) {
 /// A generic smoke test that works for any dtype that supports ones() and
 /// zeros().
 TEST_F(OpSliceCopyTensorOutTest, AllDtypesSupported) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
 #define TEST_ENTRY(ctype, dtype) test_dtype();
   ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
 #undef TEST_ENTRY
@@ -541,9 +542,9 @@ TEST_F(OpSliceCopyTensorOutTest, MismatchedDtypesDies) {
 }
 
 TEST_F(OpSliceCopyTensorOutTest, OutSizeMismatchDimDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle out with mismatched dimensions");
   TensorFactory tf;
 
   Tensor input = tf.zeros({2, 4, 7, 5});
@@ -649,9 +650,9 @@ TEST_F(OpSliceCopyTensorOutTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpSliceCopyTensorOutTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{10, 10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
   %rewrite(unary_op) */
@@ -694,9 +695,9 @@ TEST_F(OpSliceCopyTensorOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpSliceCopyTensorOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op) */
diff --git a/kernels/test/op_slice_scatter_test.cpp b/kernels/test/op_slice_scatter_test.cpp
index 14a5bd2679d..501141c5089 100644
--- a/kernels/test/op_slice_scatter_test.cpp
+++ b/kernels/test/op_slice_scatter_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -765,9 +766,9 @@ TEST_F(OpSliceScatterTensorOutTest, MismatchedOutDtypesDies) {
 }
 
 TEST_F(OpSliceScatterTensorOutTest, OutSizeMismatchDimDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle out with mismatched dimensions");
   TensorFactory tf;
 
   Tensor input = tf.zeros({2, 4, 7, 5});
@@ -783,9 +784,9 @@ TEST_F(OpSliceScatterTensorOutTest, OutSizeMismatchDimDies) {
 }
 
 TEST_F(OpSliceScatterTensorOutTest, SrcSizeMismatchDimDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle out with mismatched dimensions");
   TensorFactory tf;
 
   Tensor input = tf.zeros({2, 4, 7, 5});
diff --git a/kernels/test/op_softmax_test.cpp b/kernels/test/op_softmax_test.cpp
index 3c61acb7d29..3f515bb4dcc 100644
--- a/kernels/test/op_softmax_test.cpp
+++ b/kernels/test/op_softmax_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -131,9 +132,9 @@ TEST_F(OpSoftmaxOutTest, MismatchedDimensionsDies) {
 }
 
 TEST_F(OpSoftmaxOutTest, MismatchedDimensionSizeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched dimension size";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched dimension size");
   TensorFactory tf;
 
   Tensor x = tf.ones({3, 4});
@@ -147,9 +148,9 @@ TEST_F(OpSoftmaxOutTest, MismatchedDimensionSizeDies) {
 }
 
 TEST_F(OpSoftmaxOutTest, NegativeDim) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   TensorFactory tf;
 
   // Input tensor with shape (2, 3) and values (0, 1, 2, 3, 4, 5).
diff --git a/kernels/test/op_split_copy_test.cpp b/kernels/test/op_split_copy_test.cpp
index 2dd112b1ace..0f35e88e9f3 100644
--- a/kernels/test/op_split_copy_test.cpp
+++ b/kernels/test/op_split_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -402,7 +403,7 @@ TEST_F(OpSplitCopyTensorOutTest, OutOfRangeDimsDie) {
 }
 
 TEST_F(OpSplitCopyTensorOutTest, DtypeMismatchDies) {
-  GTEST_SKIP() << "ATen kernel can handle dtype mismatch";
+  ET_SKIP_IF(true, "ATen kernel can handle dtype mismatch");
   TensorFactory tf_int;
   TensorListFactory tlf_int;
   TensorListFactory tlf_float;
@@ -474,9 +475,9 @@ TEST_F(OpSplitCopyTensorOutTest, WrongNumOutputEntriesDies) {
 }
 
 TEST_F(OpSplitCopyTensorOutTest, WrongOutputShapeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle wrong out shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle wrong out shape");
   TensorFactory tf;
   TensorListFactory tlf;
 
diff --git a/kernels/test/op_squeeze_copy_test.cpp b/kernels/test/op_squeeze_copy_test.cpp
index 53f3465c508..bdd980f1357 100644
--- a/kernels/test/op_squeeze_copy_test.cpp
+++ b/kernels/test/op_squeeze_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -246,9 +247,9 @@ TEST_F(OpSqueezeTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpSqueezeTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{5, 5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
   %rewrite(unary_op) */
@@ -283,9 +284,9 @@ TEST_F(OpSqueezeTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpSqueezeTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op) */
diff --git a/kernels/test/op_stack_test.cpp b/kernels/test/op_stack_test.cpp
index 9a102878c0f..77fef7eed8e 100644
--- a/kernels/test/op_stack_test.cpp
+++ b/kernels/test/op_stack_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -348,9 +349,9 @@ TEST_F(OpStackOutTest, MismatchedDtypesDies) {
 }
 
 TEST_F(OpStackOutTest, OutMatchNumelWithExtraDimAtEndDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle out with mismatched dimensions");
   TensorFactory tf;
   Tensor out = tf.zeros({1, 2, 2, 1});
 
@@ -365,9 +366,9 @@ TEST_F(OpStackOutTest, OutMatchNumelWithExtraDimAtEndDies) {
 }
 
 TEST_F(OpStackOutTest, OutMatchNumelLackDimAtFrontDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle out with mismatched dimensions");
   TensorFactory tf;
   Tensor out = tf.zeros({2, 2});
 
@@ -382,9 +383,9 @@ TEST_F(OpStackOutTest, OutMatchNumelLackDimAtFrontDies) {
 }
 
 TEST_F(OpStackOutTest, OutRegularMismatchDimDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle out with mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle out with mismatched dimensions");
   TensorFactory tf;
 
   // Should be {2, 2, 3} to match the inputs when calling stack() with dim 0.
@@ -437,9 +438,9 @@ TEST_F(OpStackOutTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpStackOutTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{5, 5, 5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
   %rewrite(unary_op_tensor_list_in) */
@@ -462,9 +463,9 @@ TEST_F(OpStackOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpStackOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op_tensor_list_in) */
diff --git a/kernels/test/op_sub_test.cpp b/kernels/test/op_sub_test.cpp
index 41ebc2f2733..dedea0741b4 100644
--- a/kernels/test/op_sub_test.cpp
+++ b/kernels/test/op_sub_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -444,10 +445,9 @@ TEST_F(OpSubOutTest, MismatchedNonBroadcastableInputShapesDies) {
 }
 
 TEST_F(OpSubOutTest, MismatchedOutputShapesDies) {
-  if (SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP()
-        << "The current kernel supports implicitly resizing output tensor";
-  }
+  ET_SKIP_IF(
+      SupportedFeatures::get()->output_resize,
+      "The current kernel supports implicitly resizing output tensor");
 
   TensorFactory tf;
 
diff --git a/kernels/test/op_sum_test.cpp b/kernels/test/op_sum_test.cpp
index 58624c2a110..18c71b1080b 100644
--- a/kernels/test/op_sum_test.cpp
+++ b/kernels/test/op_sum_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -307,9 +308,9 @@ class OpSumOutTest : public OperatorTest {
 };
 
 TEST_F(OpSumOutTest, InvalidDimensionListDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   // Use a two layer switch to hanldle each possible data pair
 #define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE, OUTPUT_CTYPE, OUTPUT_DTYPE) \
   test_sum_dim_out_invalid_dimensions<                                    \
@@ -325,9 +326,9 @@ TEST_F(OpSumOutTest, InvalidDimensionListDies) {
 }
 
 TEST_F(OpSumOutTest, InvalidShapeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   // Use a two layer switch to hanldle each possible data pair
 #define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE, OUTPUT_CTYPE, OUTPUT_DTYPE) \
   test_sum_dim_out_invalid_shape<                                         \
@@ -343,9 +344,9 @@ TEST_F(OpSumOutTest, InvalidShapeDies) {
 }
 
 TEST_F(OpSumOutTest, MismatchedDTypesDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   TensorFactory tf_float;
   TensorFactory tf_int;
 
diff --git a/kernels/test/op_t_copy_test.cpp b/kernels/test/op_t_copy_test.cpp
index 142fe3050f4..6f01b4c459a 100644
--- a/kernels/test/op_t_copy_test.cpp
+++ b/kernels/test/op_t_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -39,9 +40,9 @@ TEST_F(OpTCopyTest, 1DTranspose) {
 }
 
 TEST_F(OpTCopyTest, 1DTransposeMismatchShapeDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf;
 
   Tensor t_in = tf.make({4}, {1, 2, 3, 4});
@@ -62,9 +63,9 @@ TEST_F(OpTCopyTest, 2DTranspose) {
 }
 
 TEST_F(OpTCopyTest, 2DTransposeMismatchShapeDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched shapes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched shapes");
   TensorFactory tf;
 
   Tensor t_in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6});
@@ -122,9 +123,9 @@ TEST_F(OpTCopyTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpTCopyTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
   %rewrite(unary_op) */
@@ -155,9 +156,9 @@ TEST_F(OpTCopyTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpTCopyTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op) */
diff --git a/kernels/test/op_to_copy_test.cpp b/kernels/test/op_to_copy_test.cpp
index d9798d6d573..45b2b2f6020 100644
--- a/kernels/test/op_to_copy_test.cpp
+++ b/kernels/test/op_to_copy_test.cpp
@@ -14,6 +14,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -427,9 +428,9 @@ TEST_F(OpToTest, HardcodeFloatConvertInt) {
 }
 
 TEST_F(OpToTest, MismatchedSizesDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched sizes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched sizes");
   TensorFactory tf;
   Tensor input = tf.make(/*sizes=*/{3, 1, 1, 2}, /*data=*/{1, 2, 3, 4, 5, 6});
   Tensor out = tf.zeros({3, 2, 1, 1});
@@ -446,9 +447,9 @@ TEST_F(OpToTest, MismatchedSizesDie) {
 // should not be allowed. The function is expected death if using the illegal
 // memory format.
 TEST_F(OpToTest, MismatchedMemoryFormatDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle non contiguous memory formats";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle non contiguous memory formats");
   TensorFactory tf_in;
   TensorFactory tf_out;
   Tensor input =
@@ -474,9 +475,9 @@ TEST_F(OpToTest, MismatchedMemoryFormatDies) {
 
 // Only blocking data transfer supported
 TEST_F(OpToTest, MismatchedBlockingDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle non blocking data transfer";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle non blocking data transfer");
   TensorFactory tf;
   Tensor input = tf.make(/*sizes=*/{3, 1, 1, 2}, /*data=*/{1, 2, 3, 4, 5, 6});
   Tensor out = tf.zeros(/*sizes=*/{3, 1, 1, 2});
@@ -500,9 +501,9 @@ TEST_F(OpToTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpToTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   test_dynamic_shape(
       {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_transpose_copy_test.cpp b/kernels/test/op_transpose_copy_test.cpp
index e9f9bff4acc..d3797498886 100644
--- a/kernels/test/op_transpose_copy_test.cpp
+++ b/kernels/test/op_transpose_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -158,9 +159,9 @@ TEST_F(OpTransposeIntCopyTest, OutOfBoundDimDies) {
 
 // transpose a 3d tensor into a 2d one
 TEST_F(OpTransposeIntCopyTest, MismatchedDimDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched dimensions";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched dimensions");
   TensorFactory tf;
 
   Tensor a = tf.ones(/*sizes=*/{4, 2, 3});
@@ -196,9 +197,9 @@ TEST_F(OpTransposeIntCopyTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpTransposeIntCopyTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{5, 5, 5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
   %rewrite(unary_op) */
@@ -215,9 +216,9 @@ TEST_F(OpTransposeIntCopyTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpTransposeIntCopyTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op) */
diff --git a/kernels/test/op_tril_test.cpp b/kernels/test/op_tril_test.cpp
index 9a2347e39c0..9b8540713da 100644
--- a/kernels/test/op_tril_test.cpp
+++ b/kernels/test/op_tril_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -760,9 +761,9 @@ TEST_F(OpTrilTest, InvalidInputShapesDies) {
 
 TEST_F(OpTrilTest, MismatchedOutputShapesDies) {
   // Skip ATen test since it supports `self` and `out` having different shapes.
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched output shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched output shape");
 
   TensorFactory tf;
 
@@ -788,9 +789,9 @@ TEST_F(OpTrilTest, MismatchedOutputDtypeDies) {
 
 TEST_F(OpTrilTest, InvalidTensorDims) {
   // Skip ATen test since it supports `self` and `out` having different shapes.
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched output shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched output shape");
 
   TensorFactory tf;
 
diff --git a/kernels/test/op_unbind_copy_test.cpp b/kernels/test/op_unbind_copy_test.cpp
index c98edc5e1f7..54b32ef73f6 100644
--- a/kernels/test/op_unbind_copy_test.cpp
+++ b/kernels/test/op_unbind_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -269,9 +270,9 @@ TEST_F(OpUnbindCopyIntOutTest, UnbindWorksWithZeroSizedTensors) {
 }
 
 TEST_F(OpUnbindCopyIntOutTest, UnbindFailsWithWronglyAllocatedOutput) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched output shape";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched output shape");
   TensorFactory tf;
   TensorListFactory tlf;
 
diff --git a/kernels/test/op_unsqueeze_copy_test.cpp b/kernels/test/op_unsqueeze_copy_test.cpp
index d90d69a1b24..b6fd092a982 100644
--- a/kernels/test/op_unsqueeze_copy_test.cpp
+++ b/kernels/test/op_unsqueeze_copy_test.cpp
@@ -8,6 +8,7 @@
 
 #include  // Declares the operator
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -123,9 +124,9 @@ TEST_F(OpUnsqueezeTest, EmptyInputSupported) {
 }
 
 TEST_F(OpUnsqueezeTest, InputOutputMismatchedSizesDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched sizes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched sizes");
   TensorFactory tf;
 
   Tensor input = tf.make(/*sizes=*/{3, 1, 2}, /*data=*/{1, 2, 3, 4, 5, 6});
@@ -139,9 +140,9 @@ TEST_F(OpUnsqueezeTest, InputOutputMismatchedSizesDie) {
 }
 
 TEST_F(OpUnsqueezeTest, DimOutputMismatchedSizesDie) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle mismatched sizes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can handle mismatched sizes");
   TensorFactory tf;
   Tensor input = tf.make(/*sizes=*/{3, 1, 2}, /*data=*/{1, 2, 3, 4, 5, 6});
   Tensor out = tf.ones(/*sizes=*/{3, 1, 2, 1});
@@ -264,9 +265,9 @@ TEST_F(OpUnsqueezeTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpUnsqueezeTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{5, 5, 5}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND"
   %rewrite(unary_op) */
@@ -301,9 +302,9 @@ TEST_F(OpUnsqueezeTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpUnsqueezeTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   /* %python
   out_args = "{1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND"
   %rewrite(unary_op) */
diff --git a/kernels/test/op_upsample_bilinear2d_test.cpp b/kernels/test/op_upsample_bilinear2d_test.cpp
index 95fea942e39..9cd317db55e 100644
--- a/kernels/test/op_upsample_bilinear2d_test.cpp
+++ b/kernels/test/op_upsample_bilinear2d_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -315,10 +316,9 @@ TEST_F(OpUpsampleBilinear2dTest, DType) {
 }
 
 TEST_F(OpUpsampleBilinear2dTest, MismatchedOutputSizeDies) {
-  if (SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP()
-        << "The current kernel supports implicitly resizing output tensor";
-  }
+  ET_SKIP_IF(
+      SupportedFeatures::get()->output_resize,
+      "The current kernel supports implicitly resizing output tensor");
   TensorFactory tf;
 
   const auto input = tf.ones({1, 1, 1, 2});
@@ -471,9 +471,9 @@ TEST_F(OpUpsampleBilinear2dTest, ZeroComputedOutputSizeDies) {
 TEST_F(OpUpsampleBilinear2dTest, MismatchedDimOrderDies) {
   TensorFactory tf;
 
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can implicitly convert dim order";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel can implicitly convert dim order");
 
   const auto input = tf.ones({1, 1, 1, 2});
   auto out = tf.zeros_channels_last({1, 1, 1, 4});
diff --git a/kernels/test/op_upsample_nearest2d_test.cpp b/kernels/test/op_upsample_nearest2d_test.cpp
index 76e66e666dd..b65e165e7b0 100644
--- a/kernels/test/op_upsample_nearest2d_test.cpp
+++ b/kernels/test/op_upsample_nearest2d_test.cpp
@@ -8,6 +8,7 @@
 
 #include  // Declares the operator
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -266,10 +267,9 @@ TEST_F(OpUpsampleNearest2dTest, DType) {
 }
 
 TEST_F(OpUpsampleNearest2dTest, MismatchedOutputSizeDies) {
-  if (SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP()
-        << "The current kernel supports implicitly resizing output tensor";
-  }
+  ET_SKIP_IF(
+      SupportedFeatures::get()->output_resize,
+      "The current kernel supports implicitly resizing output tensor");
   TensorFactory tf;
 
   const auto input = tf.ones({1, 1, 1, 2});
diff --git a/kernels/test/op_var_mean_test.cpp b/kernels/test/op_var_mean_test.cpp
index 05a0281a090..80571518e73 100644
--- a/kernels/test/op_var_mean_test.cpp
+++ b/kernels/test/op_var_mean_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -385,9 +386,9 @@ TEST_F(OpVarMeanCorrectionOutTest, SmokeTest) {
 }
 
 TEST_F(OpVarMeanCorrectionOutTest, KeepDim) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen supports fewer dtypes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen supports fewer dtypes");
 #define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE, OUTPUT_CTYPE, OUTPUT_DTYPE) \
   test_keepdim();
 
@@ -400,9 +401,9 @@ TEST_F(OpVarMeanCorrectionOutTest, KeepDim) {
 }
 
 TEST_F(OpVarMeanCorrectionOutTest, KeepDim_Aten) {
-  if (!torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen-specific variant of test case";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen-specific variant of test case");
 #define TEST_ENTRY(CTYPE, DTYPE) \
   test_keepdim();
 
@@ -411,9 +412,9 @@ TEST_F(OpVarMeanCorrectionOutTest, KeepDim_Aten) {
 }
 
 TEST_F(OpVarMeanCorrectionOutTest, MultipleDims) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen supports fewer dtypes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen supports fewer dtypes");
 #define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE, OUTPUT_CTYPE, OUTPUT_DTYPE) \
   test_multiple_dims();
 
@@ -426,9 +427,9 @@ TEST_F(OpVarMeanCorrectionOutTest, MultipleDims) {
 }
 
 TEST_F(OpVarMeanCorrectionOutTest, MultipleDims_Aten) {
-  if (!torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen-specific variant of test case";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen-specific variant of test case");
 #define TEST_ENTRY(CTYPE, DTYPE) \
   test_multiple_dims();
 
@@ -437,9 +438,9 @@ TEST_F(OpVarMeanCorrectionOutTest, MultipleDims_Aten) {
 }
 
 TEST_F(OpVarMeanCorrectionOutTest, NegativeDim) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen supports fewer dtypes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen supports fewer dtypes");
 #define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE, OUTPUT_CTYPE, OUTPUT_DTYPE) \
   test_negative_dim();
 
@@ -452,9 +453,9 @@ TEST_F(OpVarMeanCorrectionOutTest, NegativeDim) {
 }
 
 TEST_F(OpVarMeanCorrectionOutTest, NegativeDim_Aten) {
-  if (!torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen-specific variant of test case";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen-specific variant of test case");
 #define TEST_ENTRY(CTYPE, DTYPE) \
   test_negative_dim();
 
@@ -463,9 +464,9 @@ TEST_F(OpVarMeanCorrectionOutTest, NegativeDim_Aten) {
 }
 
 TEST_F(OpVarMeanCorrectionOutTest, NullAndEmptyDimList) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen supports fewer dtypes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen supports fewer dtypes");
 #define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE, OUTPUT_CTYPE, OUTPUT_DTYPE) \
   test_null_and_empty_dim_list<                                           \
       ScalarType::INPUT_DTYPE,                                            \
@@ -480,9 +481,9 @@ TEST_F(OpVarMeanCorrectionOutTest, NullAndEmptyDimList) {
 }
 
 TEST_F(OpVarMeanCorrectionOutTest, NullAndEmptyDimList_Aten) {
-  if (!torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen-specific variant of test case";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen-specific variant of test case");
 #define TEST_ENTRY(CTYPE, DTYPE) \
   test_null_and_empty_dim_list();
 
@@ -491,9 +492,9 @@ TEST_F(OpVarMeanCorrectionOutTest, NullAndEmptyDimList_Aten) {
 }
 
 TEST_F(OpVarMeanCorrectionOutTest, InvalidDimensionListDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
 #define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE, OUTPUT_CTYPE, OUTPUT_DTYPE) \
   test_invalid_dimensions();
 
@@ -506,9 +507,9 @@ TEST_F(OpVarMeanCorrectionOutTest, InvalidDimensionListDies) {
 }
 
 TEST_F(OpVarMeanCorrectionOutTest, InvalidDTypeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   TensorFactory tf_float;
   TensorFactory tf_int;
 
diff --git a/kernels/test/op_var_test.cpp b/kernels/test/op_var_test.cpp
index 63e7e94f982..ccd0a022af3 100644
--- a/kernels/test/op_var_test.cpp
+++ b/kernels/test/op_var_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -284,9 +285,9 @@ class OpVarCorrectionOutTest : public OperatorTest {
 };
 
 TEST_F(OpVarOutTest, InvalidDimensionListDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   // Use a two layer switch to hanldle each possible data pair
 #define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE, OUTPUT_CTYPE, OUTPUT_DTYPE) \
   test_var_out_invalid_dimensions<                                        \
@@ -302,9 +303,9 @@ TEST_F(OpVarOutTest, InvalidDimensionListDies) {
 }
 
 TEST_F(OpVarOutTest, InvalidShapeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   // Use a two layer switch to hanldle each possible data pair
 #define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE, OUTPUT_CTYPE, OUTPUT_DTYPE) \
   test_var_out_invalid_shape<                                             \
@@ -320,9 +321,9 @@ TEST_F(OpVarOutTest, InvalidShapeDies) {
 }
 
 TEST_F(OpVarOutTest, InvalidDTypeDies) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel test fails";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen kernel test fails");
   TensorFactory tf_float;
   TensorFactory tf_int;
 
@@ -356,9 +357,9 @@ TEST_F(OpVarOutTest, InvalidDTypeDies) {
 }
 
 TEST_F(OpVarOutTest, AllFloatInputFloatOutputPasses) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen supports fewer dtypes";
-  }
+  ET_SKIP_IF(
+      torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen supports fewer dtypes");
   // Use a two layer switch to hanldle each possible data pair
 #define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE, OUTPUT_CTYPE, OUTPUT_DTYPE) \
   test_var_out_dtype();
@@ -372,9 +373,9 @@ TEST_F(OpVarOutTest, AllFloatInputFloatOutputPasses) {
 }
 
 TEST_F(OpVarOutTest, AllFloatInputFloatOutputPasses_Aten) {
-  if (!torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen-specific variant of test case";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->is_aten,
+      "ATen-specific variant of test case");
   // Use a two layer switch to hanldle each possible data pair
 #define TEST_KERNEL(INPUT_CTYPE, INPUT_DTYPE, OUTPUT_CTYPE, OUTPUT_DTYPE) \
   test_var_out_dtype();
diff --git a/kernels/test/op_view_copy_test.cpp b/kernels/test/op_view_copy_test.cpp
index 73a59fa5bd5..4990f1b1a77 100644
--- a/kernels/test/op_view_copy_test.cpp
+++ b/kernels/test/op_view_copy_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -293,17 +294,17 @@ TEST_F(OpViewTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpViewTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
 }
 
 TEST_F(OpViewTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_where_test.cpp b/kernels/test/op_where_test.cpp
index c9e845e38a9..01847bc8464 100644
--- a/kernels/test/op_where_test.cpp
+++ b/kernels/test/op_where_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -447,17 +448,17 @@ TEST_F(OpWhereOutTest, DynamicShapeUpperBoundSameAsExpected) {
 }
 
 TEST_F(OpWhereOutTest, DynamicShapeUpperBoundLargerThanExpected) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {10, 10, 10}, torch::executor::TensorShapeDynamism::DYNAMIC_BOUND);
 }
 
 TEST_F(OpWhereOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape not supported");
   test_dynamic_shape(
       {1, 1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
diff --git a/kernels/test/op_zeros_test.cpp b/kernels/test/op_zeros_test.cpp
index 96586f91df5..0345906bcf5 100644
--- a/kernels/test/op_zeros_test.cpp
+++ b/kernels/test/op_zeros_test.cpp
@@ -9,6 +9,7 @@
 #include  // Declares the operator
 #include 
 #include 
+#include 
 #include 
 #include 
 #include 
@@ -75,9 +76,9 @@ TEST_F(OpZerosOutTest, DynamicShapeUpperBoundLargerThanExpected) {
 }
 
 TEST_F(OpZerosOutTest, DynamicShapeUnbound) {
-  if (!torch::executor::testing::SupportedFeatures::get()->output_resize) {
-    GTEST_SKIP() << "Dynamic shape unbound not supported";
-  }
+  ET_SKIP_IF(
+      !torch::executor::testing::SupportedFeatures::get()->output_resize,
+      "Dynamic shape unbound not supported");
   TensorFactory tf;
   Tensor expected = tf.zeros({3, 2});
 
diff --git a/kernels/test/supported_features_skip.h b/kernels/test/supported_features_skip.h
new file mode 100644
index 00000000000..b9416b56870
--- /dev/null
+++ b/kernels/test/supported_features_skip.h
@@ -0,0 +1,59 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+// `ET_SKIP_IF(cond, reason)` -- skip a kernel test when `cond` is true.
+//
+// Replaces the older inline pattern:
+//     if (SupportedFeatures::get()->is_aten) {
+//       GTEST_SKIP() << "ATen handles X";
+//     }
+// with:
+//     ET_SKIP_IF(SupportedFeatures::get()->is_aten, "ATen handles X");
+//
+// OSS:    expands to `if (cond) GTEST_SKIP() << reason;` (unchanged).
+// fbcode: expands to `if (cond) return;` so the test reports PASS, not SKIP.
+//
+// fbcode's TestX flags consistently-skipping tests as "broken" -- see
+// T208053850 and
+// https://fb.workplace.com/groups/testinfra.discuss/permalink/2044665472719153/.
+// Collapse back to the OSS form once that's resolved.
+//
+// `EXECUTORCH_INTERNAL` is set by BUCK gated on `runtime.is_oss` (see
+// `runtime/executor/targets.bzl` for the existing precedent).
+
+#if defined(EXECUTORCH_INTERNAL) && EXECUTORCH_INTERNAL == 1
+
+namespace executorch::testing::internal {
+// No-op sink so `<<` chains in the reason still parse and type-check.
+struct SkipReasonSink {
+  template 
+  const SkipReasonSink& operator<<(const T&) const {
+    return *this;
+  }
+};
+} // namespace executorch::testing::internal
+
+// `if/else` form avoids dangling-else hazards and lets the reason still
+// participate in `<<` chains.
+#define ET_SKIP_IF(cond, reason) \
+  if ((cond)) {                  \
+    return;                      \
+  } else                         \
+    ::executorch::testing::internal::SkipReasonSink{} << reason
+
+#else // !EXECUTORCH_INTERNAL
+
+#include 
+
+#define ET_SKIP_IF(cond, reason) \
+  if ((cond))                    \
+  GTEST_SKIP() << reason
+
+#endif // EXECUTORCH_INTERNAL
diff --git a/kernels/test/targets.bzl b/kernels/test/targets.bzl
index abf6329248d..bc51e336cb8 100644
--- a/kernels/test/targets.bzl
+++ b/kernels/test/targets.bzl
@@ -115,7 +115,16 @@ def define_common_targets():
     runtime.cxx_library(
         name = "supported_features_header",
         srcs = [],
-        exported_headers = {"supported_features.h": ":supported_feature_header_gen[supported_features.h]"},
+        exported_headers = {
+            "supported_features.h": ":supported_feature_header_gen[supported_features.h]",
+            "supported_features_skip.h": "supported_features_skip.h",
+        },
+        # Set EXECUTORCH_INTERNAL=1 for fbcode-internal builds so the
+        # ET_SKIP_IF helper in supported_features_skip.h compiles to an
+        # early `return;` instead of GTEST_SKIP. This avoids TestX's
+        # "ConsistentlySkipping" / broken-test signal. OSS builds keep
+        # the canonical GTEST_SKIP behavior. See header for context.
+        exported_preprocessor_flags = [] if runtime.is_oss else ["-DEXECUTORCH_INTERNAL=1"],
         visibility = [
             "//executorch/kernels/...",
         ],
diff --git a/runtime/backend/backend_options_map.h b/runtime/backend/backend_options_map.h
index 7d11911f4c4..f0c6a68c509 100644
--- a/runtime/backend/backend_options_map.h
+++ b/runtime/backend/backend_options_map.h
@@ -11,6 +11,7 @@
 #include 
 #include 
 #include 
+#include 
 
 #include 
 
@@ -171,6 +172,46 @@ class LoadBackendOptionsMap final {
     return size_;
   }
 
+  /**
+   * Non-owning view of a single (backend_id, options) entry, returned by
+   * entry_at(). The pointer / span are valid until the map is mutated or
+   * destroyed.
+   */
+  struct EntryView {
+    const char* backend_id = nullptr;
+    Span options;
+  };
+
+  /**
+   * Returns the (backend_id, options) entry at the given index for
+   * enumeration over the map's contents.
+   *
+   * @param index The entry index. Must be < size(); behavior is undefined
+   *     otherwise. Use this together with size() to walk every entry.
+   * @return EntryView referencing the entry's backend_id and options. The
+   *     view is valid until the next mutation of, or destruction of, this
+   *     map.
+   *
+   * Example:
+   * @code
+   *   for (size_t i = 0; i < map.size(); ++i) {
+   *     const auto entry = map.entry_at(i);
+   *     // use entry.backend_id and entry.options ...
+   *   }
+   * @endcode
+   */
+  EntryView entry_at(size_t index) const {
+    ET_DCHECK_MSG(
+        index < size_,
+        "entry_at index %zu out of bounds (size=%zu)",
+        index,
+        size_);
+    return EntryView{
+        entries_[index].backend_id,
+        Span(
+            entries_[index].options.data(), entries_[index].options.size())};
+  }
+
  private:
   static constexpr size_t kMaxBackends = 8;
   static constexpr size_t kMaxBackendIdLength = 64;
diff --git a/runtime/backend/test/backend_options_map_test.cpp b/runtime/backend/test/backend_options_map_test.cpp
index b98878d763b..f6466271ffa 100644
--- a/runtime/backend/test/backend_options_map_test.cpp
+++ b/runtime/backend/test/backend_options_map_test.cpp
@@ -365,3 +365,32 @@ TEST_F(LoadBackendOptionsMapTest, SetOptionsWithBuilderUpdatesExisting) {
   }
   EXPECT_EQ(num_threads2, 8); // Should be updated value
 }
+
+// Test entry_at returns each (backend_id, options) pair in insertion order
+// and the spans reference the same data the corresponding get_options
+// calls return.
+TEST_F(LoadBackendOptionsMapTest, EntryAtEnumeratesAllEntries) {
+  LoadBackendOptionsMap map;
+
+  BackendOptions<2> opts1;
+  opts1.set_option("k1", 1);
+  ASSERT_EQ(map.set_options("BackendA", opts1.view()), Error::Ok);
+
+  BackendOptions<2> opts2;
+  opts2.set_option("k2", true);
+  opts2.set_option("k3", "v");
+  ASSERT_EQ(map.set_options("BackendB", opts2.view()), Error::Ok);
+
+  ASSERT_EQ(map.size(), 2u);
+
+  const auto e0 = map.entry_at(0);
+  EXPECT_STREQ(e0.backend_id, "BackendA");
+  EXPECT_EQ(e0.options.size(), 1u);
+  // Spans returned by entry_at point at the same storage as get_options.
+  EXPECT_EQ(e0.options.data(), map.get_options("BackendA").data());
+
+  const auto e1 = map.entry_at(1);
+  EXPECT_STREQ(e1.backend_id, "BackendB");
+  EXPECT_EQ(e1.options.size(), 2u);
+  EXPECT_EQ(e1.options.data(), map.get_options("BackendB").data());
+}
diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp
index 1610804586d..8eb48480463 100644
--- a/runtime/executor/method.cpp
+++ b/runtime/executor/method.cpp
@@ -802,8 +802,19 @@ Error Method::resolve_operator(
   }
 
   // Find a kernel with the matching name and tensor meta.
-  Result op_function =
-      get_op_function_from_registry(operator_name, {meta, count});
+  // Try method-scoped registry first (if provided), then fall back to global.
+  auto resolve_op_function = [&]() -> Result {
+    if (!kernel_registry_.empty()) {
+      Result method_scoped_op_function =
+          get_op_function_from_registry(
+              operator_name, {meta, count}, kernel_registry_);
+      if (method_scoped_op_function.ok()) {
+        return method_scoped_op_function;
+      }
+    }
+    return get_op_function_from_registry(operator_name, {meta, count});
+  };
+  Result op_function = resolve_op_function();
   if (!op_function.ok()) {
     ET_LOG(
         Error,
@@ -831,7 +842,8 @@ Result Method::load(
     MemoryManager* memory_manager,
     EventTracer* event_tracer,
     const NamedDataMap* external_data_map,
-    const LoadBackendOptionsMap* backend_options) {
+    const LoadBackendOptionsMap* backend_options,
+    Span kernel_registry) {
   MemoryAllocator* temp_allocator = memory_manager->temp_allocator();
   if (temp_allocator == nullptr) {
     PlatformMemoryAllocator* platform_allocator =
@@ -844,7 +856,8 @@ Result Method::load(
     new (platform_allocator) PlatformMemoryAllocator();
     temp_allocator = platform_allocator;
   }
-  Method method(program, memory_manager, event_tracer, temp_allocator);
+  Method method(
+      program, memory_manager, event_tracer, temp_allocator, kernel_registry);
   ET_LOG(Debug, "Loading method: %s.", s_plan->name()->c_str());
   Error err = method.init(s_plan, external_data_map, backend_options);
   if (err != Error::Ok) {
diff --git a/runtime/executor/method.h b/runtime/executor/method.h
index a390341960e..834325f008c 100644
--- a/runtime/executor/method.h
+++ b/runtime/executor/method.h
@@ -23,6 +23,7 @@
 #include 
 #include 
 #include 
+#include 
 #include 
 
 // Forward declare flatbuffer types. This is a public header and must not
@@ -82,6 +83,7 @@ class Method final {
         merged_data_map_(std::move(rhs.merged_data_map_)),
         external_constants_(rhs.external_constants_),
         n_external_constants_(rhs.n_external_constants_),
+        kernel_registry_(rhs.kernel_registry_),
         init_state_(rhs.init_state_) {
     // Required: clear out fields that the dtor looks at, so that we don't free
     // anything twice.
@@ -331,7 +333,8 @@ class Method final {
       const Program* program,
       MemoryManager* memory_manager,
       EventTracer* event_tracer,
-      MemoryAllocator* temp_allocator)
+      MemoryAllocator* temp_allocator,
+      Span kernel_registry = {})
       : step_state_(),
         program_(program),
         memory_manager_(memory_manager),
@@ -348,6 +351,7 @@ class Method final {
         merged_data_map_(nullptr),
         external_constants_(nullptr),
         n_external_constants_(0),
+        kernel_registry_(kernel_registry),
         init_state_(InitializationState::Uninitialized) {}
 
   /// Static factory used by Program.
@@ -357,7 +361,8 @@ class Method final {
       MemoryManager* memory_manager,
       EventTracer* event_tracer,
       const NamedDataMap* named_data_map,
-      const LoadBackendOptionsMap* backend_options = nullptr);
+      const LoadBackendOptionsMap* backend_options = nullptr,
+      Span kernel_registry = {});
 
   /**
    * Initialize the method from its serialized representation.
@@ -403,6 +408,8 @@ class Method final {
   NamedData* external_constants_;
   size_t n_external_constants_ = 0;
 
+  Span kernel_registry_;
+
   InitializationState init_state_;
 
   /**
diff --git a/runtime/executor/program.cpp b/runtime/executor/program.cpp
index 31c02831448..4c0337e56d8 100644
--- a/runtime/executor/program.cpp
+++ b/runtime/executor/program.cpp
@@ -371,7 +371,8 @@ Result Program::load_method(
     MemoryManager* memory_manager,
     EventTracer* event_tracer,
     const NamedDataMap* named_data_map,
-    const LoadBackendOptionsMap* backend_options) const {
+    const LoadBackendOptionsMap* backend_options,
+    Span kernel_registry) const {
   EXECUTORCH_SCOPE_PROF("Program::load_method");
   internal::event_tracer_create_event_block(event_tracer, "Default");
   internal::EventTracerProfileMethodScope event_tracer_scope =
@@ -394,7 +395,8 @@ Result Program::load_method(
       memory_manager,
       event_tracer,
       named_data_map,
-      backend_options);
+      backend_options,
+      kernel_registry);
 }
 
 Result Program::method_meta(const char* method_name) const {
diff --git a/runtime/executor/program.h b/runtime/executor/program.h
index c4b96a241d7..e1208e52454 100644
--- a/runtime/executor/program.h
+++ b/runtime/executor/program.h
@@ -21,6 +21,7 @@
 #include 
 #include 
 #include 
+#include 
 #include 
 
 // Forward declare flatbuffer types. This is a public header and must not
@@ -151,7 +152,8 @@ class Program final {
       MemoryManager* memory_manager,
       EventTracer* event_tracer = nullptr,
       const NamedDataMap* named_data_map = nullptr,
-      const LoadBackendOptionsMap* backend_options = nullptr) const;
+      const LoadBackendOptionsMap* backend_options = nullptr,
+      Span kernel_registry = {}) const;
 
   /**
    * Gathers metadata for the named method.
diff --git a/runtime/kernel/test/operator_registry_test.cpp b/runtime/kernel/test/operator_registry_test.cpp
index ddc7207ed47..83a72482d70 100644
--- a/runtime/kernel/test/operator_registry_test.cpp
+++ b/runtime/kernel/test/operator_registry_test.cpp
@@ -440,6 +440,69 @@ TEST_F(OperatorRegistryTest, GetOpFunctionUsesProvidedKernelList) {
   EXPECT_EQ(run_kernel(*fallback_func), 50);
 }
 
+TEST_F(OperatorRegistryTest, ProvidedKernelListMissCanFallBackToGlobal) {
+  std::array buf{};
+  Error err = make_kernel_key(
+      {{ScalarType::Long, {0, 1, 2, 3}}}, buf.data(), buf.size());
+  ASSERT_EQ(err, Error::Ok);
+  KernelKey long_key = KernelKey(buf.data());
+
+  Kernel global_kernel = Kernel(
+      "test::provided_kernel_list_global_fallback",
+      KernelKey{},
+      [](KernelRuntimeContext& context, Span stack) {
+        (void)context;
+        *(stack[0]) = Scalar(50);
+      });
+  err = register_kernels({&global_kernel, 1});
+  ASSERT_EQ(err, Error::Ok);
+
+  Kernel scoped_kernel = Kernel(
+      "test::provided_kernel_list_global_fallback",
+      long_key,
+      [](KernelRuntimeContext& context, Span stack) {
+        (void)context;
+        *(stack[0]) = Scalar(100);
+      });
+  Span scoped_registry(&scoped_kernel, 1);
+
+  std::array dims = {0, 1, 2, 3};
+  auto dim_order_type = Span(dims.data(), dims.size());
+  std::array long_meta = {
+      TensorMeta(ScalarType::Long, dim_order_type)};
+  Span long_kernel_key(long_meta.data(), long_meta.size());
+
+  std::array float_meta = {
+      TensorMeta(ScalarType::Float, dim_order_type)};
+  Span float_kernel_key(float_meta.data(), float_meta.size());
+
+  auto run_kernel = [](OpFunction func) {
+    EValue value = Scalar(0);
+    std::array stack = {&value};
+    KernelRuntimeContext context{};
+    func(context, Span(stack.data(), stack.size()));
+    return value.toScalar().to();
+  };
+
+  Result scoped_func = get_op_function_from_registry(
+      "test::provided_kernel_list_global_fallback",
+      long_kernel_key,
+      scoped_registry);
+  ASSERT_EQ(scoped_func.error(), Error::Ok);
+  EXPECT_EQ(run_kernel(*scoped_func), 100);
+
+  Result scoped_miss = get_op_function_from_registry(
+      "test::provided_kernel_list_global_fallback",
+      float_kernel_key,
+      scoped_registry);
+  ASSERT_EQ(scoped_miss.error(), Error::OperatorMissing);
+
+  Result global_func = get_op_function_from_registry(
+      "test::provided_kernel_list_global_fallback", float_kernel_key);
+  ASSERT_EQ(global_func.error(), Error::Ok);
+  EXPECT_EQ(run_kernel(*global_func), 50);
+}
+
 TEST_F(OperatorRegistryTest, DoubleRegisterKernelsDies) {
   std::array buf_long_contiguous;
   Error err = make_kernel_key(
diff --git a/third-party/CMakeLists.txt b/third-party/CMakeLists.txt
index 7e49a18c42c..67da4833283 100644
--- a/third-party/CMakeLists.txt
+++ b/third-party/CMakeLists.txt
@@ -151,6 +151,13 @@ set(FLATCC_INSTALL
     CACHE BOOL ""
 )
 add_subdirectory(flatcc)
+if(CMAKE_C_COMPILER_ID STREQUAL "GNU"
+   AND CMAKE_C_COMPILER_VERSION VERSION_GREATER_EQUAL 15
+)
+  target_compile_options(
+    flatccrt PRIVATE -Wno-error=unterminated-string-initialization
+  )
+endif()
 # Unfortunately flatcc writes libs directly in to the source tree [1]. So to
 # ensure the target lib is created last, force flatcc_cli to build first.
 #
diff --git a/tools/cmake/preset/arm_baremetal.cmake b/tools/cmake/preset/arm_baremetal.cmake
index 882780ade1d..c12cc95233a 100644
--- a/tools/cmake/preset/arm_baremetal.cmake
+++ b/tools/cmake/preset/arm_baremetal.cmake
@@ -1,9 +1,30 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-set(CMAKE_INSTALL_PREFIX "${CMAKE_BINARY_DIR}")
+define_overridable_option(
+  EXECUTORCH_BAREMETAL_SKIP_INSTALL
+  "Skip emitting install/export rules when building bare-metal artifacts" BOOL
+  ON
+)
+
+if(EXECUTORCH_BAREMETAL_SKIP_INSTALL)
+  set(CMAKE_INSTALL_PREFIX "${CMAKE_BINARY_DIR}")
+  # Bare-metal builds consume the build tree directly. Keep the install target
+  # available (many docs/scripts still invoke it) but route the output back into
+  # the build tree so nothing is exported outside the repo.
+  unset(CMAKE_SKIP_INSTALL_RULES CACHE)
+  set(CMAKE_SKIP_INSTALL_RULES OFF)
+  set(CMAKE_SKIP_INSTALL_RULES
+      OFF
+      CACHE
+        BOOL
+        "Retain install() rules so docs/scripts can keep calling `--target install`"
+        FORCE
+  )
+endif()
+
 set_overridable_option(EXECUTORCH_BUILD_EXECUTOR_RUNNER OFF)
 set_overridable_option(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR OFF)
 set_overridable_option(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER OFF)
diff --git a/tools/cmake/preset/riscv64_linux.cmake b/tools/cmake/preset/riscv64_linux.cmake
index 32b891cd743..87894b63088 100644
--- a/tools/cmake/preset/riscv64_linux.cmake
+++ b/tools/cmake/preset/riscv64_linux.cmake
@@ -9,3 +9,29 @@ set_overridable_option(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL ON)
 set_overridable_option(EXECUTORCH_BUILD_DEVTOOLS ON)
 set_overridable_option(EXECUTORCH_ENABLE_BUNDLE_IO ON)
 set_overridable_option(EXECUTORCH_ENABLE_LOGGING ON)
+
+define_overridable_option(
+  EXECUTORCH_BUILD_RISCV_ETDUMP "Build etdump support for RISC-V" BOOL OFF
+)
+
+if("${EXECUTORCH_BUILD_RISCV_ETDUMP}")
+  set(EXECUTORCH_BUILD_DEVTOOLS ON)
+  set(EXECUTORCH_ENABLE_EVENT_TRACER ON)
+  set(FLATCC_ALLOW_WERROR OFF)
+else()
+  set(EXECUTORCH_ENABLE_EVENT_TRACER OFF)
+endif()
+
+if(EXECUTORCH_BUILD_XNNPACK)
+  if(CMAKE_COMPILER_IS_GNUCC AND CMAKE_C_COMPILER_VERSION VERSION_LESS 14)
+    message(FATAL_ERROR "XNNPACK requires GCC 14+ on riscv64")
+  endif()
+elseif(NOT DEFINED EXECUTORCH_BUILD_XNNPACK)
+  if(CMAKE_COMPILER_IS_GNUCC AND CMAKE_C_COMPILER_VERSION VERSION_GREATER_EQUAL
+                                 14
+  )
+    set(EXECUTORCH_BUILD_XNNPACK ON)
+  else()
+    message(NOTICE "XNNPACK requires GCC 14+ on riscv64")
+  endif()
+endif()
diff --git a/zephyr/samples/hello-executorch/README.md b/zephyr/samples/hello-executorch/README.md
index ab8022d5d62..45efe73d7d1 100644
--- a/zephyr/samples/hello-executorch/README.md
+++ b/zephyr/samples/hello-executorch/README.md
@@ -17,7 +17,7 @@ This is done in the example depending on the board you build for so if you build
 Set up FVP paths and macs used, this will also set `shutdown_on_eot` so the FVP auto stops after it has run the example.
 
 Config Zephyr Corstone300 FVP
-
+
 ```
 export FVP_ROOT=$PWD/modules/lib/executorch/examples/arm/arm-scratch/FVP-corstone300
 export ARMFVP_BIN_PATH=${FVP_ROOT}/models/Linux64_GCC-9.3
@@ -72,7 +72,7 @@ west build -b mps3/corstone300/fvp modules/lib/executorch/zephyr/samples/hello-e
 Set up FVP paths, libs and macs used, this will also set `shutdown_on_eot` so the FVP auto stops after it has run the example.
 
 Config Zephyr Corstone320 FVP
-
+
 ```
 export FVP_ROOT=$PWD/modules/lib/executorch/examples/arm/arm-scratch/FVP-corstone320
 export ARMFVP_BIN_PATH=${FVP_ROOT}/models/Linux64_GCC-9.3
diff --git a/zephyr/samples/mv2-ethosu/README.md b/zephyr/samples/mv2-ethosu/README.md
index a05b46835b2..70877abcafb 100644
--- a/zephyr/samples/mv2-ethosu/README.md
+++ b/zephyr/samples/mv2-ethosu/README.md
@@ -18,36 +18,66 @@ The model classifies a static RGB test input tensor with shape `[1, 3, 224, 224]
 > ~3 MiB, so the build will link but FVP execution will fail at runtime.
 > Use Corstone-320 (below) for end-to-end MV2 inference.
 
+### Setup FVP paths
+
+Set up FVP paths and macs used, this will also set `shutdown_on_eot` so the FVP auto stops after it has run the example.
+
+Config Zephyr Corstone300 FVP
+
+```
+export FVP_ROOT=$PWD/modules/lib/executorch/examples/arm/arm-scratch/FVP-corstone300
+export ARMFVP_BIN_PATH=${FVP_ROOT}/models/Linux64_GCC-9.3
+export ARMFVP_EXTRA_FLAGS="-C mps3_board.uart0.shutdown_on_eot=1 -C ethosu.num_macs=128"
+```
+
 ### Export the model
 
 Export a quantized INT8 MobileNetV2 model with Ethos-U55 delegation:
 
-
+
 ```
 python -m modules.lib.executorch.backends.arm.scripts.aot_arm_compiler --model_name=mv2_untrained --quantize --delegate --target=ethos-u55-128 --output=mv2_u55_128.pte
 ```
 
 ### Build (link-check only)
 
-
+
 ```
 west build -b mps3/corstone300/fvp modules/lib/executorch/zephyr/samples/mv2-ethosu -- -DET_PTE_FILE_PATH=mv2_u55_128.pte
 ```
 
 ## Corstone-320 FVP (Ethos-U85)
 
+### Setup FVP paths
+
+Set up FVP paths, libs and macs used, this will also set `shutdown_on_eot` so the FVP auto stops after it has run the example.
+
+These FVP command-line options are passed through the `ARMFVP_EXTRA_FLAGS`
+environment variable. The sample does not set `ARMFVP_FLAGS` in its
+`CMakeLists.txt`; the base `ARMFVP_FLAGS` come from the selected Zephyr board's
+`board.cmake`.
+
+Config Zephyr Corstone320 FVP
+
+```
+export FVP_ROOT=$PWD/modules/lib/executorch/examples/arm/arm-scratch/FVP-corstone320
+export ARMFVP_BIN_PATH=${FVP_ROOT}/models/Linux64_GCC-9.3
+export LD_LIBRARY_PATH=${FVP_ROOT}/python/lib:${ARMFVP_BIN_PATH}:${LD_LIBRARY_PATH}
+export ARMFVP_EXTRA_FLAGS="-C mps4_board.uart0.shutdown_on_eot=1 -C mps4_board.subsystem.ethosu.num_macs=256"
+```
+
 ### Export the model
 
 Export a quantized INT8 MobileNetV2 model with Ethos-U85 delegation:
 
-
+
 ```
 python -m modules.lib.executorch.backends.arm.scripts.aot_arm_compiler --model_name=mv2_untrained --quantize --delegate --target=ethos-u85-256 --output=mv2_u85_256.pte
 ```
 
 ### Build
 
-
+
 ```
 west build -b mps4/corstone320/fvp modules/lib/executorch/zephyr/samples/mv2-ethosu -- -DET_PTE_FILE_PATH=mv2_u85_256.pte
 ```