diff --git a/.ci/docker/ci_commit_pins/pytorch.txt b/.ci/docker/ci_commit_pins/pytorch.txt index 401a0594d98..6c3fe42ddf3 100644 --- a/.ci/docker/ci_commit_pins/pytorch.txt +++ b/.ci/docker/ci_commit_pins/pytorch.txt @@ -1 +1 @@ -release/2.13 +release/2.14 diff --git a/.ci/docker/common/install_docs_reqs.sh b/.ci/docker/common/install_docs_reqs.sh index ea54d90523e..2794ffb8fc9 100755 --- a/.ci/docker/common/install_docs_reqs.sh +++ b/.ci/docker/common/install_docs_reqs.sh @@ -20,7 +20,9 @@ if [ -n "$BUILD_DOCS" ]; then apt-get update apt-get install -y --no-install-recommends yarn - yarn global add katex --prefix /usr/local + # katex 0.18.5 requires commander@15 / node >= 22.12; pin to the last + # release compatible with the node 16 installed above + yarn global add katex@0.18.4 --prefix /usr/local sudo apt-get -y install doxygen diff --git a/.ci/docker/common/install_pytorch.sh b/.ci/docker/common/install_pytorch.sh index 0ac5e79cf4a..e51a8886afd 100755 --- a/.ci/docker/common/install_pytorch.sh +++ b/.ci/docker/common/install_pytorch.sh @@ -76,21 +76,31 @@ install_pytorch_and_domains() { # the image compiler cannot satisfy. The venv inherits the image's # site-packages, so PyTorch still builds against the same numpy. # - # Keep the list in sync with pytorch/pyproject.toml [build-system].requires. + # Keep in sync with pytorch/pyproject.toml [build-system].requires. local build_venv=/tmp/pytorch-build-venv rm -rf "${build_venv}" conda_run python -m venv --system-site-packages "${build_venv}" + # No pip cmake: scikit-build-core would prefer it over the image's, and it + # searches site-packages, where MKL and libomp are not. conda_run "${build_venv}/bin/pip" install build "scikit-build-core>=1.0" \ - "setuptools>=77.0.0,<82" "cmake>=3.27,<4" ninja "packaging>=24.2" \ - "typing-extensions>=4.10.0" pyyaml six - conda_run "${build_venv}/bin/python" -m build --wheel --no-isolation + ninja "packaging>=24.2" "typing-extensions>=4.10.0" pyyaml six numpy + # These images have no module scanner, and nothing here uses modules. + conda_run env CMAKE_CXX_SCAN_FOR_MODULES=OFF \ + "${build_venv}/bin/python" -m build --wheel --no-isolation rm -rf "${build_venv}" pip_install "$(echo dist/*.whl)" + # A build with no BLAS succeeds silently. Run from / to import the wheel. + (cd / && conda_run python -c " +import torch +assert torch._C.has_lapack, 'built without LAPACK' +torch.linalg.qr(torch.randn(4, 4)) +") + # Grab the pinned audio and vision commits from PyTorch TORCHAUDIO_VERSION=release/2.11 export TORCHAUDIO_VERSION - TORCHVISION_VERSION=release/0.28 + TORCHVISION_VERSION=release/0.29 export TORCHVISION_VERSION install_domains diff --git a/.ci/scripts/build-qnn-windows-msvc.ps1 b/.ci/scripts/build-qnn-windows-msvc.ps1 index a69316be6e3..c94daf453e7 100644 --- a/.ci/scripts/build-qnn-windows-msvc.ps1 +++ b/.ci/scripts/build-qnn-windows-msvc.ps1 @@ -4,13 +4,35 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +param( + [switch]$SkipX86Windows, + [switch]$SkipArm64Windows +) + $ErrorActionPreference = "Stop" -conda create --yes --quiet -n et python=3.12 -conda activate et +if ($SkipX86Windows -eq $SkipArm64Windows) { + Write-Error "Specify exactly one of -SkipArm64Windows (to build x86_64) or -SkipX86Windows (to build arm64)." + exit 1 +} -# Install CI requirements -pip install -r .ci/docker/requirements-ci.txt +if ($SkipX86Windows) { + $ArchLabel = "arm64" + py -3.12 -m venv et + .\et\Scripts\Activate.ps1 + # ARM64 prebuilt wheels are not available for some Python modules. + # To unblock the build process, only a minimal set of dependencies + # is installed via pip. `PyYAML`/`torch` for ExecuTorch's codegen, + # `requests` for download_qnn_sdk.py. + pip install pyyaml requests + pip install torch --index-url https://download.pytorch.org/whl/cpu +} else { + $ArchLabel = "x86_64" + conda create --yes --quiet -n et python=3.12 + conda activate et + # Install CI requirements + pip install -r .ci/docker/requirements-ci.txt +} # Provision the QNN SDK if ($env:QNN_SDK_ROOT -and (Test-Path -Path $env:QNN_SDK_ROOT)) { @@ -44,27 +66,26 @@ if (-not (Test-Path -Path (Join-Path $env:QNN_SDK_ROOT "include\QNN"))) { exit 1 } -# Test x86_64 Windows host build -.\backends\qualcomm\scripts\build.ps1 -SkipArm64Windows -Release +if ($SkipArm64Windows) { + .\backends\qualcomm\scripts\build.ps1 -SkipArm64Windows -Release +} else { + .\backends\qualcomm\scripts\build.ps1 -SkipX86Windows -Release +} -$x86Artifacts = @( - "build-x86_64-windows\backends\qualcomm\Release\PyQnnManagerAdaptor*.pyd", - "build-x86_64-windows\backends\qualcomm\Release\qnn_executorch_backend.dll", - "build-x86_64-windows\examples\qualcomm\executor_runner\Release\qnn_executor_runner.exe" +$Artifacts = @( + "build-$ArchLabel-windows\backends\qualcomm\Release\qnn_executorch_backend.dll", + "build-$ArchLabel-windows\examples\qualcomm\executor_runner\Release\qnn_executor_runner.exe" ) -foreach ($artifact in $x86Artifacts) { +if ($SkipArm64Windows) { + # Only run PyQnnManagerAdaptor validation for x86_64 Windows artifacts, + # since AOT is not fully supported on native ARM64 Windows. + $Artifacts += "build-x86_64-windows\backends\qualcomm\Release\PyQnnManagerAdaptor*.pyd" +} +foreach ($artifact in $Artifacts) { if (-not (Get-ChildItem -Path $artifact -ErrorAction SilentlyContinue)) { - Write-Error "ERROR: x86_64 artifact not found: $artifact" + Write-Error "ERROR: $ArchLabel artifact not found: $artifact" exit 1 } } -# The ARM64 MSVC toolchain is currently not installed in the Windows CI -# environment. Enabling this build configuration results in build failures -# due to the missing ARM64 platform definition. -# `.\backends\qualcomm\scripts\build.ps1 -SkipX86Windows -Release` -# -# Temporarily disable this build option until ARM64 MSVC support is available -# in CI. The configuration can be re-enabled in a future update. - -Write-Host "PASSED: QNN backend Windows MSVC build completed" +Write-Host "PASSED: QNN backend Windows MSVC build ($ArchLabel) completed" diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index dfa4d7dc38b..223dd82951c 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -294,6 +294,26 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then fi fi +# Downloads and compiler caches go in scratch dirs outside OUTPUT_DIR because the CI job +# templates upload OUTPUT_DIR even when the job fails. A failed export also empties +# OUTPUT_DIR, but only if it started out empty, so a local run with output_dir=. cannot +# delete the checkout. Scratch goes under RUNNER_TEMP, which the runner wipes between +# jobs, with a fallback for containers where RUNNER_TEMP is not writable. +LOCAL_MODEL_DIR=$(mktemp -d "${RUNNER_TEMP:-/tmp}/model_XXXXXX" 2>/dev/null || mktemp -d) +SCRATCH_DIRS=("$LOCAL_MODEL_DIR") +OUTPUT_DIR_WAS_EMPTY=0 +[ -n "$(ls -A -- "$OUTPUT_DIR" 2>/dev/null)" ] || OUTPUT_DIR_WAS_EMPTY=1 +cleanup() { + local rc=$? + set +e + rm -rf "${SCRATCH_DIRS[@]}" + if [ "$rc" -ne 0 ] && [ "$OUTPUT_DIR_WAS_EMPTY" = 1 ] && [ -d "$OUTPUT_DIR" ]; then + echo "Export failed with exit code $rc; removing partial output from ${OUTPUT_DIR}" + (cd -- "$OUTPUT_DIR" && find . -mindepth 1 -delete) + fi +} +trap cleanup EXIT + echo "::group::Export $MODEL_NAME" if [ -n "$EXTRA_PIP" ]; then @@ -386,8 +406,7 @@ fi if [ "$MODEL_NAME" = "voxtral_realtime" ]; then pip install safetensors huggingface_hub - # Download model weights from HuggingFace (requires HF_TOKEN for gated model) - LOCAL_MODEL_DIR="${OUTPUT_DIR}/model_weights" + # Download model weights outside OUTPUT_DIR to avoid uploading on failure (requires HF_TOKEN for gated model) python -c "from huggingface_hub import snapshot_download; snapshot_download('${HF_MODEL}', local_dir='${LOCAL_MODEL_DIR}')" # Per-component quantization flags @@ -437,7 +456,6 @@ if [ "$MODEL_NAME" = "voxtral_realtime" ]; then fi # Copy tokenizer from downloaded model weights cp "$LOCAL_MODEL_DIR/tekken.json" "${OUTPUT_DIR}/tekken.json" - rm -rf "$LOCAL_MODEL_DIR" ls -al "${OUTPUT_DIR}" echo "::endgroup::" exit 0 @@ -448,12 +466,11 @@ if [ "$MODEL_NAME" = "qwen3_5_moe" ]; then pip install safetensors huggingface_hub pip install -r examples/models/qwen3_5_moe/requirements.txt - # Download prequantized model outside OUTPUT_DIR to avoid uploading on failure - LOCAL_MODEL_DIR=$(mktemp -d) INDUCTOR_CACHE=$(mktemp -d "${RUNNER_TEMP:-/tmp}/inductor_cache_XXXXXX") INDUCTOR_TMPDIR=$(mktemp -d "${RUNNER_TEMP:-/tmp}/tmpdir_XXXXXX") - trap 'rm -rf "$LOCAL_MODEL_DIR" "$INDUCTOR_CACHE" "$INDUCTOR_TMPDIR"' EXIT + SCRATCH_DIRS+=("$INDUCTOR_CACHE" "$INDUCTOR_TMPDIR") + # Download prequantized model outside OUTPUT_DIR to avoid uploading on failure 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 @@ -521,10 +538,9 @@ fi if [ "$MODEL_NAME" = "muse_glimmer" ]; then pip install safetensors huggingface_hub gguf - LOCAL_MODEL_DIR=$(mktemp -d) INDUCTOR_CACHE=$(mktemp -d "${RUNNER_TEMP:-/tmp}/inductor_cache_XXXXXX") INDUCTOR_TMPDIR=$(mktemp -d "${RUNNER_TEMP:-/tmp}/tmpdir_XXXXXX") - trap 'rm -rf "$LOCAL_MODEL_DIR" "$INDUCTOR_CACHE" "$INDUCTOR_TMPDIR"' EXIT + SCRATCH_DIRS+=("$INDUCTOR_CACHE" "$INDUCTOR_TMPDIR") case "$QUANT_NAME" in kquant-17gb) @@ -588,14 +604,13 @@ fi if [ "$MODEL_NAME" = "gemma4_31b" ]; then pip install safetensors huggingface_hub gguf - # Download GGUF + tokenizer outside OUTPUT_DIR to avoid uploading on failure. - # The unsloth GGUF repo ships the .gguf but no tokenizer.json, so the tokenizer - # is fetched from the (non-GGUF) unsloth/gemma-4-31B-it repo. - LOCAL_MODEL_DIR=$(mktemp -d) INDUCTOR_CACHE=$(mktemp -d "${RUNNER_TEMP:-/tmp}/inductor_cache_XXXXXX") INDUCTOR_TMPDIR=$(mktemp -d "${RUNNER_TEMP:-/tmp}/tmpdir_XXXXXX") - trap 'rm -rf "$LOCAL_MODEL_DIR" "$INDUCTOR_CACHE" "$INDUCTOR_TMPDIR"' EXIT + SCRATCH_DIRS+=("$INDUCTOR_CACHE" "$INDUCTOR_TMPDIR") + # Download GGUF + tokenizer outside OUTPUT_DIR to avoid uploading on failure. + # The unsloth GGUF repo ships the .gguf but no tokenizer.json, so the tokenizer + # is fetched from the (non-GGUF) unsloth/gemma-4-31B-it repo. GGUF_FILE="gemma-4-31B-it-Q4_K_M.gguf" python -c "from huggingface_hub import hf_hub_download; hf_hub_download('unsloth/gemma-4-31B-it-GGUF', '${GGUF_FILE}', local_dir='${LOCAL_MODEL_DIR}')" python -c "from huggingface_hub import hf_hub_download; hf_hub_download('unsloth/gemma-4-31B-it', 'tokenizer.json', local_dir='${LOCAL_MODEL_DIR}')" diff --git a/.ci/scripts/gather_test_models.py b/.ci/scripts/gather_test_models.py index 65fefc4073d..db0f2671594 100755 --- a/.ci/scripts/gather_test_models.py +++ b/.ci/scripts/gather_test_models.py @@ -17,23 +17,23 @@ from examples.xnnpack import MODEL_NAME_TO_OPTIONS, QuantType DEFAULT_RUNNERS = { - "linux": "linux.2xlarge", + "linux": "mt-l-x86iavx512-8-64", "macos": "macos-m1-stable", } CUSTOM_RUNNERS = { "linux": { # This one runs OOM on smaller runner, the root cause is unclear (T163016365) - "w2l": "linux.4xlarge.memory", - "ic4": "linux.4xlarge.memory", - "resnet50": "linux.4xlarge.memory", - "llava": "linux.4xlarge.memory", - "llama3_2_vision_encoder": "linux.4xlarge.memory", - "llama3_2_text_decoder": "linux.4xlarge.memory", + "w2l": "mt-l-x86iavx512-16-128", + "ic4": "mt-l-x86iavx512-16-128", + "resnet50": "mt-l-x86iavx512-16-128", + "llava": "mt-l-x86iavx512-16-128", + "llama3_2_vision_encoder": "mt-l-x86iavx512-16-128", + "llama3_2_text_decoder": "mt-l-x86iavx512-16-128", # This one causes timeout on smaller runner, the root cause is unclear (T161064121) - "dl3": "linux.4xlarge.memory", - "emformer_join": "linux.4xlarge.memory", - "emformer_predict": "linux.4xlarge.memory", - "phi_4_mini": "linux.4xlarge.memory", + "dl3": "mt-l-x86iavx512-16-128", + "emformer_join": "mt-l-x86iavx512-16-128", + "emformer_predict": "mt-l-x86iavx512-16-128", + "phi_4_mini": "mt-l-x86iavx512-16-128", } } @@ -146,7 +146,7 @@ def export_models_for_ci() -> dict[str, dict]: "build-tool": "buck2", "model": "mv3", "backend": backend, - "runner": "linux.2xlarge", + "runner": "mt-l-x86iavx512-8-64", "timeout": DEFAULT_TIMEOUT, } models["include"].append(record) @@ -175,7 +175,7 @@ def export_models_for_ci() -> dict[str, dict]: "build-tool": "cmake", "model": name, "backend": backend, - "runner": DEFAULT_RUNNERS.get(target_os, "linux.2xlarge"), + "runner": DEFAULT_RUNNERS.get(target_os, "mt-l-x86iavx512-8-64"), "timeout": DEFAULT_TIMEOUT, } diff --git a/.ci/scripts/pytest-parallelism.sh b/.ci/scripts/pytest-parallelism.sh new file mode 100755 index 00000000000..97d50061cf3 --- /dev/null +++ b/.ci/scripts/pytest-parallelism.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# 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. + +# Pin pytest-xdist's worker count, but only when the container is allowed less +# CPU than the machine it landed on. `auto` asks psutil for the machine's +# physical cores, which inside an OSDC pod is the whole node, so the workers +# exhaust the pod's memory. nproc honours the pod's cpuset, which is what +# pytorch relies on for OMP_NUM_THREADS on the same fleet. +# +# Left alone when unconstrained. `auto` already discounts hyperthreads there, +# while nproc counts them, so overriding it would double the workers and halve +# the memory each one gets. + +if [[ -z "${PYTEST_XDIST_AUTO_NUM_WORKERS:-}" ]] && command -v nproc >/dev/null 2>&1; then + cpus_allowed="$(nproc)" + cpus_installed="$(nproc --all)" + echo "pytest-xdist: ${cpus_allowed} of ${cpus_installed} CPUs available" + if [[ "${cpus_allowed}" -lt "${cpus_installed}" ]]; then + export PYTEST_XDIST_AUTO_NUM_WORKERS="${cpus_allowed}" + echo "PYTEST_XDIST_AUTO_NUM_WORKERS=${PYTEST_XDIST_AUTO_NUM_WORKERS}" + fi +fi diff --git a/.ci/scripts/setup-samsung-linux-deps.sh b/.ci/scripts/setup-samsung-linux-deps.sh index a8339b16c7d..548f30bf161 100644 --- a/.ci/scripts/setup-samsung-linux-deps.sh +++ b/.ci/scripts/setup-samsung-linux-deps.sh @@ -165,6 +165,7 @@ install_enn_backend() { echo "NDK will be installed/used at: ${ANDROID_NDK_ROOT}" bash backends/samsung/build.sh --build all + bash examples/samsung/build.sh export EXECUTORCH_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" export PYTHONPATH="${PYTHONPATH:-}:${EXECUTORCH_ROOT}/.." diff --git a/.ci/scripts/test-cuda-build.sh b/.ci/scripts/test-cuda-build.sh index e717718be66..07e2a50b1a0 100755 --- a/.ci/scripts/test-cuda-build.sh +++ b/.ci/scripts/test-cuda-build.sh @@ -49,10 +49,16 @@ import executorch print('SUCCESS: ExecuTorch imported successfully') " + python -m pip check + # Test CUDA availability and show details - python -c " + EXPECTED_CUDA_VERSION="$cuda_version" python -c " try: + import os import torch + assert torch.version.cuda == os.environ['EXPECTED_CUDA_VERSION'], ( + torch.version.cuda, os.environ['EXPECTED_CUDA_VERSION'] + ) print('INFO: PyTorch version:', torch.__version__) print('INFO: CUDA available:', torch.cuda.is_available()) @@ -68,7 +74,8 @@ try: x = torch.randn(10, 10).to(device) y = torch.randn(10, 10).to(device) z = torch.mm(x, y) - print('SUCCESS: CUDA tensor operation completed on device:', z.device) + torch.testing.assert_close(z.cpu(), x.cpu() @ y.cpu()) + print('SUCCESS: CUDA tensor operation matched CPU on device:', z.device) print('INFO: Result tensor shape:', z.shape) print('SUCCESS: ExecuTorch CUDA integration verified') diff --git a/.ci/scripts/test-rocm-aoti.sh b/.ci/scripts/test-rocm-aoti.sh index 0f4ac9d826f..00592bc4bf6 100644 --- a/.ci/scripts/test-rocm-aoti.sh +++ b/.ci/scripts/test-rocm-aoti.sh @@ -7,7 +7,7 @@ set -euo pipefail -ROCM_VERSION="${ROCM_VERSION:-7.1}" +ROCM_VERSION="${ROCM_VERSION:-7.2}" ROCM_PATH="${ROCM_PATH:-/opt/rocm}" PYTORCH_ROCM_INDEX="${PYTORCH_ROCM_INDEX:-https://download.pytorch.org/whl/test/rocm${ROCM_VERSION}}" TORCHAO_ROCM_WHEEL_BASE="${TORCHAO_ROCM_WHEEL_BASE:-https://download.pytorch.org/whl/nightly/rocm${ROCM_VERSION}}" diff --git a/.ci/scripts/test-rocm-voxtral.sh b/.ci/scripts/test-rocm-voxtral.sh index cd0562c3808..eb00c1efc73 100644 --- a/.ci/scripts/test-rocm-voxtral.sh +++ b/.ci/scripts/test-rocm-voxtral.sh @@ -7,7 +7,7 @@ set -euo pipefail -ROCM_VERSION="${ROCM_VERSION:-7.1}" +ROCM_VERSION="${ROCM_VERSION:-7.2}" ROCM_PATH="${ROCM_PATH:-/opt/rocm}" EXPECTED_ROCM_ARCH="${EXPECTED_ROCM_ARCH:-gfx950}" EXPECTED_WARP_SIZE="${EXPECTED_WARP_SIZE:-64}" diff --git a/.ci/scripts/test_backend.sh b/.ci/scripts/test_backend.sh index 45e3322a0c8..068d5adb260 100755 --- a/.ci/scripts/test_backend.sh +++ b/.ci/scripts/test_backend.sh @@ -7,6 +7,9 @@ # LICENSE file in the root directory of this source tree. set -eux +# Cap pytest-xdist's `auto` workers to the container's CPU quota. +source .ci/scripts/pytest-parallelism.sh + SUITE=$1 FLOW=$2 ARTIFACT_DIR=$3 @@ -151,11 +154,13 @@ if [[ "$FLOW" == *nxp* ]]; then export NXP_RUNNER_PATH="$(pwd)/examples/nxp/executor_runner/build/nxp_executor_runner" fi -GOLDEN_DIR="${ARTIFACT_DIR}/golden-artifacts" -export GOLDEN_ARTIFACTS_DIR="${GOLDEN_DIR}" - EXIT_CODE=0 -PYTEST_ARGS=(-c /dev/null -n auto) +# An Ethos-U failure captures a few hundred thousand lines of Vela operator +# listings, and the runner agent throws System.OutOfMemoryException processing +# a step that size, taking the whole job down before pytest can report. The +# reason for each failure is in its exception message and traceback, which are +# unaffected. +PYTEST_ARGS=(-c /dev/null -n auto --show-capture=no) if [[ ${#PYTEST_RETRY_ARGS[@]} -gt 0 ]]; then PYTEST_ARGS+=("${PYTEST_RETRY_ARGS[@]}") fi diff --git a/.ci/scripts/test_lora.sh b/.ci/scripts/test_lora.sh index 102347a08fd..91c78ae5f3e 100644 --- a/.ci/scripts/test_lora.sh +++ b/.ci/scripts/test_lora.sh @@ -28,11 +28,14 @@ cmake_build_llama_runner() { } cleanup_files() { - echo "Deleting downloaded and generated files" - rm -rf "${HF_QWEN_PATH}/" - rm -rf "${HF_ADAPTER_PATH}/" - rm -rf *.pte *.ptd - rm result*.txt + # Only what this script generated. HF_QWEN_PATH and HF_ADAPTER_PATH point + # inside the huggingface_hub cache: on OSDC that is a shared read-only mount, + # so removing them failed the job after the tests had already passed, and + # anywhere else it discards a cache entry the next job wants. A teardown also + # must not fail a run whose tests passed. + echo "Deleting generated files" + rm -rf ./*.pte ./*.ptd || true + rm -f result*.txt || true } matches_base_response_prefix() { diff --git a/.ci/scripts/test_lora_multimethod.sh b/.ci/scripts/test_lora_multimethod.sh index f0b30bd4be1..afcde517cd6 100755 --- a/.ci/scripts/test_lora_multimethod.sh +++ b/.ci/scripts/test_lora_multimethod.sh @@ -28,11 +28,14 @@ cmake_build_llama_runner() { } cleanup_files() { - echo "Deleting downloaded and generated files" - rm -rf "${HF_QWEN_PATH}/" - rm -rf "${HF_ADAPTER_PATH}/" - rm -rf *.pte - rm -f result*.txt + # Only what this script generated. HF_QWEN_PATH and HF_ADAPTER_PATH point + # inside the huggingface_hub cache: on OSDC that is a shared read-only mount, + # so removing them failed the job after the tests had already passed, and + # anywhere else it discards a cache entry the next job wants. A teardown also + # must not fail a run whose tests passed. + echo "Deleting generated files" + rm -rf ./*.pte || true + rm -f result*.txt || true } matches_base_response_prefix() { diff --git a/.ci/scripts/tests/test_cu134_dependencies.py b/.ci/scripts/tests/test_cu134_dependencies.py new file mode 100644 index 00000000000..663cdb3b3bf --- /dev/null +++ b/.ci/scripts/tests/test_cu134_dependencies.py @@ -0,0 +1,264 @@ +# 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. + +import ast +import importlib.util +import os +import subprocess +import sys +import unittest +from pathlib import Path +from unittest.mock import patch + +from packaging.requirements import Requirement + +ROOT = Path(__file__).resolve().parents[3] + + +def load_module(name): + spec = importlib.util.spec_from_file_location(name, ROOT / f"{name}.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class TestCu134Dependencies(unittest.TestCase): + def setUp(self): + self.utils = load_module("install_utils") + self.modules = patch.dict(sys.modules, {"install_utils": self.utils}) + self.modules.start() + self.addCleanup(self.modules.stop) + self.installer = load_module("install_requirements") + + def install_commands(self, cuda, machine="x86_64", nightly=True, system="Linux"): + self.utils.determine_torch_url.cache_clear() + with ( + patch.dict(os.environ, {}, clear=True), + patch.object( + self.utils, + "_get_cuda_version", + return_value=cuda, + side_effect=RuntimeError("no nvcc") if cuda is None else None, + ), + patch.object(self.installer.platform, "machine", return_value=machine), + patch.object(self.installer.platform, "system", return_value=system), + patch.object(self.installer.sys, "platform", "linux"), + patch.object(self.installer.subprocess, "run") as run, + ): + self.installer.install_requirements(nightly) + self.installer.install_optional_example_requirements(nightly) + return [call.args[0] for call in run.call_args_list] + + def test_all_install_steps_preserve_exact_cu134_selection(self): + for machine, ao_variant in (("x86_64", "cu134"), ("aarch64", "cpu")): + with self.subTest(machine=machine): + commands = self.install_commands((13, 4), machine) + self.assertEqual(len(commands), 4) + expected = { + "torch==2.14.0.dev20260810+cu134", + "torchvision==0.29.0.dev20260811+cu134", + "torchaudio==2.11.0.dev20260811+cu134", + f"torchao==0.19.0.dev20260811+{ao_variant}", + } + for index, command in enumerate(commands): + required = ( + expected + if index >= 2 + else { + requirement + for requirement in expected + if requirement.startswith(("torch==", "torchao==")) + } + ) + self.assertTrue(required.issubset(command), command) + if index < 2: + self.assertFalse( + any( + arg.startswith(("torchvision", "torchaudio")) + for arg in command + ) + ) + self.assertIn( + "https://download.pytorch.org/whl/nightly/cu134", command + ) + self.assertNotIn( + "https://download.pytorch.org/whl/test/cu134", command + ) + self.assertNotIn("--no-deps", command) + if machine == "aarch64": + self.assertIn( + "https://download.pytorch.org/whl/nightly/cpu", command + ) + + def test_other_cuda_trains_keep_existing_pins(self): + for cuda in ((12, 6), (13, 0), (13, 2)): + for machine in ("x86_64", "aarch64"): + with self.subTest(cuda=cuda, machine=machine): + core, local, domains, examples = self.install_commands( + cuda, machine + ) + self.assertIn("torch==2.14.0", core) + self.assertIn("torchao==0.18.0.dev20260729", core) + self.assertIn("torchvision==0.29.0", domains) + self.assertIn("torchaudio==2.11.0", domains) + self.assertFalse(any("==" in arg for arg in local)) + self.assertFalse(any("==" in arg for arg in examples)) + + def test_source_pinned_torch_is_not_replaced(self): + for cuda in ((13, 2), (13, 4)): + with self.subTest(cuda=cuda): + core, _, domains, _ = self.install_commands(cuda, nightly=False) + self.assertIn("torch", core) + self.assertNotIn("torch==2.14.0.dev20260810+cu134", core) + self.assertIn("torchvision", domains) + self.assertIn("torchaudio", domains) + + def test_no_cuda_keeps_default_pins(self): + core, _, domains, _ = self.install_commands(None) + self.assertIn("torch==2.14.0", core) + self.assertIn("torchao==0.18.0.dev20260729", core) + self.assertIn("torchvision==0.29.0", domains) + self.assertIn("https://download.pytorch.org/whl/test/cpu", core) + + def test_windows_does_not_select_cu134(self): + core, _, domains, _ = self.install_commands((13, 4), system="Windows") + self.assertIn("torch==2.14.0", core) + self.assertIn("torchvision==0.29.0", domains) + self.assertIn("https://download.pytorch.org/whl/test/cpu", core) + + def test_failure_is_not_retried_with_another_cuda_train(self): + with ( + patch.object(self.utils, "_get_cuda_version", return_value=(13, 4)), + patch.object(self.installer.platform, "system", return_value="Linux"), + patch.object( + self.installer.subprocess, + "run", + side_effect=subprocess.CalledProcessError(1, "pip"), + ) as run, + ): + with self.assertRaises(subprocess.CalledProcessError): + self.installer.install_requirements(True) + self.assertEqual(run.call_count, 1) + + def torchao_requirement(self): + path = ROOT / "setup.py" + tree = ast.parse(path.read_text()) + function = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_torchao_requirement" + ) + namespace = { + "__file__": str(path), + "Path": Path, + "importlib": importlib, + "sys": sys, + "install_utils": self.utils, + } + exec( + compile(ast.Module(body=[function], type_ignores=[]), str(path), "exec"), + namespace, + ) + return namespace["_torchao_requirement"]() + + def test_package_install_preserves_source_pinned_torchao(self): + with patch.dict(sys.modules, {"install_requirements": self.installer}): + package_installer = load_module("install_executorch") + for machine in ("x86_64", "aarch64"): + with self.subTest(machine=machine): + self.utils.determine_torch_url.cache_clear() + with ( + patch.dict(os.environ, {}, clear=True), + patch.object(self.utils, "_get_cuda_version", return_value=(13, 4)), + patch.object( + self.installer.platform, "machine", return_value=machine + ), + patch.object( + self.installer.platform, "system", return_value="Linux" + ), + patch.object(self.installer.sys, "platform", "linux"), + patch.object( + sys, + "argv", + ["install_executorch", "--use-pt-pinned-commit", "--minimal"], + ), + patch.object( + package_installer, "python_is_compatible", return_value=True + ), + patch.object(package_installer, "check_and_update_submodules"), + patch.object(self.installer.subprocess, "run") as run, + ): + package_installer.main([]) + commands = [call.args[0] for call in run.call_args_list] + metadata = Requirement(self.torchao_requirement()) + self.assertEqual(len(commands), 3) + self.assertIn(".", commands[-1]) + core = commands[0] + torchao = Requirement( + next(arg for arg in core if arg.startswith("torchao==")) + ) + version = next(iter(torchao.specifier)).version + self.assertIn(version, metadata.specifier) + self.assertIn("torch", core) + self.assertFalse(any(arg.startswith("torch==") for arg in core)) + + def test_cu134_keeps_explicit_torchao_source_build(self): + with patch.dict(sys.modules, {"install_requirements": self.installer}): + package_installer = load_module("install_executorch") + for source_flag in ( + "EXECUTORCH_BUILD_KERNELS_TORCHAO", + "TORCHAO_BUILD_EXPERIMENTAL_MPS", + ): + with self.subTest(source_flag=source_flag): + self.utils.determine_torch_url.cache_clear() + with ( + patch.dict(os.environ, {source_flag: "1"}, clear=True), + patch.object(self.utils, "_get_cuda_version", return_value=(13, 4)), + patch.object( + self.installer.platform, "system", return_value="Linux" + ), + patch.object(self.installer.sys, "platform", "linux"), + patch.object(sys, "argv", ["install_executorch"]), + patch.object( + package_installer, "python_is_compatible", return_value=True + ), + patch.object(package_installer, "check_and_update_submodules"), + patch.object(self.installer.subprocess, "run") as run, + ): + package_installer.main([]) + metadata = Requirement(self.torchao_requirement()) + commands = [call.args[0] for call in run.call_args_list] + self.assertEqual(len(commands), 5) + self.assertIn("third-party/ao", commands[1]) + self.assertIn(".", commands[2]) + for command in commands: + self.assertFalse( + any(arg.startswith("torchao==") for arg in command) + ) + self.assertIn("torch==2.14.0.dev20260810+cu134", commands[-1]) + self.assertIn("0.18.0+git03ca489", metadata.specifier) + + def test_wheel_torchao_bound_matches_selected_train(self): + for cuda, expected in ( + ((13, 4), "torchao>=0.19.0.dev20260811,<0.20"), + ((13, 2), "torchao>=0.18.0.dev20260729,<0.19"), + (None, "torchao>=0.18.0.dev20260729,<0.19"), + ): + self.utils.determine_torch_url.cache_clear() + with ( + patch.object( + self.utils, + "_get_cuda_version", + return_value=cuda, + side_effect=RuntimeError("no nvcc") if cuda is None else None, + ), + patch.object(self.installer.platform, "system", return_value="Linux"), + ): + self.assertEqual(self.torchao_requirement(), expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/.ci/scripts/tests/test_cuda_workflow.py b/.ci/scripts/tests/test_cuda_workflow.py new file mode 100644 index 00000000000..0ce37edc917 --- /dev/null +++ b/.ci/scripts/tests/test_cuda_workflow.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. + +import os +import subprocess +import sys +import unittest +from pathlib import Path + +import yaml + + +ROOT = Path(__file__).resolve().parents[3] +WORKFLOW = yaml.safe_load((ROOT / ".github" / "workflows" / "cuda.yml").read_text()) + + +def _all_keys(value): + if isinstance(value, dict): + for key, child in value.items(): + yield key + yield from _all_keys(child) + elif isinstance(value, list): + for child in value: + yield from _all_keys(child) + + +def _model_quant(entry): + return (entry["model"]["repo"], entry["model"]["name"], entry["quant"]) + + +class CudaWorkflowTest(unittest.TestCase): + def test_build_matrix_preserves_existing_cuda_versions(self): + job = WORKFLOW["jobs"]["test-cuda-builds"] + self.assertEqual( + job["strategy"]["matrix"]["cuda-version"], ["12.6", "13.0", "13.4"] + ) + self.assertEqual(job["with"]["gpu-arch-version"], "${{ matrix.cuda-version }}") + + def test_cuda134_driver_uses_matching_workflow_and_action_revision(self): + job = WORKFLOW["jobs"]["test-cuda-builds"] + workflow, revision = job["uses"].split("@") + self.assertEqual( + workflow, "pytorch/test-infra/.github/workflows/linux_job_v2.yml" + ) + self.assertRegex(revision, r"^[0-9a-f]{40}$") + self.assertEqual(job["with"]["test-infra-ref"], revision) + self.assertEqual( + job["with"]["driver-version"], + "${{ matrix.cuda-version == '13.4' && '615.71.09' || '580.65.06' }}", + ) + self.assertEqual( + job["with"]["driver-download-url"], + "${{ matrix.cuda-version == '13.4' && " + "'https://download.nvidia.com/XFree86/Linux-x86_64/615.71.09/" + "NVIDIA-Linux-x86_64-615.71.09.run' || '' }}", + ) + + def test_cuda134_runtime_update_precedes_build_and_propagates_failure(self): + script = WORKFLOW["jobs"]["test-cuda-builds"]["with"]["script"] + stubs = """ +conda() { printf 'CONDA %s\n' "$*"; return "$CONDA_STATUS"; } +source() { printf 'BUILD %s\n' "$*"; } +""" + for version, conda_status in ( + ("12.6", 0), + ("13.0", 0), + ("13.4", 0), + ("13.4", 1), + ): + with self.subTest(version=version, conda_status=conda_status): + result = subprocess.run( + [ + "bash", + "-c", + stubs + script.replace("${{ matrix.cuda-version }}", version), + ], + env={**os.environ, "CONDA_STATUS": str(conda_status)}, + capture_output=True, + text=True, + ) + self.assertEqual(result.returncode, conda_status, result.stderr) + expected = [] + if version == "13.4": + expected.append( + "CONDA install -y -n base -c conda-forge " + "libstdcxx-ng=16.2.0 libgcc-ng=16.2.0" + ) + if conda_status == 0: + expected.append(f"BUILD .ci/scripts/test-cuda-build.sh {version}") + self.assertEqual(result.stdout.splitlines(), expected) + + def test_cuda_probe_checks_the_result_and_torch_train(self): + script = (ROOT / ".ci/scripts/test-cuda-build.sh").read_text() + probes = [block.split('\n"', 1)[0] for block in script.split('python -c "')[1:]] + probe = next(block for block in probes if "import torch" in block) + fake_torch = """ +import os +import sys +from types import SimpleNamespace +class Tensor: + device = 'cuda' + shape = (10, 10) + def to(self, device): + return self + def cpu(self): + return self + def __matmul__(self, other): + return self +def assert_close(actual, expected): + print('RESULT CHECKED') + assert os.environ['INVALID_CUDA_RESULT'] == '0', 'CUDA result mismatch' +sys.modules['torch'] = SimpleNamespace( + __version__='test', version=SimpleNamespace(cuda='13.4'), + cuda=SimpleNamespace( + is_available=lambda: True, device_count=lambda: 1, + current_device=lambda: 0, get_device_name=lambda: 'test', + ), + device=lambda name: name, randn=lambda *args: Tensor(), + mm=lambda x, y: Tensor(), + testing=SimpleNamespace(assert_close=assert_close), +) +""" + for expected, invalid_result, succeeds in ( + ("13.4", "0", True), + ("13.0", "0", False), + ("13.4", "1", False), + ): + with self.subTest(expected=expected, invalid_result=invalid_result): + result = subprocess.run( + [sys.executable, "-c", fake_torch + probe], + env={ + **os.environ, + "EXPECTED_CUDA_VERSION": expected, + "INVALID_CUDA_RESULT": invalid_result, + }, + capture_output=True, + text=True, + ) + self.assertEqual(result.returncode == 0, succeeds, result.stdout) + if succeeds: + self.assertIn("RESULT CHECKED", result.stdout) + + def test_pybind_runs_inline_for_the_expected_matrix_cells(self): + job = WORKFLOW["jobs"]["test-model-cuda-e2e"] + matrix = job["strategy"]["matrix"] + pybind_rows = [row for row in matrix["include"] if "pybind_model" in row] + + actual = { + (*_model_quant(row), row["pybind_model"], row["pybind_quantized"]) + for row in pybind_rows + } + expected = { + ( + "google", + "gemma-3-4b-it", + "quantized-int4-tile-packed", + "gemma3-4b", + True, + ), + ( + "Qwen", + "Qwen3-0.6B", + "non-quantized", + "qwen3-0.6b", + False, + ), + ( + "Qwen", + "Qwen3-0.6B", + "quantized-int4-tile-packed", + "qwen3-0.6b", + True, + ), + } + self.assertEqual(expected, actual) + + excluded = {_model_quant(row) for row in matrix["exclude"]} + active = { + (model["repo"], model["name"], quant) + for model in matrix["model"] + for quant in matrix["quant"] + } - excluded + self.assertTrue({_model_quant(row) for row in pybind_rows} <= active) + + script = job["with"]["script"] + self.assertIn('if [ -n "${{ matrix.pybind_model }}" ]', script) + self.assertIn("test_huggingface_optimum_model.py", script) + self.assertIn("--run_only", script) + self.assertGreaterEqual(script.count('"${MODEL_DIR}"'), 2) + + def test_model_e2e_does_not_transfer_artifacts(self): + self.assertNotIn("test-cuda-pybind", WORKFLOW["jobs"]) + keys = set(_all_keys(WORKFLOW["jobs"]["test-model-cuda-e2e"])) + self.assertNotIn("upload-artifact", keys) + self.assertNotIn("download-artifact", keys) diff --git a/.ci/scripts/tests/test_filter_cuda_matrix.py b/.ci/scripts/tests/test_filter_cuda_matrix.py index d9b39024e8e..3f21762c44f 100644 --- a/.ci/scripts/tests/test_filter_cuda_matrix.py +++ b/.ci/scripts/tests/test_filter_cuda_matrix.py @@ -10,8 +10,12 @@ # with what the project can publish. Two of its comments record past bugs it now guards against, and # a regression in any of them would surface only as a broken release, so each gate is pinned here. +import contextlib import importlib.util +import io import json +import os +import subprocess import unittest from pathlib import Path from unittest import mock @@ -21,29 +25,21 @@ ROOT = Path(__file__).resolve().parents[3] -def _load_filter(): - """Load the script by path, since .github/scripts is not an importable package.""" - path = ROOT / ".github" / "scripts" / "filter_cuda_matrix.py" - spec = importlib.util.spec_from_file_location("filter_cuda_matrix", path) +def _load_module(name, path): + spec = importlib.util.spec_from_file_location(name, path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module -FILTER = _load_filter() +FILTER = _load_module( + "filter_cuda_matrix", ROOT / ".github" / "scripts" / "filter_cuda_matrix.py" +) +INSTALL_UTILS = _load_module("install_utils", ROOT / "install_utils.py") def _full_matrix(): - """Every supported python and CUDA pair. - - The filter refuses anything less: one gate rejects a matrix that would leave a CUDA train - unpublished, another rejects a missing python and CUDA combination. Built from the module's own - lists so it cannot go stale when either grows. - - That also means it shrinks when either list shrinks, and every gate keeps passing. Measured: - deleting cu132 and 3.13 from the filter left all sixteen cases here green. TestPublishedSets - below is what notices that, so this fixture does not have to. - """ + """Every supported pair; TestPublishedSets separately guards against shrinking the lists.""" return { "include": [ {"python_version": python, "desired_cuda": cuda} @@ -165,9 +161,6 @@ def test_unsupported_cuda_is_dropped(self): class TestGates(unittest.TestCase): def _exit_message(self, matrix, limit="false", extra=None): """The stderr text of the gate that fired, so a case can name which one it hit.""" - import contextlib - import io - argv = ["--matrix", json.dumps(matrix), "--limit-pr-builds", limit] + ( extra or [] ) @@ -191,18 +184,14 @@ def test_unparseable_matrix_exits_nonzero(self): FILTER.main(argv) self.assertNotEqual(raised.exception.code, 0) - def test_absent_train_exits_nonzero(self): - # A supported train the generator offers nothing for would publish no wheel at all. + def test_absent_train_is_skipped_not_fatal(self): + # A supported train the generator offers nothing for is one PyTorch stopped shipping. The + # release skips it and publishes the rest, so one dropped train cannot take the others down. # - # Patching the supported list rather than deleting rows, because deleting every row for one - # train also creates missing combinations, so both gates fire and the test cannot tell which - # one it exercised. Adding an extra supported train makes it absent while every offered - # combination stays complete. - # These two gates cannot be separated by input: any matrix leaving a train absent also - # leaves every combination for that train missing, so the later gate always catches what the - # earlier one would. Measured. So each gate gets its own case, and the case asserts on the - # message rather than only on a nonzero exit, which is the only way to tell them apart. + # Offering every train but the last leaves that train absent while every offered combination + # stays complete, which is exactly the shape of an upstream drop. offered = FILTER.SUPPORTED_CUDA_VERSIONS[:-1] + dropped = FILTER.SUPPORTED_CUDA_VERSIONS[-1] matrix = { "include": [ {"python_version": python, "desired_cuda": cuda} @@ -210,14 +199,58 @@ def test_absent_train_exits_nonzero(self): for cuda in offered ] } - message = self._exit_message(matrix) - self.assertIn("publish no wheel for that CUDA version", message) + stdout = io.StringIO() + stderr = io.StringIO() + with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): + FILTER.main(["--matrix", json.dumps(matrix)]) + emitted = json.loads(stdout.getvalue()) + self.assertIn("the generator offered no row", stderr.getvalue()) + self.assertIn(dropped, stderr.getvalue()) + published = sorted({row["desired_cuda"] for row in emitted["include"]}) + self.assertEqual(published, sorted(offered)) + self.assertNotIn(dropped, published) + + def test_dropped_train_still_publishes_the_others(self): + # Losing cu126 from the generator must not block the remaining supported trains. + if "cu126" not in FILTER.SUPPORTED_CUDA_VERSIONS: + self.skipTest("cu126 is not a published train") + survivors = [c for c in FILTER.SUPPORTED_CUDA_VERSIONS if c != "cu126"] + matrix = { + "include": [ + {"python_version": python, "desired_cuda": cuda} + for python in FILTER.SUPPORTED_PYTHON_VERSIONS + for cuda in survivors + ] + } + emitted = _emitted(_run(matrix)) + published = sorted({row["desired_cuda"] for row in emitted["include"]}) + self.assertEqual(published, sorted(survivors)) + self.assertNotIn("cu126", published) + # Every survivor keeps all its pythons, so what publishes is complete, just narrower. + self.assertEqual( + len(emitted["include"]), + len(survivors) * len(FILTER.SUPPORTED_PYTHON_VERSIONS), + ) def test_missing_combination_exits_nonzero(self): + # A train that IS offered but missing one python is a real break, not an upstream drop: the + # release would ship that train incomplete. Deleting one row from a full matrix leaves its + # train present, so this exercises the incomplete-train gate rather than the skip above. matrix = _full_matrix() del matrix["include"][0] message = self._exit_message(matrix) - self.assertIn("combination(s) produced no row", message) + self.assertIn("incomplete train", message) + + def test_offered_train_with_only_unsupported_pythons_exits_nonzero(self): + for cuda in FILTER.SUPPORTED_CUDA_VERSIONS: + with self.subTest(cuda=cuda): + matrix = _full_matrix() + for row in matrix["include"]: + if row["desired_cuda"] == cuda: + row["python_version"] = "3.15" + message = self._exit_message(matrix) + self.assertIn("incomplete train", message) + self.assertIn(f"3.10/{cuda}", message) def test_jetpack_not_published_exits_nonzero(self): # Refused explicitly rather than allowed to fall through to an empty result, so the reason a @@ -247,7 +280,66 @@ class TestPublishedSets(unittest.TestCase): """ def test_published_cuda_versions(self): - self.assertEqual(FILTER.SUPPORTED_CUDA_VERSIONS, ["cu126", "cu130", "cu132"]) + self.assertEqual( + FILTER.SUPPORTED_CUDA_VERSIONS, ["cu126", "cu130", "cu132", "cu134"] + ) + + def test_published_cuda_versions_are_supported_by_the_installer(self): + supported = { + f"cu{major}{minor}" + for major, minor in INSTALL_UTILS.SUPPORTED_CUDA_VERSIONS + } + self.assertLessEqual(set(FILTER.SUPPORTED_CUDA_VERSIONS), supported) + + def test_supported_toolkits_select_the_matching_torch_index(self): + base_url = "https://download.pytorch.org/whl/nightly" + self.addCleanup(INSTALL_UTILS._get_cuda_version.cache_clear) + self.addCleanup(INSTALL_UTILS.determine_torch_url.cache_clear) + for major, minor in INSTALL_UTILS.SUPPORTED_CUDA_VERSIONS: + with self.subTest(cuda=(major, minor)): + INSTALL_UTILS._get_cuda_version.cache_clear() + INSTALL_UTILS.determine_torch_url.cache_clear() + detected = subprocess.CompletedProcess( + args=[], + returncode=0, + stdout=f"Cuda compilation tools, release {major}.{minor}, V{major}.{minor}.0", + ) + with mock.patch.object( + INSTALL_UTILS.platform, "system", return_value="Linux" + ), mock.patch.object( + INSTALL_UTILS.subprocess, "run", return_value=detected + ): + self.assertEqual( + INSTALL_UTILS.determine_torch_url(base_url), + f"{base_url}/cu{major}{minor}", + ) + self.assertTrue(INSTALL_UTILS.is_cuda_available()) + + def test_published_cuda_versions_have_gpu_architectures(self): + script = ROOT / ".ci" / "scripts" / "wheel" / "cuda_arch_list.sh" + for machine in ("x86_64", "aarch64"): + for cuda in FILTER.SUPPORTED_CUDA_VERSIONS: + with self.subTest(machine=machine, cuda=cuda): + result = subprocess.run( + [ + "bash", + "-c", + 'uname() { printf "%s\\n" "$MACHINE"; }; ' + 'source "$1"; executorch_cuda_arch_list', + "bash", + str(script), + ], + env={ + **os.environ, + "MACHINE": machine, + "CU_VERSION": cuda, + "EXECUTORCH_BUILD_CUDA": "1", + }, + capture_output=True, + text=True, + ) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn("8.0", result.stdout.split()) def test_published_python_versions(self): self.assertEqual( diff --git a/.ci/scripts/tests/test_is_aten_target.py b/.ci/scripts/tests/test_is_aten_target.py new file mode 100644 index 00000000000..8a468b2c77b --- /dev/null +++ b/.ci/scripts/tests/test_is_aten_target.py @@ -0,0 +1,209 @@ +# 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 the ATen detection in the Buck macro layer. + +Here rather than as a build test, because the open source Buck build cannot query the +Vulkan backend at all, so the targets this decision matters most for are never built by +CI. The decision itself is a pure function of a target's keyword arguments, so it can be +exercised directly. + +The detection had a real defect that this covers. A target can name a third-party +dependency either by its short name, which lands in ``external_deps``, or through +``external_dep_location``, which hands back the resolved label and lands in an ordinary +``deps`` list. Only the first was checked, so the Vulkan operator tests, which use the +second, compiled at the older standard against headers that need the newer one. +""" + +import types +import unittest +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[3] +MACROS = ( + REPO_ROOT / "shim_et" / "xplat" / "executorch" / "build" / "runtime_wrapper.bzl" +) + + +class _Select: + __slots__ = ("values",) + + def __init__(self, values: dict[str, list[str]]) -> None: + self.values = values + + +# Representative dependency-map results, including platform-specific values and +# aliases shared by ATen and non-ATen dependency names. +RESOLVED = { + "c10": _Select( + { + "ovr_config//os:android": ["fbsource//xplat/caffe2/c10:c10"], + "DEFAULT": ["fbsource//xplat/caffe2/c10:c10_ovrsource"], + } + ), + "libtorch": ["//third-party:libtorch"], + "libtorch_python": ["//third-party:libtorch_python"], + "torch-core-cpp": ["//third-party:libtorch"], + "gtest_aten": ["fbsource//third-party/googletest:gtest"], + "gmock_aten": ["fbsource//third-party/googletest:gmock"], +} +FALLTHROUGH = "@fallthrough@" + + +def _load_macro_function(name: str): + """Execute the real macro text, with the little of Starlark it uses shimmed.""" + text = MACROS.read_text() + start = text.index("def _has_pytorch_dep") + end = text.index("def _cxx_library_common", start) + + env = types.SimpleNamespace( + EXTERNAL_DEP_FALLTHROUGH=FALLTHROUGH, + resolve_external_dep=lambda name: RESOLVED.get(name, FALLTHROUGH), + ) + + def _apply(obj, function): + """Stand-in for selects.apply: run over each list the object holds.""" + if isinstance(obj, _Select): + return _Select({key: function(value) for key, value in obj.values.items()}) + if isinstance(obj, dict): + return {key: function(value) for key, value in obj.items()} + return function(obj) + + namespace = { + # Starlark's type() returns a name, and the macros compare against "string". + "type": lambda value: "string" if isinstance(value, str) else "other", + "env": env, + "selects": types.SimpleNamespace(apply=_apply), + } + exec(compile(text[start:end], str(MACROS), "exec"), namespace) + return namespace[name] + + +def _load_is_aten_target(): + return _load_macro_function("_is_aten_target") + + +class TestIsAtenTarget(unittest.TestCase): + def setUp(self) -> None: + self.is_aten_target = _load_is_aten_target() + + def test_resolved_label_in_deps(self) -> None: + """The Vulkan operator tests name libtorch this way.""" + self.assertTrue( + self.is_aten_target( + { + "name": "compute_graph_op_tests_bin", + "deps": [ + "//third-party/googletest:gtest_main", + "//executorch/backends/vulkan:vulkan_graph_runtime", + "//third-party:libtorch", + ], + } + ) + ) + + def test_resolved_label_in_exported_deps(self) -> None: + self.assertTrue( + self.is_aten_target( + {"name": "some_lib", "exported_deps": ["//third-party:libtorch"]} + ) + ) + + def test_resolved_label_returned_inside_a_select(self) -> None: + for dep in [ + "fbsource//xplat/caffe2/c10:c10", + "fbsource//xplat/caffe2/c10:c10_ovrsource", + ]: + with self.subTest(dep=dep): + self.assertTrue( + self.is_aten_target({"name": "some_lib", "deps": [dep]}) + ) + + def test_short_name_in_external_deps(self) -> None: + for name in RESOLVED: + with self.subTest(name=name): + self.assertTrue( + self.is_aten_target({"name": "some_test", "external_deps": [name]}) + ) + + def test_plain_target_is_not_aten(self) -> None: + self.assertFalse( + self.is_aten_target( + { + "name": "op_add_test", + "deps": [ + "//executorch/runtime/core:core", + "//third-party/googletest:gtest_main", + ], + } + ) + ) + + def test_plain_gtest_target_is_not_aten(self) -> None: + self.assertFalse( + self.is_aten_target( + { + "name": "some_test", + "deps": ["fbsource//third-party/googletest:gtest"], + } + ) + ) + + def test_executorch_label_alone_is_not_aten(self) -> None: + """Every label under the project contains the word torch.""" + self.assertFalse( + self.is_aten_target( + {"name": "evalue_test", "deps": ["//executorch/test/utils:utils"]} + ) + ) + + def test_resolved_label_inside_a_select(self) -> None: + """A dep list can be a select, which cannot be walked like a list.""" + self.assertTrue( + self.is_aten_target( + { + "name": "some_test", + "deps": { + "DEFAULT": ["//third-party:libtorch"], + "ovr_config//os:windows": [], + }, + } + ) + ) + + def test_select_without_aten_is_not_aten(self) -> None: + self.assertFalse( + self.is_aten_target( + { + "name": "some_test", + "deps": {"DEFAULT": ["//executorch/runtime/core:core"]}, + } + ) + ) + + +class TestPatchTestCompilerFlags(unittest.TestCase): + def test_inherits_platform_standard(self) -> None: + patch_test_compiler_flags = _load_macro_function("_patch_test_compiler_flags") + for name in ["some_test", "some_aten_test"]: + with self.subTest(name=name): + kwargs = { + "name": name, + "compiler_flags": ["-DTEST"], + "fbobjc_compiler_flags": ["-DAPPLE_TEST"], + } + + result = patch_test_compiler_flags(kwargs) + + self.assertFalse( + any(flag.startswith("-std=") for flag in result["compiler_flags"]) + ) + self.assertEqual(["-DAPPLE_TEST"], result["fbobjc_compiler_flags"]) + self.assertIn("-Wno-error", result["compiler_flags"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/.ci/scripts/tests/test_runtime_path_filter.py b/.ci/scripts/tests/test_runtime_path_filter.py new file mode 100644 index 00000000000..74af1b03066 --- /dev/null +++ b/.ci/scripts/tests/test_runtime_path_filter.py @@ -0,0 +1,331 @@ +# 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 the runtime search path filter packaging applies to shipped libraries. + +Here rather than in the wheel checks, because the decision under test is a pure function of one +string. The wheel checks can only see it after a full wheel build, and only on the platform that +built one, so a filter that dropped the wrong entry reached the published artifact before anything +ran that would notice. + +Two properties are covered, and they pull in opposite directions, which is why both are needed. +The filter must drop the MKL arch directories torch's exported link interface leaves anchored at +the filesystem root, and it must keep the absolute torch directory that is a library's only route +to torch when no relative one was recorded. A filter that satisfies either alone is wrong: the +first way ships unusable paths, the second way stops the extensions importing. + +The functions are read out of setup.py rather than restated, so the test exercises what ships. +setup.py calls setup() at module scope, so it is loaded by compiling the definitions this needs +instead of importing it, which would exit during setuptools argument parsing. +""" + +import ast +import importlib.util +import re +from pathlib import Path, PurePosixPath + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[3] + +# Compiled from setup.py, so a change to the filter is exercised here rather than duplicated. +_WANTED = ( + "_MKL_ARCH_DIRECTORIES", + "_is_cuda_toolkit_directory", + "_is_unresolved_math_library_directory", + "_is_usable_runtime_path", +) + + +def _setup_source() -> str: + """setup.py's text, decoded as UTF-8. + + The encoding is named because `read_text()` defaults to the locale's, and both files this test + parses contain non-ASCII characters. On a Windows runner that resolves to a code page, which + mangles them, and the mangled text is what gets parsed. + """ + return (REPO_ROOT / "setup.py").read_text(encoding="utf-8") + + +def _setup_namespace() -> dict: + """The runtime path helpers from setup.py, compiled without running setup().""" + tree = ast.parse(_setup_source()) + wanted = [] + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + if node.name in _WANTED: + wanted.append(node) + elif isinstance(node, ast.Assign): + names = [t.id for t in node.targets if isinstance(t, ast.Name)] + if any(name in _WANTED for name in names): + wanted.append(node) + found = set() + for node in wanted: + found.add( + node.name + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + else next(t.id for t in node.targets if isinstance(t, ast.Name)) + ) + missing = sorted(set(_WANTED) - found) + assert not missing, ( + f"setup.py no longer defines {missing} at module scope, so this test would silently " + "check nothing. Update the names here to match." + ) + namespace = {"re": re, "PurePosixPath": PurePosixPath} + exec( + compile(ast.Module(body=wanted, type_ignores=[]), "", "exec"), + namespace, + ) + return namespace + + +@pytest.fixture(scope="module") +def setup_helpers() -> dict: + return _setup_namespace() + + +# setup.py's own arch names, read at import so the parametrized cases below are driven by the +# shipped constant rather than a second copy of it. Adding an arch to setup.py then extends both +# the reject and the accept cases, which is what keeps packaging and the release check in step. +_MKL_ARCH_DIRECTORIES = _setup_namespace()["_MKL_ARCH_DIRECTORIES"] + +# What the linker records when MKL's prefix resolves empty, leaving its arch subdirectory +# concatenated onto nothing. Two of the three name a Windows layout, in a Linux wheel. +UNRESOLVED_MATH_DIRECTORIES = tuple(f"/lib/{arch}" for arch in _MKL_ARCH_DIRECTORIES) + + +@pytest.mark.parametrize("entry", UNRESOLVED_MATH_DIRECTORIES) +def test_drops_math_directories_with_an_empty_prefix(entry, setup_helpers): + assert setup_helpers["_is_unresolved_math_library_directory"](entry) is True + # A trailing separator is the same directory, which the path type normalises on its own. Asserted + # because patchelf prints entries as recorded, so that spelling can genuinely arrive. + assert setup_helpers["_is_unresolved_math_library_directory"](entry + "/") is True + + +@pytest.mark.parametrize( + "entry", + [ + # A real MKL installation spells the same arch directory below a prefix, and that + # directory genuinely exists, so dropping it would break a library resolving through it. + "/opt/intel/mkl/lib/intel64", + "/opt/intel/oneapi/mkl/latest/lib/intel64", + "/usr/lib/intel64", + "/home/user/lib/intel64", + # The arch name as a parent rather than as the entry itself. + "/lib/intel64/extra", + # Ordinary system directories, which differ from the bad entries only in the last part. + "/lib", + "/lib64", + "/lib/x86_64-linux-gnu", + # Relative entries are decided before this predicate is reached, but it must not claim + # one, or a wheel's own hop into a directory named intel64 would be dropped. + "$ORIGIN/../../lib", + "$ORIGIN/lib/intel64", + ], +) +def test_keeps_directories_that_name_a_real_prefix(entry, setup_helpers): + assert setup_helpers["_is_unresolved_math_library_directory"](entry) is False + # Also through the production predicate, not just the narrow one. Asserting only the narrow + # predicate let _is_usable_runtime_path start rejecting an ordinary system directory with the + # whole suite green, since nothing checked that these entries actually survive the filter. + assert setup_helpers["_is_usable_runtime_path"](entry, True, True) is True + + +# Read off the pybindings extension in the published x86_64 CPU nightly, in the recorded order. The +# MKL block sits after five relative hops and before one. On this wheel the entries are simply dead, +# since nothing resolves through them, so dropping them costs a few wasted lookups at load time and +# removes paths a user cannot have. On the CUDA wheel the same block precedes the hop into the CUDA +# runtime, which is where the ordering matters. +SHIPPED_RUNTIME_PATH = [ + "$ORIGIN/../../../torch/lib", + "$ORIGIN/../../src/executorch/lib", + "$ORIGIN/../../backends/qualcomm", + "$ORIGIN/../../lib", + "$ORIGIN/../../../lib64", + "/lib/intel64", + "/lib/intel64_win", + "/lib/win-x64", + "$ORIGIN/../../backends/cuda", +] + + +def _relative_torch_route_predicate(): + """setup.py's own has_relative_torch_route expression, lifted out of its function. + + Read rather than restated, because it decides the third argument to the filter under test. + A copy here would let the test keep passing after that expression changed, which is the one + failure a regression guard must not have. + """ + tree = ast.parse(_setup_source()) + for node in ast.walk(tree): + if not ( + isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "has_relative_torch_route" + for target in node.targets + ) + ): + continue + # The expression iterates a name bound in setup.py's own scope, so it is rebound to this + # function's argument by compiling it as the body of a one-argument lambda. + source = ast.unparse(node.value) + iterated = "for entry in found" + assert iterated in source, ( + f"setup.py computes has_relative_torch_route as {source!r}, which no longer iterates " + "the name this test rebinds. Update the rebinding rather than leaving it a no-op." + ) + return eval( + f"lambda entries: {source.replace(iterated, 'for entry in entries')}" + ) + raise AssertionError( + "setup.py no longer computes has_relative_torch_route, so this test would pass the " + "filter an argument the shipped code never produces." + ) + + +def _filtered(entries, setup_helpers): + is_usable = setup_helpers["_is_usable_runtime_path"] + has_relative_torch_route = _relative_torch_route_predicate()(entries) + return [ + entry + for entry in entries + # True is safe_to_drop_toolkit_paths: this wheel ships no CUDA, so an absolute toolkit path + # in it names only the build machine. Packaging derives the same value from the built tree. + if is_usable(entry, True, has_relative_torch_route) + ] + + +def test_shipped_library_keeps_no_absolute_entry(setup_helpers): + kept = _filtered(SHIPPED_RUNTIME_PATH, setup_helpers) + assert [entry for entry in kept if entry.startswith("/")] == [] + + +def test_shipped_library_keeps_every_relative_hop(setup_helpers): + # Asserted separately from the absence of absolute entries, because a filter that dropped + # everything would satisfy that one while leaving the library unable to find its siblings. + kept = _filtered(SHIPPED_RUNTIME_PATH, setup_helpers) + assert kept == [ + entry for entry in SHIPPED_RUNTIME_PATH if not entry.startswith("/") + ] + + +def test_keeps_the_absolute_torch_directory_when_it_is_the_only_route(setup_helpers): + # The case that stops this being a blanket "drop everything absolute": an extension links + # torch and, with no relative route recorded, reaches it only through the directory the + # linker found it in. Dropping that would stop the extension importing. + entries = [ + "/opt/conda/envs/py_3.12/lib/python3.12/site-packages/torch/lib", + "/lib/intel64", + ] + assert _filtered(entries, setup_helpers) == [entries[0]] + + +def test_drops_the_absolute_torch_directory_when_a_relative_route_exists(setup_helpers): + # The other side of the same rule, and the only case that exercises the route expression at all. + # Without a relative entry in the list the route is False whatever that expression says, so the + # test above cannot tell a correct expression from an inverted one. + entries = [ + "$ORIGIN/../../../torch/lib", + "/opt/conda/envs/py_3.12/lib/python3.12/site-packages/torch/lib", + "/lib/intel64", + ] + assert _filtered(entries, setup_helpers) == [entries[0]] + + +_MATH_DIRECTORY_REASON = "a maths library directory whose prefix resolved empty" +_BUILD_DIRECTORY_REASON = "inside a build of this project" +_UNREACHABLE_REASON = "an absolute directory the wheel has a relative route to" + + +def _release_check_decision(): + """The release check's own per-entry decision, imported rather than replayed. + + Loaded as a module so the unit test exercises the function the wheel check calls, rather than + that function's source text. Reading the text let the rejection be deleted outright. + """ + path = REPO_ROOT / ".ci" / "scripts" / "wheel" / "test_shared_libraries.py" + spec = importlib.util.spec_from_file_location("_release_check_under_test", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module._unusable_runtime_path_kind + + +@pytest.mark.parametrize("arch", _MKL_ARCH_DIRECTORIES) +def test_release_check_rejects_the_math_directories(arch): + # The check and packaging are separate code, so a wheel built without patchelf keeps these + # entries and only the check would catch it. Driven off setup.py's own constant, and asserting + # the REASON rather than merely a rejection: the check rejects every absolute path it does not + # recognise, so an arch it was never taught about would otherwise pass through the catch-all. + decide = _release_check_decision() + assert ( + decide(f"/lib/{arch}", "_C.cpython-312-x86_64-linux-gnu.so") + == _MATH_DIRECTORY_REASON + ) + + +def test_release_check_rejects_each_kind_for_its_own_reason(): + # One assertion per rejecting branch, by reason, so deleting any single branch fails here. + # Asserting only "not None" let the build-directory branch and the catch-all each be removed + # on their own with every test green. + decide = _release_check_decision() + assert decide("/home/u/pip-out/lib", "_C.so") == _BUILD_DIRECTORY_REASON + assert decide("/opt/rocm/lib", "_C.so") == _UNREACHABLE_REASON + assert decide("", "_C.so") is not None + + +def test_release_check_rejects_a_build_directory_before_the_allowlist(): + # Order is load bearing and nothing else holds it. A torch directory inside a CI worker tree + # must be rejected, which only happens because the build-directory branch runs before the + # suffix allowlist gets to accept the /torch/lib ending. + decide = _release_check_decision() + entry = "/home/ec2-user/actions-runner/_work/executorch/pytorch/torch/lib" + assert decide(entry, "_C.so") == _BUILD_DIRECTORY_REASON + + +@pytest.mark.parametrize("arch", _MKL_ARCH_DIRECTORIES) +def test_release_check_still_accepts_a_real_mkl_installation(arch, setup_helpers): + # The two must agree on every arch in the shared constant. Packaging KEEPS a prefixed one, + # because the environment provides it, so a check that rejected it would fail a wheel packaging + # deliberately allowed and the builder could not satisfy both. Parametrized off the constant, so + # adding an arch to setup.py without teaching the check about it fails here. + entry = f"/opt/intel/mkl/lib/{arch}" + assert setup_helpers["_is_usable_runtime_path"](entry, True, True) is True + assert ( + _release_check_decision()(entry, "_C.cpython-312-x86_64-linux-gnu.so") is None + ) + + +def test_the_wheel_scan_consults_the_classifier(): + # The unit tests above call the classifier directly, and the wheel scan is the only thing that + # applies it to a real library. That scan needs an installed wheel, so it cannot run here; + # what is checkable is the wiring, and severing it left all tests green. Asserted on the AST so + # a rename or an accidental deletion fails rather than silently disabling the enforcement. + path = REPO_ROOT / ".ci" / "scripts" / "wheel" / "test_shared_libraries.py" + tree = ast.parse(path.read_text(encoding="utf-8")) + scan = next( + ( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) + and node.name == "test_no_absolute_runtime_paths" + ), + None, + ) + assert scan is not None, "the wheel scan this check enforces no longer exists" + called = { + node.func.id + for node in ast.walk(scan) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } + assert "_unusable_runtime_path_kind" in called, ( + "test_no_absolute_runtime_paths no longer calls _unusable_runtime_path_kind, so the wheel " + "scan would report every shipped library clean while these unit tests still pass." + ) + + +def test_release_check_accepts_a_relative_entry(): + # A relative hop is the normal case and must never be rejected, whatever the absolute rules do. + assert _release_check_decision()("$ORIGIN/../../lib", "_C.so") is None diff --git a/.ci/scripts/tests/test_wheel_test_modules.py b/.ci/scripts/tests/test_wheel_test_modules.py new file mode 100644 index 00000000000..39c1d3f250f --- /dev/null +++ b/.ci/scripts/tests/test_wheel_test_modules.py @@ -0,0 +1,867 @@ +# 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 the test modules the full wheel drops. + +The wheel used to carry every test file in the repository, about 9.7 MB of Python that nothing +in an installed wheel can reach. A test case is only ever loaded by pytest from a path in the +checkout, never through the installed name, so shipping it buys nothing. + +Shared helpers are the opposite. The suites here import each other by installed name, for +example `from executorch.backends.arm.test.tester.test_pipeline import TosaPipelineFP`, so a +helper has to ship or collection breaks. That is why the keep set is computed from the import +graph and not from file names: `test_pipeline.py` and `test_add.py` are indistinguishable by +name and only one of them can go. + +setup.py is read rather than imported. It calls setup() at module scope, so importing it under a +test runner hands setup() the runner's own arguments and the session dies on an invalid command +name. +""" + +import ast +import functools +import os +import re +import shutil +import subprocess +import tempfile +import unittest +from pathlib import Path +from typing import Dict, FrozenSet, List, Optional, Set, Tuple + +from setuptools import find_namespace_packages + +SETUP_PY = Path(__file__).resolve().parents[3] / "setup.py" +REPO_ROOT = SETUP_PY.parent + +# Enough of setup.py to exercise the keep set, and nothing that builds anything. +_WANTED = ( + "_WALK_SKIP_DIRS", + "_TEST_DIR_NAMES", + "_CI_ENTRY_POINTS", + "_CI_ENTRY_POINT_DIRS", + "_SHADER_TEMPLATE_MARKERS", + "_is_shader_template", + "_VENDORED_DIR_NAMES", + "_VENDORED_SUBMODULE_FALLBACK", + "_top_level_package_dirs", + "_first_party_module", + "_is_test_module", + "_module_name", + "_import_targets", + "_scan_imports", + "_GENERATED_DIR_NAMES", + "_unshipped_directories", + "_import_graph", + "_reachable_test_modules", + "_vendored_prefixes", + "_is_vendored_path", + "_full_packages", + "_minimal_packages", +) + + +def _setup_py_module() -> ast.Module: + # Name the encoding: these tests are collected on Windows too (pytest-windows.ini line 19), + # where the default is cp1252, and setup.py holds a non-ascii apostrophe that would decode to + # the wrong characters without a word rather than raising. + return ast.parse(SETUP_PY.read_text(encoding="utf-8")) + + +def _load_from_setup_py() -> Dict[str, object]: + """Run only the named definitions from setup.py, not its build logic.""" + selected: List[ast.stmt] = [] + found: Set[str] = set() + for node in _setup_py_module().body: + if isinstance(node, ast.FunctionDef) and node.name in _WANTED: + selected.append(node) + found.add(node.name) + elif isinstance(node, ast.Assign): + names = { + target.id + for target in node.targets + if isinstance(target, ast.Name) and target.id in _WANTED + } + if names: + selected.append(node) + found |= names + + assert found == set( + _WANTED + ), f"setup.py no longer defines {sorted(set(_WANTED) - found)}, so this test checks nothing" + + namespace: Dict[str, object] = { + "__file__": str(SETUP_PY), + "ast": ast, + "os": os, + "Path": Path, + "functools": functools, + "subprocess": subprocess, + "Dict": Dict, + "FrozenSet": FrozenSet, + "List": List, + "Optional": Optional, + "Set": Set, + "Tuple": Tuple, + "find_namespace_packages": find_namespace_packages, + } + exec( + compile(ast.Module(body=selected, type_ignores=[]), str(SETUP_PY), "exec"), + namespace, + ) + return namespace + + +_NAMESPACE = _load_from_setup_py() +_is_test_module = _NAMESPACE["_is_test_module"] +_import_graph = _NAMESPACE["_import_graph"] +_reachable_test_modules = _NAMESPACE["_reachable_test_modules"] +_CI_ENTRY_POINTS = _NAMESPACE["_CI_ENTRY_POINTS"] + + +@functools.lru_cache(maxsize=None) +def _graph() -> Tuple[Set[str], Dict[str, Set[str]], Set[str]]: + return _import_graph(REPO_ROOT / "src" / "executorch") + + +def _a_dropped_test_module() -> str: + """A real test module the keep set excludes, so the wiring tests assert on real data. + + Must be a leaf inside a test package, because find_package_modules is only given a chance to + drop something when the package it is asked about is itself under a test directory. + """ + modules, _edges, _dynamic = _graph() + keep = _reachable_test_modules() + dropped = sorted( + name + for name in modules + if name not in keep + and _is_test_module(name.rsplit(".", 1)[0]) + and name.rsplit(".", 1)[1] != "__init__" + ) + assert dropped, "nothing is dropped, so the wiring tests would be vacuous" + return dropped[0] + + +_UNREACHABLE_TEST_MODULE = _a_dropped_test_module() + + +@functools.lru_cache(maxsize=None) +def _reachable_from_imports_only() -> FrozenSet[str]: + """The keep set the import graph produces on its own, with no directory entries applied. + + Used to tell a load-bearing directory entry from a redundant one: a module the graph already + reaches would ship whether or not its directory is listed. + """ + modules, edges, dynamic = _graph() + referenced = set(dynamic) + referenced.update(_CI_ENTRY_POINTS) + for targets in edges.values(): + referenced.update(targets) + return frozenset(name for name in referenced if _is_test_module(name)) & modules + + +@functools.lru_cache(maxsize=None) +def _entry_point_dir_drivers() -> Dict[str, str]: + """Why each `_CI_ENTRY_POINT_DIRS` entry exists, as the file that drives it. + + These directories cannot be re-derived from the source, which is the whole reason they are + listed by hand: each is walked by something that never spells out a module name, so there is + no import to find and no literal to grep for. What CAN be checked is that the thing doing the + walking still exists and still refers to the directory. If a driver is deleted or stops + mentioning its directory, the entry has outlived its reason and this pairing fails. + + Keyed by the dotted prefix, valued by a repository-relative path. + """ + return { + "executorch.backends.mlx.custom_kernel_ops": ".github/workflows/mlx.yml", + "executorch.backends.webgpu.test": "backends/webgpu/scripts/test_webgpu_native_ci.sh", + "executorch.backends.test.suite": "backends/test/suite/runner.py", + "executorch.examples.models.llava.test": "examples/models/llava/README.md", + } + + +@functools.lru_cache(maxsize=None) +def _tracked_shell_scripts() -> Tuple[str, ...]: + """Shell scripts this repository actually owns, as repository-relative paths. + + Asked of git rather than found by walking. CI checks other repositories out INSIDE this one, + for example a `pytorch/` sibling clone, and a walk cannot tell those files from ours. It found + `pytorch/.ci/pytorch/test.sh` and reported a module belonging to a different project, so the + walk failed on CI while passing in every local checkout. + + An archive with no git available yields nothing, which makes this test vacuous rather than + wrong. It is a drift guard, so silence in an environment that cannot check is the safe way to + fail. + """ + try: + listed = subprocess.run( + ["git", "-C", str(REPO_ROOT), "ls-files", "-z", "*.sh"], + capture_output=True, + text=True, + check=False, + ) + except OSError: + # No git on PATH, as in an unpacked source archive. + return () + if listed.returncode: + return () + return tuple(name for name in listed.stdout.split("\0") if name) + + +def _fake_prune(build_lib, source_root): + """A CustomBuildPy whose prune runs, with build_lib and the source tree given directly. + + The real method resolves the source tree from setup.py's own location, so the lifted body is + bound to a stand-in whose __file__ points at the fixture instead. + """ + classes = [ + node + for node in _setup_py_module().body + if isinstance(node, ast.ClassDef) and node.name == "CustomBuildPy" + ] + assert len(classes) == 1, "setup.py no longer defines CustomBuildPy" + bodies = [ + node + for node in classes[0].body + if isinstance(node, ast.FunctionDef) and node.name == "_prune_unstaged_files" + ] + assert len(bodies) == 1, "the stale-file prune is gone" + + namespace = { + "os": os, + "Path": Path, + "__file__": str(source_root.parent / "setup.py"), + } + exec(compile(ast.unparse(bodies[0]), "prune", "exec"), namespace) + + class Stub: + editable_mode = False + packages = ["executorch", "executorch.pkg"] + + def __init__(self): + self.build_lib = str(build_lib) + + def find_all_modules(self): + # Deliberately omits stale.py, which is what marks it unwanted. + return [("executorch", "__init__", ""), ("executorch.pkg", "__init__", "")] + + def get_package_dir(self, package): + return str(source_root / Path(*package.split("."))) + + def find_data_files(self, package, src_dir): + return [] + + Stub._prune_unstaged_files = namespace["_prune_unstaged_files"] + return Stub() + + +def _fake_build_py(): + """A CustomBuildPy whose overrides run, without configuring a real distribution. + + The overrides are lifted from setup.py and bound to a stand-in so they can be CALLED. The + point is to exercise the real bodies: a test that only reads their syntax passes on code that + never runs, which is the hole this helper exists to close. + """ + classes = [ + node + for node in _setup_py_module().body + if isinstance(node, ast.ClassDef) and node.name == "CustomBuildPy" + ] + assert len(classes) == 1, "setup.py no longer defines CustomBuildPy" + wanted = ("find_package_modules", "find_data_files") + overrides = [ + node + for node in classes[0].body + if isinstance(node, ast.FunctionDef) and node.name in wanted + ] + assert {node.name for node in overrides} == set( + wanted + ), f"CustomBuildPy no longer overrides {sorted(set(wanted) - {n.name for n in overrides})}" + + package = _UNREACHABLE_TEST_MODULE.rsplit(".", 1)[0] + leaf = _UNREACHABLE_TEST_MODULE.rsplit(".", 1)[1] + modules = [(package, "__init__", "x"), (package, leaf, "y")] + + class Stub: + editable_mode = False + + def __init__(self) -> None: + self._data_files_to_return: List[str] = [] + + # Stands in for build_py's own implementations, which need a configured distribution. + def _super_find_package_modules(self, _package, _package_dir): + return list(modules) + + def _super_find_data_files(self, _package, _src_dir): + return list(self._data_files_to_return) + + namespace = dict(_NAMESPACE) + namespace["os"] = os + # `super()` needs a real base, so give the lifted bodies one that returns the fixtures above. + source = "\n".join( + ast.unparse(node) + .replace( + "super().find_package_modules(package, package_dir)", + "self._super_find_package_modules(package, package_dir)", + ) + .replace( + "super().find_data_files(package, src_dir)", + "self._super_find_data_files(package, src_dir)", + ) + for node in overrides + ) + exec(compile(source, "overrides", "exec"), namespace) + for name in wanted: + setattr(Stub, name, namespace[name]) + return Stub(), package, modules + + +class TestDroppedTestModules(unittest.TestCase): + def test_something_is_actually_dropped(self) -> None: + """The rule removes a substantial number of modules. + + Without this, every assertion below is vacuous on a keep set that happens to contain + everything, and the whole change could be reverted with the suite still green. + """ + modules, _, _ = _graph() + tests = {name for name in modules if _is_test_module(name)} + keep = _reachable_test_modules() + self.assertGreater(len(tests), 500, "no test modules discovered at all") + self.assertLess( + len(keep), + len(tests) // 2, + f"keeping {len(keep)} of {len(tests)} test modules, so almost nothing is dropped", + ) + + def test_shared_helpers_are_kept(self) -> None: + """Modules the suites import by installed name still ship. + + These are the ones whose removal breaks collection rather than a single test. Each is + imported from outside its own directory, which is what makes the installed name matter. + """ + keep = _reachable_test_modules() + for helper in ( + "executorch.backends.arm.test.tester.test_pipeline", + "executorch.backends.xnnpack.test.tester.tester", + "executorch.backends.test.harness.stages", + "executorch.backends.test.graph_builder", + "executorch.exir.backend.test.op_partitioner_demo", + ): + self.assertIn(helper, keep) + + def test_leaf_cases_are_dropped(self) -> None: + """A test case nothing imports does not ship. + + Chosen from different suites, because one backend getting this right says nothing about + the others. + """ + keep = _reachable_test_modules() + modules, _, _ = _graph() + for leaf in ( + "executorch.backends.arm.test.ops.test_add", + "executorch.backends.xnnpack.test.ops.test_bilinear2d", + ): + self.assertIn( + leaf, modules, f"{leaf} no longer exists, pick another example" + ) + self.assertNotIn(leaf, keep) + + def test_relative_imports_are_followed(self) -> None: + """A submodule reached only by a relative import is kept. + + backends/test/harness/stages/__init__.py does `from .export import Export`, so treating + a relative import as reaching nothing new drops stages.export and breaks every importer + of that package. This is a regression guard: it failed exactly that way once. + """ + self.assertIn( + "executorch.backends.test.harness.stages.export", _reachable_test_modules() + ) + + def test_dynamic_imports_are_followed(self) -> None: + """A module named only as a string to importlib is kept. + + backends/mlx/test/run_all_tests.py does + `importlib.import_module(".test_ops", package=__package__)`, which an import scan that + only reads import statements cannot see. Note test_ops is also named like a leaf, so a + file name rule would drop it. + """ + self.assertIn( + "executorch.backends.mlx.test.test_ops", _reachable_test_modules() + ) + + def test_ci_entry_points_are_kept(self) -> None: + """The modules only a workflow names are kept.""" + keep = _reachable_test_modules() + for name in _CI_ENTRY_POINTS: + self.assertIn(name, keep) + + def test_ci_entry_points_still_match_the_workflows(self) -> None: + """The hand-written CI list has not drifted from what the workflows actually run. + + The list is explicit rather than scanned at build time, because a source distribution + carries no .github directory and a scan there would silently keep nothing. The cost of + being explicit is drift, so it is checked here instead. + """ + pattern = re.compile(r"executorch(?:\.[A-Za-z0-9_]+)+") + referenced: Set[str] = set() + # The whole tree, not just .github and .ci. A workflow often calls a script that lives + # beside the backend it tests, and those name modules too: + # backends/webgpu/scripts/test_webgpu_native_ci.sh runs six of them by dotted name. + skip = { + ".git", + "pip-out", + "cmake-out", + "third-party", + "third_party", + "__pycache__", + } + for dirpath, dirnames, filenames in os.walk(REPO_ROOT, followlinks=False): + dirnames[:] = [d for d in dirnames if d not in skip] + for filename in filenames: + # Markdown too: a README documenting `python -m executorch.x.test.y` is a + # promise to users, and dropping that module breaks the documented command. + # Python as well, because several modules document their own `python -m` + # invocation in a docstring rather than in a README, and that is the same + # promise written somewhere else. + if not filename.endswith( + (".yml", ".yaml", ".sh", ".ps1", ".md", ".py") + ): + continue + path = Path(dirpath) / filename + if path.resolve() == Path(__file__).resolve(): + # This file names dropped modules as examples of what the rule removes, so + # reading itself would report them as promised and contradict its own tests. + continue + text = path.read_text(encoding="utf-8", errors="replace") + referenced.update(pattern.findall(text)) + + modules, _, _ = _graph() + + # A reference like executorch.a.test.b.SomeClass.some_method is one dotted run to the + # regex, and it is not a module, so trim each match back to its longest real module + # prefix. Without this the class-suffixed entries silently drop out of the comparison + # and the guard protects fewer names than it appears to. + def longest_module(name: str) -> str: + parts = name.split(".") + while parts: + candidate = ".".join(parts) + if candidate in modules: + return candidate + parts.pop() + return name + + expected = { + trimmed + for trimmed in (longest_module(name) for name in referenced) + if _is_test_module(trimmed) and trimmed in modules + } + missing = sorted(expected - set(_CI_ENTRY_POINTS) - _reachable_test_modules()) + self.assertEqual( + missing, + [], + f"a workflow names these test modules but nothing keeps them: {missing}", + ) + + def test_parent_packages_of_kept_modules_are_kept(self) -> None: + """Every kept module's package chain is kept, or the dotted path cannot resolve.""" + keep = _reachable_test_modules() + modules, _, _ = _graph() + for name in keep: + parts = name.split(".") + for end in range(2, len(parts)): + parent = ".".join(parts[:end]) + if _is_test_module(parent) and parent in modules: + self.assertIn(parent, keep, f"{parent} missing but {name} is kept") + + def test_shader_templates_do_not_ship(self) -> None: + """Shader codegen inputs are dropped, and op definitions are not. + + The cmake build expands these into SPIR-V and WGSL headers, so the wheel already carries + the compiled result. Matched on content, so the two examples below are the real + distinction: one is a template, the other is read at run time through + importlib.resources and must survive. + """ + is_template = _NAMESPACE["_is_shader_template"] + self.assertTrue( + is_template("backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml") + ) + self.assertTrue( + is_template("backends/webgpu/runtime/ops/binary_op/binary_op.yaml") + ) + for needed in ( + "exir/dialects/edge/edge.yaml", + "kernels/portable/functions.yaml", + "backends/cadence/aot/functions.yaml", + ): + self.assertFalse(is_template(needed), f"{needed} would stop shipping") + + def test_build_py_is_wired_to_the_custom_class(self) -> None: + """setup() receives CustomBuildPy, not the stock build_py. + + Every other test here exercises the class directly, so all of them stay green when the + cmdclass entry is pointed back at setuptools' own build_py. That single edit disables + the module filter, the data file filter and the prune at once, and the wheel then ships + everything again. + """ + assignments = [ + node + for node in ast.walk(_setup_py_module()) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "setup" + ] + self.assertEqual(len(assignments), 1, "expected exactly one setup() call") + + mapping = [kw.value for kw in assignments[0].keywords if kw.arg == "cmdclass"] + self.assertEqual(len(mapping), 1, "setup() no longer passes cmdclass") + wired = { + key.value: value.id + for key, value in zip(mapping[0].keys, mapping[0].values) + if isinstance(key, ast.Constant) and isinstance(value, ast.Name) + } + self.assertEqual( + wired.get("build_py"), + "CustomBuildPy", + "build_py is not wired to CustomBuildPy, so none of the filters run", + ) + + def test_both_package_lists_are_anchored_on_this_file(self) -> None: + """Neither package list depends on the working directory. + + A cwd-relative `where` returns nothing when the build runs from anywhere but the + repository root, and an empty package list makes the prune treat every staged file as + unwanted. The full list was anchored for this reason; the minimal one has to match. + + Both lists are CALLED from a directory that is not the repository root, because reading + the syntax of the `where=` argument only proves it is not a literal. Swapping the anchor + for `Path.cwd()` leaves the syntax test green and breaks every build started elsewhere. + """ + original = os.getcwd() + os.chdir(tempfile.gettempdir()) + try: + full = _NAMESPACE["_full_packages"]() + minimal = _NAMESPACE["_minimal_packages"]() + finally: + os.chdir(original) + self.assertIn("executorch", full) + self.assertGreater( + len(full), 100, "the full list collapsed when built from another directory" + ) + self.assertIn("executorch", minimal) + self.assertGreater( + len(minimal), + 1, + "the minimal list collapsed when built from another directory", + ) + + def test_stale_staged_files_are_pruned(self) -> None: + """A rebuild removes what an earlier build staged and this one does not want. + + build_py only copies, so without this a second build into the same directory keeps + every file the first one put there and the wheel packages it. The failure is silent: + the build succeeds and the wheel quietly contains the dropped files. + + The prune is CALLED against a real staging directory, because checking that the method + and its call site exist leaves an early `return` inside the body undetected, and the + prune then does nothing while this test stays green. + """ + staging = Path(tempfile.mkdtemp(prefix="prunetest-")) + self.addCleanup(shutil.rmtree, staging, ignore_errors=True) + source = staging / "src" + (source / "executorch" / "pkg").mkdir(parents=True) + for name in ("executorch/__init__.py", "executorch/pkg/__init__.py"): + (source / name).write_text("") + # Exists in the source tree and is NOT in build_py's file list, so the prune wants it + # gone. That is the whole contract. + (source / "executorch" / "pkg" / "stale.py").write_text( + "# left by an earlier build\n" + ) + build_lib = staging / "lib" + shutil.copytree(source, build_lib) + # Generated by a later build command, absent from src/, and must survive. + (build_lib / "executorch" / "pkg" / "generated.py").write_text( + "# from a template\n" + ) + + command = _fake_prune(build_lib, source) + command._prune_unstaged_files() + + remaining = sorted(p.name for p in (build_lib / "executorch" / "pkg").iterdir()) + self.assertNotIn( + "stale.py", + remaining, + "the prune left a file the current build does not want", + ) + self.assertIn( + "generated.py", + remaining, + "the prune deleted a file another command generated", + ) + self.assertIn("__init__.py", remaining, "the prune deleted a wanted module") + + def test_build_py_applies_the_keep_set(self) -> None: + """The drop is actually wired into the build, checked by CALLING the override. + + An earlier version of this test read the override's syntax tree instead. That passes on + code that is present but never runs, so an early `return modules` at the top of the + override left the filter dead with every assertion here still true. Build a real command + and look at what it returns. + """ + command, package, modules = _fake_build_py() + result = command.find_package_modules(package, "unused") + returned = {entry[1] for entry in result} + offered = {entry[1] for entry in modules} + self.assertIn("__init__", returned, "a kept package must still import") + self.assertTrue( + offered - returned, + "find_package_modules returned everything it was offered, so nothing is dropped", + ) + self.assertNotIn( + _UNREACHABLE_TEST_MODULE.rsplit(".", 1)[1], + returned, + "an unreachable test module was not dropped", + ) + + def test_shader_filter_is_wired_into_find_data_files(self) -> None: + """The shader classifier is actually CALLED, not merely correct. + + test_shader_templates_do_not_ship above checks the predicate. That is not the same thing: + deleting the filtering line in find_data_files leaves the predicate perfect and unused, + and every shader template ships again. + """ + command, _package, _modules = _fake_build_py() + template = "backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml" + needed = "kernels/portable/functions.yaml" + root = str(REPO_ROOT) + command._data_files_to_return = [ + os.path.join(root, template), + os.path.join(root, needed), + ] + kept = command.find_data_files("executorch", root) + self.assertNotIn( + os.path.join(root, template), + kept, + "find_data_files does not drop shader templates, so the filter is not wired in", + ) + self.assertIn( + os.path.join(root, needed), + kept, + "find_data_files dropped a yaml the runtime reads", + ) + + def test_ci_entry_point_dirs_are_all_load_bearing(self) -> None: + """Every listed directory still has a driver, and still keeps something. + + Nothing referenced `_CI_ENTRY_POINT_DIRS`, so an entry could be deleted with the whole + suite green: removing the backend suite line silently stopped 86 modules shipping. Two + checks close that. Each entry must be paired with the file that walks it, which fails when + an entry is added or removed without updating the pairing, and each entry must keep modules + the import graph cannot reach on its own, which fails when an entry becomes dead weight. + """ + listed = set(_NAMESPACE["_CI_ENTRY_POINT_DIRS"]) + self.assertTrue(listed, "the list is empty, so nothing is protected") + + drivers = _entry_point_dir_drivers() + self.assertEqual( + listed, + set(drivers), + "_CI_ENTRY_POINT_DIRS and its list of drivers disagree. Add the new entry with the " + "file that walks it, or drop the driver for the entry that went away", + ) + + for prefix, driver in sorted(drivers.items()): + path = REPO_ROOT / driver + self.assertTrue( + path.is_file(), + f"{prefix} is kept for {driver}, which no longer exists, so the entry may be " + "obsolete", + ) + tail = prefix.split(".")[-1] + self.assertIn( + tail, + path.read_text(encoding="utf-8", errors="replace"), + f"{driver} no longer mentions {tail}, so it may have stopped driving {prefix}", + ) + + # And the other direction: an entry that keeps nothing new is dead weight. + reached_anyway = _reachable_from_imports_only() + keep = _reachable_test_modules() + for entry in sorted(listed): + covered = { + name for name in keep if name == entry or name.startswith(f"{entry}.") + } + self.assertTrue( + covered - reached_anyway, + f"{entry} keeps nothing the import graph does not already reach, so the entry " + "is redundant and should be removed", + ) + + def test_ci_entry_points_cover_constructed_module_names(self) -> None: + """A runner that BUILDS a dotted name is covered too. + + The drift test above searches for a literal dotted name, so it cannot see a script that + assembles one, and a directory whose tests are only reached that way would be dropped + with nothing to warn about. + + A script that runs from a checkout by design is exempt, and says so in its own header. + `backends/apple/coreai/run_all_tests.sh` is the current example: it cds to the repository + root, so it always finds the files on disk and never needs them installed. + """ + pattern = re.compile(r"find\s+([A-Za-z0-9_./-]+)\s+-name\s+'?test_\*\.py'?") + keep = _reachable_test_modules() + listed = _NAMESPACE["_CI_ENTRY_POINT_DIRS"] + unprotected = [] + for relative in _tracked_shell_scripts(): + path = REPO_ROOT / relative + text = path.read_text(encoding="utf-8", errors="replace") + walked = pattern.findall(text) + if not walked: + continue + if "not a landing artifact" in text: + continue + for entry in walked: + dotted = "executorch." + entry.strip("./").replace("/", ".") + covered = any( + dotted == prefix or dotted.startswith(f"{prefix}.") + for prefix in listed + ) or any(name.startswith(f"{dotted}.") for name in keep) + if not covered: + unprotected.append(f"{relative} -> {dotted}") + self.assertEqual( + unprotected, + [], + "a script discovers test modules under these paths by building dotted names, and " + "nothing keeps them. Either add the directory to _CI_ENTRY_POINT_DIRS in setup.py, " + "or say in the script's header that it is not a landing artifact if it only ever " + f"runs from a checkout: {unprotected}", + ) + + def test_unprefixed_first_party_imports_count_as_references(self) -> None: + """`from backends.x import y` keeps y, the same as the prefixed spelling. + + This repository imports itself both ways: most code says `executorch.backends.x`, but the + Arm suites say `backends.arm.test...`, which resolves because the repository root is on + sys.path. Both name the same file. Following only the prefixed spelling dropped two shared + helpers with eight importers between them, which is the invariant this change exists to + preserve. + """ + first_party = _NAMESPACE["_first_party_module"] + self.assertEqual( + first_party("backends.arm.test.common"), + "executorch.backends.arm.test.common", + ) + self.assertEqual( + first_party("executorch.exir.tests.common"), "executorch.exir.tests.common" + ) + # A third-party module whose first component is not one of ours stays out. + self.assertIsNone(first_party("torch.nn.functional")) + self.assertIsNone(first_party("numpy")) + + keep = _reachable_test_modules() + for helper in ( + "executorch.backends.arm.test._custom_vgf_test_utils", + "executorch.backends.arm.test.runtime._vgf_runtime_test_utils", + ): + self.assertIn( + helper, + keep, + f"{helper} is imported without the executorch prefix and must still ship", + ) + + def test_importers_outside_the_shipped_tree_are_followed(self) -> None: + """A file the wheel does not carry can still import one that it does. + + `src/executorch` is a subset of the checkout, so an importer in a directory that is never + packaged is invisible to a walk of the shipped tree alone. Its imports still have to keep + their targets: test/end2end/test_end2end.py imports two model helpers out of exir/tests. + """ + keep = _reachable_test_modules() + importer = REPO_ROOT / "test" / "end2end" / "test_end2end.py" + self.assertTrue( + importer.is_file(), "this test needs a different example importer" + ) + for helper in ( + "executorch.exir.tests.dynamic_shape_models", + "executorch.exir.tests.transformer", + ): + self.assertIn( + helper, + keep, + f"{helper} is imported from outside the shipped tree and must still ship", + ) + + def test_vendored_trees_are_not_read_as_import_evidence(self) -> None: + """A vendored submodule's own imports do not keep anything. + + The package list excludes vendored trees, so nothing in one ships. The import scan has to + agree, or the two disagree about the same directory: a submodule checked out under an + ordinary name, rather than under `third-party`, was read as first-party and its imports + kept test modules the wheel never carries. + + Skipping by directory name alone is not enough, which is why this asserts on the scan's + output rather than on the skip list. + """ + modules, _edges, _dynamic = _graph() + is_vendored = _NAMESPACE["_is_vendored_path"] + vendored = sorted( + name for name in modules if is_vendored(name.replace(".", "/")) + ) + self.assertEqual( + vendored, + [], + "the import scan read these vendored modules as first-party, so their imports can " + f"keep test modules nothing shipped reaches: {vendored[:5]}", + ) + + def test_generated_directories_are_not_read_as_import_evidence(self) -> None: + """A build tree or an in-tree virtualenv does not vote on what ships. + + Those hold an INSTALLED copy of this package, so reading one lets the last wheel decide + what the next carries: a file that shipped once keeps itself alive. A clean checkout has + none of them, so the guard is exercised here by creating one. + """ + unshipped = _NAMESPACE["_unshipped_directories"] + root = REPO_ROOT / "src" / "executorch" + planted = REPO_ROOT / ".venv" + created = not planted.exists() + if created: + (planted / "lib").mkdir(parents=True) + self.addCleanup(shutil.rmtree, planted, ignore_errors=True) + walked = {entry.name for entry in unshipped(root)} + self.assertNotIn( + ".venv", + walked, + "a generated directory is read as import evidence, so an installed copy of this " + "package can keep test modules alive across builds", + ) + self.assertIn( + "test", walked, "the guard also dropped a real unshipped directory" + ) + + def test_editable_installs_are_left_alone(self) -> None: + """An editable install still exposes every test module. + + It maps the package root to a directory, so the suites resolve from the source tree + whatever is listed, and dropping modules there would only make the two install modes + disagree for no benefit. + """ + classes = [ + node + for node in _setup_py_module().body + if isinstance(node, ast.ClassDef) and node.name == "CustomBuildPy" + ] + overrides = [ + node + for node in classes[0].body + if isinstance(node, ast.FunctionDef) and node.name == "find_package_modules" + ] + source = ast.unparse(overrides[0]) + self.assertIn("editable_mode", source) + + +if __name__ == "__main__": + unittest.main() diff --git a/.ci/scripts/tests/test_wheel_vendored_packages.py b/.ci/scripts/tests/test_wheel_vendored_packages.py new file mode 100644 index 00000000000..a3d81a940b0 --- /dev/null +++ b/.ci/scripts/tests/test_wheel_vendored_packages.py @@ -0,0 +1,494 @@ +# 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 the packages the full wheel publishes. + +The wheel used to carry the Python files and codegen scripts of every vendored third-party +checkout, because the full build passed no `packages` list and setuptools then discovered +everything under src/executorch. Those files exist to build the C++ targets, so nothing in +an installed wheel imports them. + +Asserting on the discovery result rather than on a built wheel, because the behaviour under +test is a pure function of the source tree plus the exclude patterns, and a full build takes +minutes to exercise one filter. `.ci/scripts/test_minimal_wheel.sh` already covers the +built-artifact side for the minimal wheel. + +setup.py is read rather than imported. It calls setup() at module scope, so importing it under +a test runner hands setup() the runner's own arguments and the session dies on an invalid +command name. +""" + +import ast +import functools +import os +import subprocess +import tempfile +import unittest +from pathlib import Path +from typing import Dict, FrozenSet, List, Set, Tuple + +from setuptools import find_namespace_packages +from setuptools.command.build_py import build_py + +SETUP_PY = Path(__file__).resolve().parents[3] / "setup.py" +REPO_ROOT = SETUP_PY.parent +# Discovery is anchored on this file's location, not on the working directory, so the result +# does not depend on where the runner was started. +PACKAGE_ROOT = str(SETUP_PY.parent / "src") + + +# The helpers this test drives, shared by both loaders below. +_HELPERS = ( + # _VENDORED_DIR_NAMES is not used directly below, but _is_vendored_path closes over it. + "_VENDORED_DIR_NAMES", + "_VENDORED_SUBMODULE_FALLBACK", + "_vendored_prefixes", + "_is_vendored_path", + # CustomBuildPy calls this, so the class cannot be exec'd without it. + "_SHADER_TEMPLATE_MARKERS", + "_is_shader_template", + "_full_packages", +) + + +def _setup_py_module() -> ast.Module: + # Name the encoding: these tests are collected on Windows too (pytest-windows.ini line 19), + # where the default is cp1252, and setup.py holds a non-ascii apostrophe that would decode to + # the wrong characters without a word rather than raising. + return ast.parse(SETUP_PY.read_text(encoding="utf-8")) + + +def _load_from_setup_py(root: Path = None) -> Dict[str, object]: + """The vendored-path helpers and the package list builder, from setup.py's source. + + Only those definitions are executed, so none of setup.py's module level build logic runs. + """ + wanted = _HELPERS + + selected: List[ast.stmt] = [] + found = set() + for node in _setup_py_module().body: + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in wanted: + selected.append(node) + found.add(node.name) + elif isinstance(node, ast.Assign): + names = { + target.id + for target in node.targets + if isinstance(target, ast.Name) and target.id in wanted + } + if names: + selected.append(node) + found |= names + + assert found == set( + wanted + ), f"setup.py no longer defines {sorted(set(wanted) - found)}, so this test checks nothing" + + namespace: Dict[str, object] = { + "__file__": str((root or SETUP_PY.parent) / "setup.py"), + "Path": Path, + "List": List, + "Tuple": Tuple, + "functools": functools, + "subprocess": subprocess, + "find_namespace_packages": find_namespace_packages, + } + exec( + compile(ast.Module(body=selected, type_ignores=[]), str(SETUP_PY), "exec"), + namespace, + ) + return namespace + + +@functools.lru_cache(maxsize=None) +def _load_build_py() -> Dict[str, object]: + """CustomBuildPy plus the helpers it calls, so analyze_manifest can be driven directly. + + Only the class body and those helpers run. Its methods reference names from setup.py's own + imports, so the ones analyze_manifest touches are supplied here. + """ + wanted = {"CustomBuildPy"} | set(_HELPERS) + selected: List[ast.stmt] = [] + for node in _setup_py_module().body: + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in wanted: + selected.append(node) + elif isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id in wanted + for target in node.targets + ): + selected.append(node) + + namespace: Dict[str, object] = { + "__file__": str(SETUP_PY), + "os": os, + "ast": ast, + "Path": Path, + "functools": functools, + "subprocess": subprocess, + "build_py": build_py, + "Dict": Dict, + "FrozenSet": FrozenSet, + "List": List, + "Set": Set, + "Tuple": Tuple, + "find_namespace_packages": find_namespace_packages, + } + exec( + compile(ast.Module(body=selected, type_ignores=[]), str(SETUP_PY), "exec"), + namespace, + ) + return namespace + + +_NAMESPACE = _load_from_setup_py() +_vendored_prefixes = _NAMESPACE["_vendored_prefixes"] +_is_vendored_path = _NAMESPACE["_is_vendored_path"] +_full_packages = _NAMESPACE["_full_packages"] + + +def _discovered_packages() -> List[str]: + """Everything setuptools finds, before any of this change's filtering.""" + return sorted( + find_namespace_packages( + where=PACKAGE_ROOT, include=["executorch", "executorch.*"] + ) + ) + + +def _vendored(packages: List[str]) -> List[str]: + return [ + package for package in packages if _is_vendored_path(package.replace(".", "/")) + ] + + +class TestFullWheelPackages(unittest.TestCase): + def test_the_tree_has_vendored_packages_to_exclude(self) -> None: + """Fail rather than skip when there is nothing to exclude. + + Every other test here is vacuous on a tree with no vendored checkouts: an empty + package list contains no vendored package, so the exclusion would look correct even + if it had been deleted. Assert the premise instead of quietly passing on it. + """ + discovered = _discovered_packages() + self.assertNotEqual( + discovered, [], f"no packages discovered under {PACKAGE_ROOT}" + ) + self.assertNotEqual( + _vendored(discovered), + [], + "no vendored third-party packages in this tree, so the exclusion below cannot " + "be shown to do anything. Initialize the submodules before running this.", + ) + + def test_no_vendored_package_ships(self) -> None: + """No package from another repository is published. + + Compares against what discovery finds rather than re-filtering the helper's own output. + Filtering the result with the same predicate the helper already applied is a tautology: + it is empty whatever the helper did, so it would pass even with the exclusion removed. + """ + discovered = set(_discovered_packages()) + shipped = set(_full_packages()) + dropped = discovered - shipped + + leaked = sorted(shipped & set(_vendored(discovered))) + # Only the count and a few names, because a regression here leaks hundreds of + # packages and the default diff would bury the message. + self.assertEqual( + len(leaked), + 0, + f"the wheel would publish {len(leaked)} vendored packages, " + f"e.g. {leaked[:3]}", + ) + # And the helper really removed them, rather than discovery never having found them. + self.assertEqual( + dropped, + set(_vendored(discovered)), + "the set the helper drops is not the set of vendored packages on disk", + ) + + def test_the_exclusion_is_load_bearing(self) -> None: + """Discovery without the exclusion finds the packages the exclusion removes.""" + self.assertLess( + len(_full_packages()), + len(_discovered_packages()), + "the exclusion dropped nothing, so it is no longer doing any work", + ) + + def test_setup_passes_the_package_list(self) -> None: + """The helper is actually wired into the full build. + + Without this, every test above still passes when the assignment that hands the list + to setuptools is deleted, which is the whole of the change. The sibling wheel test + asserts its own wiring the same way and for the same reason. + + The search is limited to the else branch of the minimal-build check, because an + unrestricted walk also matches an assignment that can never run: moved into the + minimal branch it is overwritten by the next line, and wrapped in a false condition + it is dead, and both of those leave the full wheel discovering everything. + """ + minimal_checks = [ + node + for node in _setup_py_module().body + if isinstance(node, ast.If) + and isinstance(node.test, ast.Call) + and isinstance(node.test.func, ast.Name) + and node.test.func.id == "_is_minimal_build" + ] + self.assertEqual( + len(minimal_checks), + 1, + "expected exactly one module level `if _is_minimal_build():`", + ) + + assigned = [ + node + for node in minimal_checks[0].orelse + if isinstance(node, ast.Assign) + for target in node.targets + if isinstance(target, ast.Subscript) + and isinstance(target.value, ast.Name) + and target.value.id == "setup_kwargs" + and isinstance(target.slice, ast.Constant) + and target.slice.value == "packages" + and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) + and node.value.func.id == "_full_packages" + ] + self.assertEqual( + len(assigned), + 1, + "setup.py does not assign _full_packages() to setup_kwargs['packages'], " + "so the full build falls back to discovering every package", + ) + + def test_first_party_packages_still_ship(self) -> None: + """A named first-party package survives the exclusion. + + Every other test here asks whether unwanted packages left. This one asks whether + wanted ones stayed, which is the failure mode a too-greedy filter produces and the + one nothing else would notice. + """ + packages = _full_packages() + for package in ( + "executorch.exir", + "executorch.backends.xnnpack", + "executorch.extension.pybindings", + "executorch.devtools", + ): + self.assertIn(package, packages) + + def test_only_submodule_sections_are_read(self) -> None: + """A `path` line outside a submodule section is not an exclusion prefix. + + Written against a file with a stray entry rather than by comparing git's output to git's + own output. That comparison holds for any reader on today's clean file, so it would pass + just as well for a line scanner that accepts `path =` from any section, which is the + failure this is meant to catch: one stray line silently removes a real package. + """ + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / ".gitmodules").write_text( + '[submodule "real"]\n' + "\tpath = extension/llm/tokenizers\n" + "[core]\n" + "\tpath = executorch/exir\n" + ) + self.assertEqual( + _load_from_setup_py(root)["_vendored_prefixes"](), + ("extension/llm/tokenizers",), + "a path line outside a submodule section became an exclusion prefix", + ) + + def test_prefixes_are_normalized(self) -> None: + """A legal but unusual spelling in .gitmodules still matches the real directory. + + git treats a trailing slash, a leading ./ and a doubled separator as the same path, + so storing the raw text would silently disable the exclusion for that entry. Asserted + against a written file rather than against today's values, because today's are already + tidy and would pass either way. + """ + for spelling in ( + "extension/llm/tokenizers/", + "./extension/llm/tokenizers", + "extension//llm/tokenizers", + ): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + # Only the path is read, so the entry needs no url. + (root / ".gitmodules").write_text( + f'[submodule "t"]\n\tpath = {spelling}\n' + ) + # The helper reads .gitmodules beside its own setup.py, so it is loaded + # against the temporary tree rather than the real one. + prefixes = _load_from_setup_py(root)["_vendored_prefixes"]() + self.assertEqual( + prefixes, + ("extension/llm/tokenizers",), + f"{spelling!r} did not normalize", + ) + + def test_the_fallback_matches_gitmodules(self) -> None: + """The hardcoded fallback still lists the same submodules the file does. + + It is only used when .gitmodules cannot be read, which is the case in a source + distribution, so nothing else would notice it drifting out of date. + """ + self.assertEqual( + _vendored_prefixes(), _NAMESPACE["_VENDORED_SUBMODULE_FALLBACK"] + ) + + # And it is actually returned when the file is missing, which is the only case it + # exists for. Without this the fallback could be replaced by an empty tuple and the + # comparison above would still hold. + with tempfile.TemporaryDirectory() as tmp: + self.assertEqual( + _load_from_setup_py(Path(tmp))["_vendored_prefixes"](), + _NAMESPACE["_VENDORED_SUBMODULE_FALLBACK"], + "with no .gitmodules the submodule exclusion silently does nothing", + ) + + def test_a_submodule_name_with_a_space_is_read(self) -> None: + """A submodule whose NAME contains a space still yields its path. + + git prints " " and permits spaces in the name, so splitting on the first + space truncates the key and leaves a value that matches no directory, turning the + exclusion off for that entry. + """ + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / ".gitmodules").write_text( + '[submodule "my module"]\n\tpath = extension/llm/tokenizers\n' + ) + self.assertEqual( + _load_from_setup_py(root)["_vendored_prefixes"](), + ("extension/llm/tokenizers",), + ) + + def test_a_broken_gitmodules_falls_back(self) -> None: + """An unreadable .gitmodules reaches the fallback rather than excluding nothing. + + git exits non-zero with empty output on a bad section header or on conflict markers. + Reading that as "this repository has no submodules" would turn the exclusion off with + no warning, which is the one failure the fallback exists to prevent. + """ + for broken in ( + '[submodule "x"\n\tpath = extension/llm/tokenizers\n', + '<<<<<<< HEAD\n[submodule "x"]\n\tpath = a/b\n=======\n', + "", + ): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / ".gitmodules").write_text(broken) + prefixes = _load_from_setup_py(root)["_vendored_prefixes"]() + self.assertEqual( + prefixes, + _NAMESPACE["_VENDORED_SUBMODULE_FALLBACK"], + f"a broken .gitmodules ({broken[:20]!r}) silently excluded nothing", + ) + + def test_manifest_filter_actually_drops_vendored_data(self) -> None: + """The data-file half of the fix removes files, through the real build code path. + + `packages` only governs Python modules. Non-Python files arrive through the + package_data manifest, and setuptools attributes a file under an unlisted directory to + its nearest listed parent, so vendored data returns unless the manifest is filtered too. + + Drives CustomBuildPy.analyze_manifest itself rather than reimplementing the filter here. + Checking the predicate in isolation is not enough: inverting the editable guard or + short-circuiting the condition leaves the predicate correct and the build unfiltered, + and both of those left an earlier version of this test green. + """ + namespace = _load_build_py() + build_py_class = namespace["CustomBuildPy"] + + vendored = ( + "src/executorch/backends/xnnpack/third-party/generate-cpuinfo-wrappers.py" + ) + # A shader template goes through this same filter, and it needs its own example here: + # deleting the shader line leaves the vendored assertions below green, so the manifest + # half of the shader fix was unprotected. + shader = "src/executorch/backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml" + ordinary = "setup.py" + # All of them have to exist on disk, because the filter also drops anything that is not a + # file, and a missing path would be removed for that reason instead of this one. + for relative in (vendored, shader, ordinary): + self.assertTrue( + (REPO_ROOT / relative).is_file(), + f"{relative} is gone, so this test needs a different example", + ) + + # analyze_manifest calls up into setuptools first, which needs the full command + # machinery. Only the filtering after that call is under test, so the parent's method + # is replaced with a no-op for the duration and the manifest seeded directly. This runs + # the shipped code path rather than a copy of it, which is the point: a filter that has + # been turned off still reads correctly in the source. + parent = build_py_class.__mro__[1] + original = parent.analyze_manifest + parent.analyze_manifest = lambda self: None + try: + stub = build_py_class.__new__(build_py_class) + stub.editable_mode = False + stub.manifest_files = {"executorch": [vendored, shader, ordinary]} + stub.analyze_manifest() + kept = stub.manifest_files["executorch"] + finally: + parent.analyze_manifest = original + + self.assertNotIn( + vendored, kept, "a vendored data file survived the manifest filter" + ) + self.assertNotIn(shader, kept, "a shader template survived the manifest filter") + self.assertIn(ordinary, kept, "the filter dropped an ordinary file") + + def test_is_vendored_path_matches_whole_components(self) -> None: + """The filter matches a path component, not a substring.""" + self.assertTrue( + _is_vendored_path( + "src/executorch/backends/xnnpack/third-party/XNNPACK/a.py" + ) + ) + self.assertTrue(_is_vendored_path("src/executorch/x/third_party/y.yaml")) + self.assertFalse(_is_vendored_path("src/executorch/exir/program/_program.py")) + # "third-party" as part of a longer name is a different directory. + self.assertFalse(_is_vendored_path("src/executorch/x/third-party-tools/y.py")) + + def test_submodules_outside_a_vendored_dir_are_recognized(self) -> None: + """A submodule checked out under an ordinary name is still another repository. + + These are not matched by the directory name, so they are read from .gitmodules. Their + nested copies also cannot satisfy the imports the code uses: the FACTO helper imports + facto.specdb from the top level, and the tokenizers ship as a declared dependency. + """ + prefixes = _vendored_prefixes() + self.assertIn("backends/cadence/utils/FACTO", prefixes) + self.assertIn("extension/llm/tokenizers", prefixes) + for prefix in ("backends/cadence/utils/FACTO", "extension/llm/tokenizers"): + self.assertTrue(_is_vendored_path(f"executorch/{prefix}")) + self.assertTrue(_is_vendored_path(f"src/executorch/{prefix}/setup.py")) + self.assertFalse( + _is_vendored_path("executorch/extension/llm/custom_ops/op_sdpa.py") + ) + + def test_root_level_submodules_are_not_listed(self) -> None: + """A submodule at the repository root is not a wheel path. + + Those are build tooling, never copied into the package, and listing one would put a + bare single-word name into the matcher. That would then drop any directory sharing the + name, anywhere in the tree, which is a much wider rule than intended. + """ + for prefix in _vendored_prefixes(): + self.assertIn( + "/", + prefix, + f"{prefix!r} is a root-level submodule and must not be listed", + ) + self.assertFalse(_is_vendored_path("executorch/some/nested/shim")) + + +if __name__ == "__main__": + unittest.main() diff --git a/.ci/scripts/unittest-linux-cmake.sh b/.ci/scripts/unittest-linux-cmake.sh index 0f750e1fe13..83f2f1464ee 100755 --- a/.ci/scripts/unittest-linux-cmake.sh +++ b/.ci/scripts/unittest-linux-cmake.sh @@ -7,6 +7,9 @@ # LICENSE file in the root directory of this source tree. set -eux +# Cap pytest-xdist's `auto` workers to the container's CPU quota. +source .ci/scripts/pytest-parallelism.sh + # Some ARM/TOSA-adjacent tests import modules that require tosa_serializer. # Install from a local tosa-tools checkout when available. If absent in this # checkout layout, clone the pinned upstream tag and install from there. diff --git a/.ci/scripts/utils.sh b/.ci/scripts/utils.sh index 234e162e48e..3d573daa395 100644 --- a/.ci/scripts/utils.sh +++ b/.ci/scripts/utils.sh @@ -106,8 +106,8 @@ install_pytorch_and_domains() { local python_version=$(python -c 'import platform; v=platform.python_version_tuple(); print(f"{v[0]}{v[1]}")') local torch_release=$(cat version.txt) # Download key must match the upload key below (basename of dist/*.whl, - # which always carries setup.py's resolved +gitHASH). Branch-ref pins - # like `release/2.13` would otherwise produce `+gitrelease` here and + # which always carries the build's resolved +gitHASH). Branch-ref pins + # like `release/2.14` would otherwise produce `+gitrelease` here and # never hit the cache. local torch_short_hash=$(git rev-parse --short=7 HEAD) local torch_wheel_path="cached_artifacts/pytorch/executorch/pytorch_wheels/${system_name}/${python_version}" @@ -127,18 +127,31 @@ install_pytorch_and_domains() { if [[ "${torch_wheel_not_found}" == "1" ]]; then echo "No cached wheel found, continue with building PyTorch at ${TORCH_VERSION}" - # Install PyTorch's own build-time deps so the source build does not - # silently inherit them from whatever else happens to be in the env - # (e.g. executorch's requirements-ci.txt). - pip install -r requirements-build.txt git submodule update --init --recursive if [[ "$(uname -m)" == "aarch64" ]]; then export BUILD_IGNORE_SVE_UNAVAILABLE=1 fi - USE_DISTRIBUTED=1 python setup.py bdist_wheel + # PyTorch dropped setup.py. Build in a throwaway environment that can see the + # active one, so its build requirements, which pin a cmake that would be + # preferred over the one on PATH, cannot disturb what is installed here. + # + # Keep in sync with pytorch/pyproject.toml [build-system].requires. + local build_venv=/tmp/pytorch-build-venv + rm -rf "${build_venv}" + python -m venv --system-site-packages "${build_venv}" + "${build_venv}/bin/pip" install build "scikit-build-core>=1.0" ninja \ + "packaging>=24.2" "typing-extensions>=4.10.0" pyyaml six numpy + USE_DISTRIBUTED=1 "${build_venv}/bin/python" -m build --wheel --no-isolation + rm -rf "${build_venv}" pip install "$(echo dist/*.whl)" - - # Invariant: the basename setup.py just produced must match the cache + # A build with no BLAS succeeds silently, so check rather than assume. + (cd / && python -c " +import torch +assert torch._C.has_lapack, 'built without LAPACK' +torch.linalg.qr(torch.randn(4, 4)) +") + + # Invariant: the basename the build just produced must match the cache # URL we'd reconstruct on the next run. If they diverge (someone edits # torch_wheel_name above, or PyTorch renames its wheels), the cache # will silently miss and every macOS run will fall back to a ~30-min @@ -178,7 +191,7 @@ install_pytorch_and_domains() { # Grab the pinned audio and vision commits from PyTorch TORCHAUDIO_VERSION=release/2.11 export TORCHAUDIO_VERSION - TORCHVISION_VERSION=release/0.28 + TORCHVISION_VERSION=release/0.29 export TORCHVISION_VERSION install_domains diff --git a/.ci/scripts/wheel/cuda_arch_list.sh b/.ci/scripts/wheel/cuda_arch_list.sh index 29484056633..9e74af6c644 100644 --- a/.ci/scripts/wheel/cuda_arch_list.sh +++ b/.ci/scripts/wheel/cuda_arch_list.sh @@ -33,11 +33,13 @@ # that never claimed the device. So these lists are narrower than torch at the bottom on purpose. _cuda_arch_x86_64_cu130="8.0 8.6 8.9 9.0 10.0 12.0" _cuda_arch_x86_64_cu132="${_cuda_arch_x86_64_cu130}" +_cuda_arch_x86_64_cu134="${_cuda_arch_x86_64_cu130}" # The architectures the published aarch64 PyTorch CUDA build covers, read from its own library on an ARM # machine, for the same reason as the x86_64 rows above. Includes the ARM module whose train matches. _cuda_arch_aarch64_cu130="8.0 9.0 10.0 11.0 12.0" _cuda_arch_aarch64_cu132="${_cuda_arch_aarch64_cu130}" +_cuda_arch_aarch64_cu134="${_cuda_arch_aarch64_cu130}" # The older CUDA train. # @@ -104,6 +106,7 @@ executorch_cuda_arch_list() { 126) printf '%s' "${_cuda_arch_aarch64_cu126}" ;; 130) printf '%s' "${_cuda_arch_aarch64_cu130}" ;; 132) printf '%s' "${_cuda_arch_aarch64_cu132}" ;; + 134) printf '%s' "${_cuda_arch_aarch64_cu134}" ;; *) _executorch_unknown_train "${train}" ;; esac ;; @@ -112,6 +115,7 @@ executorch_cuda_arch_list() { 126) printf '%s' "${_cuda_arch_x86_64_cu126}" ;; 130) printf '%s' "${_cuda_arch_x86_64_cu130}" ;; 132) printf '%s' "${_cuda_arch_x86_64_cu132}" ;; + 134) printf '%s' "${_cuda_arch_x86_64_cu134}" ;; *) _executorch_unknown_train "${train}" ;; esac ;; diff --git a/.ci/scripts/wheel/test_cpp_sdk.py b/.ci/scripts/wheel/test_cpp_sdk.py index 29b2021d398..202a400e6c7 100644 --- a/.ci/scripts/wheel/test_cpp_sdk.py +++ b/.ci/scripts/wheel/test_cpp_sdk.py @@ -991,7 +991,7 @@ def test_every_shipped_header_compiles(work_dir: Path) -> None: # These say in their own text that they must not be included directly, and name the header to # include instead. Including one anyway is a use error rather than a packaging defect. "c10/util/complex_math.h", - "c10/util/complex_utils.h", + "torch/headeronly/util/complex_utils.h", ) source = work_dir / "header_probe.cpp" diff --git a/.ci/scripts/wheel/test_shared_libraries.py b/.ci/scripts/wheel/test_shared_libraries.py index 2ed6464105b..04646c7549c 100644 --- a/.ci/scripts/wheel/test_shared_libraries.py +++ b/.ci/scripts/wheel/test_shared_libraries.py @@ -2017,6 +2017,79 @@ def _names_a_build_directory(entry: str) -> bool: ) +# Absolute directories a shipped library may name. PyTorch's own is allowed because the wheel +# neither declares nor bundles PyTorch, so an absolute path is the only way to reach it. The maths +# library arch directories are allowed because a real installation spells them below a prefix, as +# /opt/intel/mkl/lib/intel64, which the environment genuinely provides. +# +# Matched as a suffix. A substring test exempted any path merely CONTAINING one of these, so a +# directory such as /home/user/torch/lib.backup/stage passed without reaching the build-directory +# classifier. Both "_win" spellings are listed explicitly now that the match is anchored. +_ALLOWED_ABSOLUTE_SUFFIXES = ( + "/torch/lib", + "/lib/intel64", + "/lib/intel64_win", + "/lib/win-x64", +) + +# The same arch directories with an EMPTY prefix, which is what PyTorch's exported link interface +# records when its MKL_ROOT resolves to nothing. They point nowhere, and they sit ahead of the +# relative hops packaging appends, which is the shadowing a CUDA toolkit prefix is rejected for. +# Matched exactly rather than folded into the allowlist above, because the two differ only in the +# prefix and a suffix match cannot tell them apart. +_UNRESOLVED_MATH_DIRECTORIES = ( + "/lib/intel64", + "/lib/intel64_win", + "/lib/win-x64", +) + +# Held for a wheel that bundles PyTorch's libraries rather than declaring them: such a copy records +# the CUDA toolkit directory of the machine that built IT, which is not this project's to fix. +# +# No wheel ships one today, so this clause never fires. It stays as a guard for a future +# bundling change; if that never comes, delete it rather than leaving an unexercised exemption. +_VENDORED_PREFIXES = ( + "libtorch", + "libc10", + "libshm", + "libcaffe2", + "libgomp", + "libiomp", +) + + +def _unusable_runtime_path_kind(entry: str, library_name: str) -> str | None: + """Why a recorded runtime search path is one a user cannot use, or None if it is fine. + + A library must not name an absolute directory the wheel has a relative route to. The one that + shipped was a CUDA toolkit prefix recorded on the build machine: it sat ahead of the relative + hop, so a user with a toolkit at the same prefix resolved the runtime from there instead of from + the declared dependency, and the builder always has one, so nothing exercised the hop. Stated as + a property rather than a list of known-bad directories, because a list only catches what someone + already thought of and that prefix was not on one. + + Order matters. The build-directory branch runs before the suffix allowlist so that a torch + directory inside a CI worker tree is rejected rather than accepted for its /torch/lib ending. + + Module scope so a unit test can call this directly and compare the reason it returns. Inline in + the caller's loop it could only be reached by building a wheel. + """ + if not entry: + # The loader reads an empty entry as the process working directory. + return "the process working directory" + if not entry.startswith("/") or library_name.startswith(_VENDORED_PREFIXES): + return None + if _names_a_build_directory(entry): + return "inside a build of this project" + if entry.rstrip("/") in _UNRESOLVED_MATH_DIRECTORIES: + return "a maths library directory whose prefix resolved empty" + if any( + entry.rstrip("/").endswith(allowed) for allowed in _ALLOWED_ABSOLUTE_SUFFIXES + ): + return None + return "an absolute directory the wheel has a relative route to" + + def test_no_absolute_runtime_paths() -> None: """No shipped library may search a directory a user does not have. @@ -2059,53 +2132,6 @@ def test_no_absolute_runtime_paths() -> None: package_dir = _installed_package_dir() - # This project's libraries must not name an absolute directory the wheel has a relative route to. The - # one that shipped was a CUDA toolkit prefix recorded on the build machine: it sat ahead of the relative - # hop, so a user with a toolkit at the same prefix resolved the CUDA runtime from there instead of from - # the declared dependency, and the builder always has one, so nothing exercised the hop. - # - # Stated as a property rather than a list of known-bad directories, because a list only catches what - # someone already thought of and that prefix was not on one. - # - # PyTorch's own directory is allowed: the wheel neither declares nor bundles PyTorch, so an absolute - # path is the only way to reach it. The maths library directories are allowed too. They arrive as - # -L flags in PyTorch's exported link interface, which CMake mirrors into the runtime path, so every - # library here that links PyTorch carries them. They point nowhere on any machine: measured on the - # link line as -L/lib/intel64 -L/lib/intel64_win -L/lib/win-x64, which is a prefix variable that - # resolved empty leaving the concatenation at the filesystem root. - # - # Matched as a suffix, the same way packaging decides what to strip at setup.py:1300. A substring - # test exempted any path merely CONTAINING one of these, so a directory such as - # /home/user/torch/lib.backup/stage passed without ever reaching the build-directory classifier. - # Both "_win" spellings are listed explicitly now that the match is anchored. - # - # A torch directory inside a CI worker tree, such as - # /home/ec2-user/actions-runner/_work/.../pytorch/torch/lib, is rejected rather than allowed: the - # build-directory classifier sees the worker components and the allowlist never gets to accept the - # /torch/lib suffix. Packaging strips the same entry, because every extension that names Torch now - # records a relative route to it. - allowed_absolute = ( - "/torch/lib", - "/lib/intel64", - "/lib/intel64_win", - "/lib/win-x64", - ) - # Held for a wheel that bundles PyTorch's libraries rather than declaring them: such a copy records - # the CUDA toolkit directory of the machine that built IT, which is not this project's to fix. - # - # No wheel ships one today. Six wheels across manylinux and macOS contain zero files with these - # prefixes, because the wheel declares torch as a dependency and there is no auditwheel step, so - # this clause is currently never false. It stays as a guard for a future bundling change; if that - # never comes, delete it rather than leaving an unexercised exemption in the check. - vendored_prefixes = ( - "libtorch", - "libc10", - "libshm", - "libcaffe2", - "libgomp", - "libiomp", - ) - offenders = {} inspected = 0 with_a_runtime_path = 0 @@ -2122,27 +2148,9 @@ def test_no_absolute_runtime_paths() -> None: with_a_runtime_path += 1 bad = [] for entry in entries: - if not entry: - bad.append("") - elif ( - entry.startswith("/") - and not library.name.startswith(vendored_prefixes) - and ( - _names_a_build_directory(entry) - or not any( - entry.rstrip("/").endswith(allowed) - for allowed in allowed_absolute - ) - ) - ): - # Named separately so the message says which kind it is: a build directory and a - # toolkit prefix are the same defect with different causes. - kind = ( - "inside a build of this project" - if _names_a_build_directory(entry) - else "an absolute directory the wheel has a relative route to" - ) - bad.append(f"{entry} ({kind})") + kind = _unusable_runtime_path_kind(entry, library.name) + if kind is not None: + bad.append(f"{entry or ''} ({kind})") if bad: offenders[str(library.relative_to(package_dir))] = bad diff --git a/.github/scripts/filter_cuda_matrix.py b/.github/scripts/filter_cuda_matrix.py index f7ac3a9debf..95d7a02120b 100644 --- a/.github/scripts/filter_cuda_matrix.py +++ b/.github/scripts/filter_cuda_matrix.py @@ -39,7 +39,7 @@ # not on that list is rejected whether or not it appears here. DISABLED_PYTHON_VERSIONS: List[str] = ["3.13t", "3.14t", "3.15", "3.15t"] -# CUDA versions to publish. +# CUDA versions to publish, when the generator offers them. # # Chosen so that every consumer row can find a matching wheel rather than by what is # convenient to verify. A delegate built against one of these has to be able to depend on an @@ -48,13 +48,17 @@ # # cu126 the floor, and what Jetson devices are limited to # cu130 the generator's stable choice, and the default for accelerator consumers -# cu132 the newest, which consumers building against a current TensorRT need +# cu132 a current TensorRT build target +# cu134 the newest, which consumers building against the latest CUDA need +# +# Skip wholly absent trains so an upstream removal cannot block the remaining releases. +# Offered trains must still cover every supported Python version. # # cu132 is included because omitting it would leave a published consumer row with no # ExecuTorch wheel to pair with. It is executable on a device one minor behind, since CUDA # minor versions are compatible, so a cu132 wheel has been run end to end on a CUDA 13.0 # device. The packaging properties are checked on every row regardless. -SUPPORTED_CUDA_VERSIONS: List[str] = ["cu126", "cu130", "cu132"] +SUPPORTED_CUDA_VERSIONS: List[str] = ["cu126", "cu130", "cu132", "cu134"] # Python versions to publish, stated rather than derived for the same reason the CUDA # versions are. Deriving them from the rows that survived the filter made the release @@ -184,45 +188,30 @@ def main(argv: List[str]) -> None: if args.limit_pr_builds.lower() == "true" and items: items = only_pull_request_row(items) elif items and not is_jetpack: - # A release has to publish every combination this policy advertises. Comparing the result against - # what the generator offered cannot catch anything, because both sides apply the same conditions, so - # the difference is empty by construction and the check never fires. The policy's own list is the - # thing to compare against: a CUDA version the generator stopped offering otherwise disappears from - # the release silently, and a missing job is a green check for a wheel that was never built. - # - # The generic rows only. A JetPack release advertises the single pair its own lists name rather than - # every supported CUDA version, so checking it against this list would fail a correct release. - # - # Both axes come from this policy's own lists, not from the matrix. Reading the generator's python - # axis pulled in rows this policy never builds, and deriving it from the rows that survived went - # blind to a python that disappeared from every supported train. The generator lives in another - # repository and its axes move independently of what this policy promises to publish. built = {(item["python_version"], item["desired_cuda"]) for item in items} - # A train that produced no row at all is missing for every python, so reporting it per python - # would read as a python problem. Named on its own instead, and first, because the per-pair - # report below would otherwise bury it. - absent_trains = sorted( - set(SUPPORTED_CUDA_VERSIONS) - {cuda for _, cuda in built} - ) + # Filtering out every Python row must not disguise an offered train as absent. + offered_trains = { + item["desired_cuda"] + for item in matrix.get("include", []) + if item["desired_cuda"] in SUPPORTED_CUDA_VERSIONS + } + absent_trains = sorted(set(SUPPORTED_CUDA_VERSIONS) - offered_trains) if absent_trains: print( - f"this policy publishes {SUPPORTED_CUDA_VERSIONS}, but the generator offered no row " - f"this filter could keep for {absent_trains}, so a release would publish no wheel for " - "that CUDA version at all", + f"the generator offered no row for {absent_trains}, so they are skipped this run; " + f"publishing {sorted(offered_trains)}", file=sys.stderr, ) - sys.exit(1) missing = sorted( f"{python}/{cuda}" for python in SUPPORTED_PYTHON_VERSIONS - for cuda in SUPPORTED_CUDA_VERSIONS + for cuda in offered_trains if (python, cuda) not in built ) if missing: print( - f"this policy publishes {SUPPORTED_CUDA_VERSIONS} for each of " - f"{SUPPORTED_PYTHON_VERSIONS}, but {len(missing)} combination(s) produced no row, so a " - f"release would publish no wheel for them: {missing}", + f"a published train is missing some of {SUPPORTED_PYTHON_VERSIONS}, so a release " + f"would ship an incomplete train: {missing}", file=sys.stderr, ) sys.exit(1) diff --git a/.github/workflows/_docker-image.yml b/.github/workflows/_docker-image.yml index a328a8fc3ca..9dc4813e16e 100644 --- a/.github/workflows/_docker-image.yml +++ b/.github/workflows/_docker-image.yml @@ -37,10 +37,22 @@ jobs: - name: Checkout ExecuTorch uses: actions/checkout@v4 with: - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + # build-cadence-runner.yml reaches this workflow on pull_request_target, + # where github.sha is the base branch tip. Resolving the tag from there + # while the test jobs check out the fork head would test a docker change + # against the old image. + ref: ${{ (github.event_name == 'pull_request' || github.event_name == 'pull_request_target') && github.event.pull_request.head.sha || github.sha }} + # Nothing here runs the checked-out code, only `git rev-parse` over it, + # but on pull_request_target that code is the fork's and the token is + # the base repository's, so do not leave one next to the other. + persist-credentials: false - name: Compute the docker tag id: hash run: | set -eu - echo "ci-docker-hash=$(git rev-parse HEAD:.ci/docker)" >> "$GITHUB_OUTPUT" + # Assigned on its own line: set -e does not catch a failure inside a + # command substitution, so `echo "x=$(git rev-parse ...)"` would write + # git's error text as the tag and exit 0. + CI_DOCKER_HASH="$(git rev-parse HEAD:.ci/docker)" + echo "ci-docker-hash=${CI_DOCKER_HASH}" >> "$GITHUB_OUTPUT" diff --git a/.github/workflows/_llm_server.yml b/.github/workflows/_llm_server.yml index e1ef5a30db9..6492a3106ee 100644 --- a/.github/workflows/_llm_server.yml +++ b/.github/workflows/_llm_server.yml @@ -4,20 +4,25 @@ on: workflow_call: inputs: docker-image: - description: Docker image to use for Linux tests. + description: Name of the docker image to use, without registry or tag suffix. required: false type: string default: ci-image:executorch-ubuntu-22.04-clang12 jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + linux: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ${{ inputs.docker-image }} + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/${{ inputs.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: recursive ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 60 diff --git a/.github/workflows/_test_arduino_library.yml b/.github/workflows/_test_arduino_library.yml index a3e73958463..5001202c069 100644 --- a/.github/workflows/_test_arduino_library.yml +++ b/.github/workflows/_test_arduino_library.yml @@ -14,15 +14,20 @@ on: default: 90 jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + run: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: job-name: arduino-library - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: ${{ inputs.timeout }} @@ -53,9 +58,9 @@ jobs: # moving branch: reproducible, and one less thing to fail transiently. ARDUINO_CLI_VERSION=1.5.1 ARDUINO_LINT_VERSION=1.3.0 - # RUNNER_TEMP and GITHUB_WORKSPACE both point at host paths this - # container cannot write to, and HOME may too, so keep the toolchain - # and every arduino-cli directory somewhere local to the container. + # HOME is /github/home, which arduino-cli's directories should not share + # with the runner's own state, so keep the toolchain and every + # arduino-cli directory somewhere local to the container. ARDUINO_CI_DIR=/tmp/.arduino-ci export ARDUINO_DIRECTORIES_USER="${ARDUINO_CI_DIR}/user" export ARDUINO_DIRECTORIES_DATA="${ARDUINO_CI_DIR}/data" diff --git a/.github/workflows/_test_backend.yml b/.github/workflows/_test_backend.yml index 18c70a31531..063c0ebb161 100644 --- a/.github/workflows/_test_backend.yml +++ b/.github/workflows/_test_backend.yml @@ -40,15 +40,26 @@ on: description: 'Runner type for Linux jobs' required: false type: string - default: linux.4xlarge.memory + default: mt-l-x86iavx512-16-128 docker-image: - description: 'Docker image for Linux jobs' + description: 'Name of the docker image to use, without registry or tag suffix' required: false type: string default: ci-image:executorch-ubuntu-22.04-clang12 + save-goldens: + description: 'Write golden .pte/.bin files in the models suite and package them; only for regenerating the Android test fixture' + required: false + type: boolean + default: false jobs: + docker-image: + name: Resolve CI docker image + if: ${{ inputs.run-linux }} + uses: ./.github/workflows/_docker-image.yml + test-backend-linux: + needs: docker-image if: ${{ inputs.run-linux }} strategy: fail-fast: false @@ -57,21 +68,27 @@ jobs: suite: [models, operators] exclude: ${{ fromJSON(inputs.exclude) }} - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read with: ref: ${{ inputs.ref }} runner: ${{ inputs.runner-linux }} - docker-image: ${{ inputs.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/${{ inputs.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: recursive timeout: ${{ inputs.timeout }} upload-artifact: test-report-${{ inputs.backend }}-${{ matrix.flow }}-${{ matrix.suite }} script: | set -eux + if [[ "${{ inputs.save-goldens }}" == "true" && "${{ matrix.suite }}" == "models" ]]; then + export GOLDEN_ARTIFACTS_DIR="${RUNNER_ARTIFACT_DIR}/golden-artifacts" + fi source .ci/scripts/test_backend.sh "${{ matrix.suite }}" "${{ matrix.flow }}" "${RUNNER_ARTIFACT_DIR}" package-golden-artifacts: - if: ${{ inputs.run-linux }} + if: ${{ inputs.run-linux && inputs.save-goldens }} needs: test-backend-linux runs-on: linux.2xlarge steps: @@ -107,13 +124,6 @@ jobs: echo "No golden artifacts found." fi - - name: Upload combined golden artifacts - uses: actions/upload-artifact@v4 - with: - name: golden-artifacts-${{ inputs.backend }} - path: golden_artifacts_*.zip - if-no-files-found: ignore - - name: Upload golden artifacts to S3 uses: seemethere/upload-artifact-s3@v5 if: ${{ hashFiles('golden_artifacts_*.zip') != '' }} diff --git a/.github/workflows/_test_cadence.yml b/.github/workflows/_test_cadence.yml index 2e98d21db1c..d899e2ae585 100644 --- a/.github/workflows/_test_cadence.yml +++ b/.github/workflows/_test_cadence.yml @@ -8,7 +8,7 @@ on: workflow_call: inputs: docker-image: - description: 'Docker image to use' + description: 'Name of the docker image to use, without registry or tag suffix' required: false type: string default: ci-image:executorch-ubuntu-22.04-clang12 @@ -16,7 +16,7 @@ on: description: 'Runner type' required: false type: string - default: linux.8xlarge.memory + default: mt-l-x86iavx512-32-256 ref: description: 'Git ref to checkout' required: false @@ -29,12 +29,20 @@ on: default: 90 jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + test-aot: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read with: job-name: test-aot runner: ${{ inputs.runner }} - docker-image: ${{ inputs.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/${{ inputs.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: recursive ref: ${{ inputs.ref }} timeout: ${{ inputs.timeout }} @@ -50,11 +58,15 @@ jobs: python -m pytest backends/cadence/aot/tests/ -v -n auto --reruns 2 --reruns-delay 1 test-ops: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read with: job-name: test-ops runner: ${{ inputs.runner }} - docker-image: ${{ inputs.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/${{ inputs.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: recursive ref: ${{ inputs.ref }} timeout: ${{ inputs.timeout }} diff --git a/.github/workflows/_test_cortex_m_e2e.yml b/.github/workflows/_test_cortex_m_e2e.yml index 0510b017723..1957a2b0c56 100644 --- a/.github/workflows/_test_cortex_m_e2e.yml +++ b/.github/workflows/_test_cortex_m_e2e.yml @@ -23,8 +23,16 @@ on: default: 120 jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + run: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read strategy: matrix: model: ${{ fromJSON(inputs.models) }} @@ -32,8 +40,8 @@ jobs: fail-fast: false with: job-name: ${{ matrix.model }}-${{ matrix.target }} - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: ${{ inputs.timeout }} diff --git a/.github/workflows/_test_cortex_m_ops.yml b/.github/workflows/_test_cortex_m_ops.yml index a9e2e6180c3..885f5fb64d9 100644 --- a/.github/workflows/_test_cortex_m_ops.yml +++ b/.github/workflows/_test_cortex_m_ops.yml @@ -18,16 +18,24 @@ on: default: 120 jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + run: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read strategy: matrix: target: ${{ fromJSON(inputs.targets) }} fail-fast: false with: job-name: cortex-m-ops-${{ matrix.target }} - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: ${{ inputs.timeout }} diff --git a/.github/workflows/_test_riscv.yml b/.github/workflows/_test_riscv.yml index 223a146e3d8..954cb94d596 100644 --- a/.github/workflows/_test_riscv.yml +++ b/.github/workflows/_test_riscv.yml @@ -37,11 +37,19 @@ on: type: string jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + run: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-24.04-gcc14 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-24.04-gcc14-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: ${{ inputs.timeout }} diff --git a/.github/workflows/add-unanswered-to-project.yml b/.github/workflows/add-unanswered-to-project.yml index 4a5702b2826..acbdb906b2b 100644 --- a/.github/workflows/add-unanswered-to-project.yml +++ b/.github/workflows/add-unanswered-to-project.yml @@ -11,11 +11,12 @@ on: pull_request: paths: - .github/workflows/add-unanswered-to-project.yml + jobs: add_to_project: runs-on: ubuntu-latest steps: - - name: Add open issues and open, non-draft PRs to org project (excluding certain authors and bots) + - name: Add open issues and open, non-draft PRs to org project and label external PRs (excluding certain authors and bots) uses: actions/github-script@v7 with: github-token: ${{ secrets.ET_EXT_CONTRIB }} @@ -26,36 +27,36 @@ jobs: // List of authors to exclude const excludedAuthors = new Set([ - "nil-is-all", "tanvirislam-meta", "cbilgin", "kimishpatel", "psiddh", "digantdesai", "SS-JIA", "ahmtox", "mcr229", - "shoumikhin", "manuelcandales", "metascroy", "cccclai", "rohansjoshi", "kirklandsign", "abhinaykukkadapu", "JacobSzwejbka", - "Conarnar", "rascani", "kiymetakdemir", "JCNTH", "lucylq", "larryliu0820", "BujSet", "Gasoonjia", "Juntian777", "guangy10", - "jackzhxng", "GregoryComer", "leafs1", "swolchok", "mergennachin", "tarun292", "byjlw", "jathu", "Jack-Khuu", "georgehong", - "zhenyan-zhang-meta", "silverguo", "harishs88ss", "AlannaBurke", "Doggeral", "laithsakka", "Reubend", "dbort", "huydhn", "mcremon-meta", - "trivedivivek", "angelayi", "helunwencser", "hsharma35", "zhxchen17", "iseeyuan", "svekars", "nathanaelsee", "dulinriley", - "jerryzh168", "cmodi-meta", "bigfootjon", "sxu", "ydwu4", "Riandy", "tugsbayasgalan", "bsoyluoglu", "yangw-dev", + "nil-is-all", "tanvirislam-meta", "cbilgin", "kimishpatel", "psiddh", "digantdesai", "SS-JIA", "ahmtox", "mcr229", + "shoumikhin", "manuelcandales", "metascroy", "cccclai", "rohansjoshi", "kirklandsign", "abhinaykukkadapu", "JacobSzwejbka", + "Conarnar", "rascani", "kiymetakdemir", "JCNTH", "lucylq", "larryliu0820", "BujSet", "Gasoonjia", "Juntian777", "guangy10", + "jackzhxng", "GregoryComer", "leafs1", "swolchok", "mergennachin", "tarun292", "byjlw", "jathu", "Jack-Khuu", "georgehong", + "zhenyan-zhang-meta", "silverguo", "harishs88ss", "AlannaBurke", "Doggeral", "laithsakka", "Reubend", "dbort", "huydhn", "mcremon-meta", + "trivedivivek", "angelayi", "helunwencser", "hsharma35", "zhxchen17", "iseeyuan", "svekars", "nathanaelsee", "dulinriley", + "jerryzh168", "cmodi-meta", "bigfootjon", "sxu", "ydwu4", "Riandy", "tugsbayasgalan", "bsoyluoglu", "yangw-dev", "YIWENX14", "namanahuja", "yushangdi", "limintang", "pianpwk", "viveknayakatmeta", "andreanicastro", "JakeStevens", "gmagogsfm", "zonglinpeng", "eigen-k", "derekxu", "salilsdesai", "skrtskrtfb", "pssrawat", "r-barnes", "kalpit-meta-1", "Will-MingLun-Li", "KapJI", "piyengar", "j-bahr", "BoyuanFeng", "fgasperij", "DariusHolmgren", "sammarden-meta", "kushrast", "meta-emilian", "Rittzz", "jeanschmidt", "copyrightly", "mikekgfb", "vmpuri", - "zonglinpengmeta", "maggiemoss", "aorenste", "hoangminhle98", "Solumin", "meyering", "rchen152", "AishwaryaSivaraman", - "migeed-z", "ebgraham", "Esteb37", "nausicaasnow", "Camyll", "ezyang", "huiyujie", "dltn", "cjhopman", "blackm00n", - "agunapal", "SamGondelman", "Ninja91", "ivayloen", "DrJessop", "rodrigos01meta", "akrieger", "cmt0", "yiming0416", - "ethansfng", "ThomasJannaud", "nirvanagth", "marcinkwiatkowski", "3l1", "omerjerk", "nitish2112", "yipjustin", - "ejnguyen", "andrewor14", "phaiting", "mgiordy", "LeeOHzzZ", "adicatana", "Polyomino", "ezrilow", "navsud", - "michaelmaitland", "RahulC7", "seyeong-han", "thdusdl1219", "jaejunku", "felixweilbach", "apullin", "trviv", "junluan01", - "mvartani-meta", "abeakkas", "elpdumont", "corporateshark", "bdemirb", "GeorgeTzoupis", "AdithyaReddy9", "drinkmorewaterr", - "aliafzal", "YifanShenSZ", "RdoubleA", "Olivia-liu", "Abhi-hpp", "Vysarat","azad-meta", "junpi", - "pytorchbot", "pytorchmergebot", "pytorchupdatebot", "facebook-github-bot", "app/dependabot", - "Erik-Lundell", "zingo", "AdrianLundell", "oscarandersson8218", "per", "Sebastian-Larsson", "SaoirseARM", "robell", - "mansnils", "martinlsm", "freddan80", "YufengShi-dudu", "tom-arm", "perheld", "Jerry-Ge", "gggekov", "fumchin", "wwwind", - "benkli01", "Tessil", "maddun01", "Michiel-Olieslagers", "armwaheed", "agrima1304", "emmakujala", "annietllnd", - "MatthiasHertel80", "AlexTawseArm", "jmahbs", "morgolock", "Christoffer-JL", "ArmRyan", "xingguo01", "tgonzalezorlandoarm", + "zonglinpengmeta", "maggiemoss", "aorenste", "hoangminhle98", "Solumin", "meyering", "rchen152", "AishwaryaSivaraman", + "migeed-z", "ebgraham", "Esteb37", "nausicaasnow", "Camyll", "ezyang", "huiyujie", "dltn", "cjhopman", "blackm00n", + "agunapal", "SamGondelman", "Ninja91", "ivayloen", "DrJessop", "rodrigos01meta", "akrieger", "cmt0", "yiming0416", + "ethansfng", "ThomasJannaud", "nirvanagth", "marcinkwiatkowski", "3l1", "omerjerk", "nitish2112", "yipjustin", + "ejnguyen", "andrewor14", "phaiting", "mgiordy", "LeeOHzzZ", "adicatana", "Polyomino", "ezrilow", "navsud", + "michaelmaitland", "RahulC7", "seyeong-han", "thdusdl1219", "jaejunku", "felixweilbach", "apullin", "trviv", "junluan01", + "mvartani-meta", "abeakkas", "elpdumont", "corporateshark", "bdemirb", "GeorgeTzoupis", "AdithyaReddy9", "drinkmorewaterr", + "aliafzal", "YifanShenSZ", "RdoubleA", "Olivia-liu", "Abhi-hpp", "Vysarat","azad-meta", "junpi", + "pytorchbot", "pytorchmergebot", "pytorchupdatebot", "facebook-github-bot", "app/dependabot", + "Erik-Lundell", "zingo", "AdrianLundell", "oscarandersson8218", "per", "Sebastian-Larsson", "SaoirseARM", "robell", + "mansnils", "martinlsm", "freddan80", "YufengShi-dudu", "tom-arm", "perheld", "Jerry-Ge", "gggekov", "fumchin", "wwwind", + "benkli01", "Tessil", "maddun01", "Michiel-Olieslagers", "armwaheed", "agrima1304", "emmakujala", "annietllnd", + "MatthiasHertel80", "AlexTawseArm", "jmahbs", "morgolock", "Christoffer-JL", "ArmRyan", "xingguo01", "tgonzalezorlandoarm", "chizkiyahu", "sarah-blades", "itsMarco-G", "usamahz", "Rob-Hughes-Arm", "swha815", "FabulousSuperDude", - "haowhsu-quic", "shewu-quic", "winskuo-quic", "chunit-quic", "DannyYuyang-quic", "chuntl", "thchenqti", "jethroqti", - "chenweng-quic", "qti-horodnic", "qti-mmadhava", "quic-boyuc", "zhaoxul-qti", - "cymbalrush", "DenisVieriu97", "billmguo", - "StrycekSimon", "jirioc", "robert-kalmar", "skywall", "MartinPavella", "roman-janik-nxp", "novak-vaclav", "irtrukhina", - "neuropilot-captain", "dijopaul", "cad-rlc", "cad-audio", "ynimmaga", "daniil-lyakhov", + "haowhsu-quic", "shewu-quic", "winskuo-quic", "chunit-quic", "DannyYuyang-quic", "chuntl", "thchenqti", "jethroqti", + "chenweng-quic", "qti-horodnic", "qti-mmadhava", "quic-boyuc", "zhaoxul-qti", + "cymbalrush", "DenisVieriu97", "billmguo", + "StrycekSimon", "jirioc", "robert-kalmar", "skywall", "MartinPavella", "roman-janik-nxp", "novak-vaclav", "irtrukhina", + "neuropilot-captain", "dijopaul", "cad-rlc", "cad-audio", "ynimmaga", "daniil-lyakhov", "emmanuel-ferdman", "cavusmustafa", "anzr299", "suryasidd", "Jiseong-oh", "alexdean08", // explicitly include the dependabot bot login seen in PRs "dependabot[bot]" @@ -69,7 +70,10 @@ jobs: // Labels on PRs to exclude from being added to the project const excludedPrLabels = new Set(["fb-exported", "meta-exported"]); - + + // Label applied to PRs that pass every external-contributor check below + const communityLabel = "community: contribution"; + // Simple cache for user -> boolean (member of excluded org) const orgsCache = new Map(); const companyCache = new Map(); @@ -125,6 +129,31 @@ jobs: return false; } } + + function hasCommunityLabel(item) { + if (!item || !item.labels) return false; + return item.labels.some(l => l && l.name && l.name.toLowerCase() === communityLabel); + } + + async function addCommunityLabel(pr) { + if (hasCommunityLabel(pr)) { + console.log(`PR #${pr.number} already has "${communityLabel}"`); + return; + } + try { + await github.rest.issues.addLabels({ + owner, + repo, + issue_number: pr.number, + labels: [communityLabel] + }); + console.log(`Labeled PR #${pr.number} with "${communityLabel}"`); + } catch (error) { + // Labeling is best-effort: a failure here must not stop the rest of the run. + console.log(`Error labeling PR #${pr.number}: ${error.message}`); + } + } + async function addItem(contentId, type, number) { try { await github.graphql(` @@ -156,6 +185,7 @@ jobs: filter: 'all' } ); + for (const issue of issues) { if (issue.pull_request) { console.log(`Skipping PR #${issue.number} (listed in issues)`); @@ -185,6 +215,7 @@ jobs: state: 'open', } ); + for (const pr of prs) { if (pr.draft) { console.log(`Skipping PR #${pr.number} (draft)`); @@ -207,6 +238,7 @@ jobs: continue; } await addItem(pr.node_id, 'pr', pr.number); + await addCommunityLabel(pr); } } catch (error) { core.setFailed(`Workflow failed: ${error.message}`); diff --git a/.github/workflows/build-cmsis-pack.yml b/.github/workflows/build-cmsis-pack.yml index c9c40d670f1..9974ceea4ee 100644 --- a/.github/workflows/build-cmsis-pack.yml +++ b/.github/workflows/build-cmsis-pack.yml @@ -40,7 +40,7 @@ on: type: string concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true jobs: diff --git a/.github/workflows/build-presets.yml b/.github/workflows/build-presets.yml index 89d36cd6b0a..48c745121be 100644 --- a/.github/workflows/build-presets.yml +++ b/.github/workflows/build-presets.yml @@ -9,7 +9,7 @@ on: workflow_dispatch: concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true jobs: diff --git a/.github/workflows/check-c10-sync.yml b/.github/workflows/check-c10-sync.yml index 73a4837adf6..db62e15cc4b 100644 --- a/.github/workflows/check-c10-sync.yml +++ b/.github/workflows/check-c10-sync.yml @@ -8,7 +8,7 @@ on: - runtime/core/portable_type/c10/** concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true jobs: diff --git a/.github/workflows/check-labels.yml b/.github/workflows/check-labels.yml index ebaa38cf0bd..b103f734a20 100644 --- a/.github/workflows/check-labels.yml +++ b/.github/workflows/check-labels.yml @@ -25,7 +25,7 @@ on: required: true concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }} cancel-in-progress: true jobs: diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 21c0fe6e844..0c134b30e05 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -1,6 +1,6 @@ # Test ExecuTorch CUDA Build Compatibility # This workflow tests whether ExecuTorch can be successfully built with CUDA support -# across different CUDA versions (12.6, 13.0) using the command: +# across different CUDA versions (12.6, 13.0, 13.4) using the command: # ./install_executorch.sh # # Intentionally skipped CUDA version 13.2 check due to ci image unsupported. @@ -64,10 +64,10 @@ jobs: strategy: fail-fast: false matrix: - cuda-version: ["12.6", "13.0"] + cuda-version: ["12.6", "13.0", "13.4"] name: test-executorch-cuda-build-${{ matrix.cuda-version }} - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@c3c4c4d48e97dbaaabd9b65496afd8c8d8dad713 permissions: id-token: write contents: read @@ -76,14 +76,19 @@ jobs: runner: linux.g5.4xlarge.nvidia.gpu gpu-arch-type: cuda gpu-arch-version: ${{ matrix.cuda-version }} + test-infra-ref: c3c4c4d48e97dbaaabd9b65496afd8c8d8dad713 + driver-version: ${{ matrix.cuda-version == '13.4' && '615.71.09' || '580.65.06' }} + driver-download-url: ${{ matrix.cuda-version == '13.4' && 'https://download.nvidia.com/XFree86/Linux-x86_64/615.71.09/NVIDIA-Linux-x86_64-615.71.09.run' || '' }} use-custom-docker-registry: false submodules: recursive ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} script: | set -eux - # Test ExecuTorch CUDA build - ExecuTorch will automatically detect CUDA version - # and install the appropriate PyTorch wheel + if [ "${{ matrix.cuda-version }}" = "13.4" ]; then + # The image's older Conda runtime aborts during CUDA thread-local cleanup. + conda install -y -n base -c conda-forge 'libstdcxx-ng=16.2.0' 'libgcc-ng=16.2.0' + fi source .ci/scripts/test-cuda-build.sh "${{ matrix.cuda-version }}" # This job will fail if any of the CUDA versions fail @@ -112,7 +117,7 @@ jobs: echo "CUDA build results: ${{ needs.test-cuda-builds.result }}" exit 1 else - echo "SUCCESS: All ExecuTorch CUDA builds (12.6, 13.0) completed successfully!" + echo "SUCCESS: All ExecuTorch CUDA builds completed successfully!" fi test-models-cuda: @@ -160,7 +165,7 @@ jobs: done export-cuda-target-smem-cross-arch: - name: export-cuda-target-smem-cross-arch-a100 + name: export-cuda-multi-arch-a100 needs: [changed-files, run-decision] if: | contains(needs.changed-files.outputs.changed-files, 'backends/cuda') || @@ -178,15 +183,13 @@ jobs: gpu-arch-version: "13.0" use-custom-docker-registry: false submodules: recursive - upload-artifact: cuda-target-smem-cross-arch + upload-artifact: cuda-multi-arch-a100 ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} script: | set -eux PYTHON_EXECUTABLE=python ./install_executorch.sh export LD_LIBRARY_PATH="/opt/conda/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" - # A10G (sm_86) supports 99 KiB of opt-in shared memory per block. - export ET_CUDA_TARGET_SMEM_BYTES=101376 # The A100 exporter has an AVX-512 host CPU while the A10G runner has # AVX2. Keep the generated AOTI host wrapper compatible with both. export ATEN_CPU_CAPABILITY=avx2 @@ -199,16 +202,33 @@ jobs: assert capability == (8, 0), f"Expected A100 sm_80 exporter, got {capability}" PY + # Export an sm80 native PTE with no PTX. This artifact must fail on the + # A10G because its metadata has neither an sm86 native variant nor PTX. + model_dir="${RUNNER_ARTIFACT_DIR}/native-sm80/linear" + mkdir -p "${model_dir}" + TORCHINDUCTOR_CACHE_DIR="${RUNNER_TEMP}/inductor-native-sm80" \ + python -m examples.cuda.scripts.export \ + --model_name=linear \ + --output_dir="${model_dir}" \ + --cuda_include_ptx=OFF \ + --seed=0 + + # Export the explicit PTX fallback separately. A10G supports 99 KiB of + # opt-in shared memory per block, so constrain the A100 export to it. for model in linear sdpa; do - model_dir="${RUNNER_ARTIFACT_DIR}/${model}" + model_dir="${RUNNER_ARTIFACT_DIR}/fallback-sm80/${model}" mkdir -p "${model_dir}" - python -m examples.cuda.scripts.export \ + ET_CUDA_TARGET_SMEM_BYTES=101376 \ + TORCHINDUCTOR_CACHE_DIR="${RUNNER_TEMP}/inductor-fallback-sm80-${model}" \ + python -m examples.cuda.scripts.export \ --model_name="${model}" \ - --output_dir "${model_dir}" + --output_dir="${model_dir}" \ + --cuda_include_ptx=ON \ + --seed=0 done test-cuda-target-smem-cross-arch: - name: test-cuda-target-smem-cross-arch-a10g + name: merge-and-test-cuda-multi-arch-a10g needs: [export-cuda-target-smem-cross-arch] if: needs.export-cuda-target-smem-cross-arch.result == 'success' uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main @@ -222,7 +242,8 @@ jobs: gpu-arch-version: "13.0" use-custom-docker-registry: false submodules: recursive - download-artifact: cuda-target-smem-cross-arch + download-artifact: cuda-multi-arch-a100 + upload-artifact: cuda-multi-arch-merged ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} script: | set -eux @@ -243,20 +264,141 @@ jobs: ) PY + # Independently export an sm86 native PTE with no PTX. A fixed seed + # makes its external weights byte-identical to both A100 exports. + model_dir="${RUNNER_ARTIFACT_DIR}/native-sm86/linear" + mkdir -p "${model_dir}" + TORCHINDUCTOR_CACHE_DIR="${RUNNER_TEMP}/inductor-native-sm86" \ + python -m examples.cuda.scripts.export \ + --model_name=linear \ + --output_dir="${model_dir}" \ + --cuda_include_ptx=OFF \ + --seed=0 + + sm80_ptd="${RUNNER_ARTIFACT_DIR}/native-sm80/linear/aoti_cuda_blob.ptd" + sm86_ptd="${RUNNER_ARTIFACT_DIR}/native-sm86/linear/aoti_cuda_blob.ptd" + fallback_ptd="${RUNNER_ARTIFACT_DIR}/fallback-sm80/linear/aoti_cuda_blob.ptd" + cmp "${sm80_ptd}" "${sm86_ptd}" + cmp "${sm80_ptd}" "${fallback_ptd}" + sha256sum "${sm80_ptd}" "${sm86_ptd}" "${fallback_ptd}" + + merged_dir="${RUNNER_ARTIFACT_DIR}/merged" + mkdir -p "${merged_dir}" + python -m executorch.backends.cuda.merge_ptes \ + --input-pte "${RUNNER_ARTIFACT_DIR}/native-sm80/linear/linear.pte" \ + --input-pte "${RUNNER_ARTIFACT_DIR}/native-sm86/linear/linear.pte" \ + --input-ptd "${sm80_ptd}" \ + --input-ptd "${sm86_ptd}" \ + --fallback-pte "${RUNNER_ARTIFACT_DIR}/fallback-sm80/linear/linear.pte" \ + --fallback-ptd "${fallback_ptd}" \ + --output-pte "${merged_dir}/linear.pte" \ + --output-ptd "${merged_dir}/aoti_cuda_blob.ptd" \ + | tee "${merged_dir}/provenance.txt" + cmp "${sm80_ptd}" "${merged_dir}/aoti_cuda_blob.ptd" + cmake -DCMAKE_BUILD_TYPE=Release \ -DEXECUTORCH_BUILD_CUDA=ON \ -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -DEXECUTORCH_ENABLE_LOGGING=ON \ -DPYTHON_EXECUTABLE=python \ -Bcmake-out . cmake --build cmake-out --target executor_runner -j4 for model in linear sdpa; do - model_dir="${RUNNER_ARTIFACT_DIR}/${model}" + model_dir="${RUNNER_ARTIFACT_DIR}/fallback-sm80/${model}" ./cmake-out/executor_runner \ --model_path "${model_dir}/${model}.pte" \ --data_path "${model_dir}/aoti_cuda_blob.ptd" done + # The local native PTE and merged PTE must run on sm86. + ./cmake-out/executor_runner \ + --model_path "${RUNNER_ARTIFACT_DIR}/native-sm86/linear/linear.pte" \ + --data_path "${sm86_ptd}" + ./cmake-out/executor_runner \ + --model_path "${merged_dir}/linear.pte" \ + --data_path "${merged_dir}/aoti_cuda_blob.ptd" + + # The sm80 native-only PTE must be rejected on sm86. + set +e + ./cmake-out/executor_runner \ + --model_path "${RUNNER_ARTIFACT_DIR}/native-sm80/linear/linear.pte" \ + --data_path "${sm80_ptd}" \ + >"${RUNNER_ARTIFACT_DIR}/native-sm80-on-sm86.log" 2>&1 + native_status=$? + set -e + cat "${RUNNER_ARTIFACT_DIR}/native-sm80-on-sm86.log" + test "${native_status}" -ne 0 + grep -q "no native or PTX variant compatible with sm86" \ + "${RUNNER_ARTIFACT_DIR}/native-sm80-on-sm86.log" + + test-cuda-merged-pte-a100: + name: test-cuda-merged-pte-a100 + needs: [test-cuda-target-smem-cross-arch] + if: needs.test-cuda-target-smem-cross-arch.result == 'success' + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read + with: + timeout: 90 + runner: mt-l-x86iavx512-11-125-a100 + gpu-arch-type: cuda + gpu-arch-version: "13.0" + use-custom-docker-registry: false + submodules: recursive + download-artifact: cuda-multi-arch-merged + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + script: | + set -eux + + PYTHON_EXECUTABLE=python ./install_executorch.sh + export LD_LIBRARY_PATH="/opt/conda/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + + python - <<'PY' + import torch + + capability = torch.cuda.get_device_capability() + assert capability == (8, 0), f"Expected A100 sm_80 runner, got {capability}" + PY + + cmake -DCMAKE_BUILD_TYPE=Release \ + -DEXECUTORCH_BUILD_CUDA=ON \ + -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -DEXECUTORCH_ENABLE_LOGGING=ON \ + -DPYTHON_EXECUTABLE=python \ + -Bcmake-out . + cmake --build cmake-out --target executor_runner -j4 + + sm80_ptd="${RUNNER_ARTIFACT_DIR}/native-sm80/linear/aoti_cuda_blob.ptd" + sm86_ptd="${RUNNER_ARTIFACT_DIR}/native-sm86/linear/aoti_cuda_blob.ptd" + fallback_ptd="${RUNNER_ARTIFACT_DIR}/fallback-sm80/linear/aoti_cuda_blob.ptd" + merged_ptd="${RUNNER_ARTIFACT_DIR}/merged/aoti_cuda_blob.ptd" + cmp "${sm80_ptd}" "${sm86_ptd}" + cmp "${sm80_ptd}" "${fallback_ptd}" + cmp "${sm80_ptd}" "${merged_ptd}" + + # The local native PTE and merged PTE must run on sm80. + ./cmake-out/executor_runner \ + --model_path "${RUNNER_ARTIFACT_DIR}/native-sm80/linear/linear.pte" \ + --data_path "${sm80_ptd}" + ./cmake-out/executor_runner \ + --model_path "${RUNNER_ARTIFACT_DIR}/merged/linear.pte" \ + --data_path "${merged_ptd}" + + # The sm86 native-only PTE must be rejected on sm80. + set +e + ./cmake-out/executor_runner \ + --model_path "${RUNNER_ARTIFACT_DIR}/native-sm86/linear/linear.pte" \ + --data_path "${sm86_ptd}" \ + >"${RUNNER_ARTIFACT_DIR}/native-sm86-on-sm80.log" 2>&1 + native_status=$? + set -e + cat "${RUNNER_ARTIFACT_DIR}/native-sm86-on-sm80.log" + test "${native_status}" -ne 0 + grep -q "no native or PTX variant compatible with sm80" \ + "${RUNNER_ARTIFACT_DIR}/native-sm86-on-sm80.log" + unittest-cuda: name: unittest-cuda needs: [changed-files, run-decision] @@ -370,6 +512,11 @@ jobs: conda install -y -c conda-forge 'libstdcxx-ng>=12' export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH + python -m pytest \ + backends/cuda/tests/test_cuda_weight_metadata.py \ + backends/cuda/tests/test_merge_ptes.py \ + -v -o "addopts=" + cmake --preset llm-release-cuda -DEXECUTORCH_BUILD_TESTS=ON cmake --build cmake-out --target test_cuda_allocator test_cuda_mutable_state test_cuda_weight_cache -j$(nproc) ctest --test-dir cmake-out -R test_cuda_allocator --output-on-failure -V @@ -504,6 +651,25 @@ jobs: repo: "openai" name: "whisper-large-v3-turbo" quant: "non-quantized" + include: + - model: + repo: "google" + name: "gemma-3-4b-it" + quant: "quantized-int4-tile-packed" + pybind_model: "gemma3-4b" + pybind_quantized: true + - model: + repo: "Qwen" + name: "Qwen3-0.6B" + quant: "non-quantized" + pybind_model: "qwen3-0.6b" + pybind_quantized: false + - model: + repo: "Qwen" + name: "Qwen3-0.6B" + quant: "quantized-int4-tile-packed" + pybind_model: "qwen3-0.6b" + pybind_quantized: true with: timeout: 240 secrets-env: EXECUTORCH_HF_TOKEN @@ -512,7 +678,6 @@ jobs: gpu-arch-version: "13.0" use-custom-docker-registry: false submodules: recursive - upload-artifact: ${{ matrix.model.repo }}-${{ matrix.model.name }}-cuda-${{ matrix.quant }} ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} script: | set -eux @@ -547,10 +712,30 @@ jobs: echo "::endgroup::" fi + MODEL_DIR="$(mktemp -d "${RUNNER_TEMP:-/tmp}/cuda_model_XXXXXX")" RUN_EXPORT=1 source .ci/scripts/test_model_e2e.sh cuda \ "${{ matrix.model.repo }}/${{ matrix.model.name }}" \ "${{ matrix.quant }}" \ - "${RUNNER_ARTIFACT_DIR}" + "${MODEL_DIR}" + + if [ -n "${{ matrix.pybind_model }}" ]; then + echo "::group::Run CUDA model with pybind" + conda install -y -c conda-forge 'libstdcxx-ng>=12' + export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH + strings /opt/conda/lib/libstdc++.so.6 | grep GLIBCXX_3.4.29 + + PYBIND_ARGS=() + if [ "${{ matrix.pybind_quantized }}" = "true" ]; then + PYBIND_ARGS+=(--quantize) + fi + python .ci/scripts/test_huggingface_optimum_model.py \ + --model "${{ matrix.pybind_model }}" \ + --recipe cuda \ + --model_dir "${MODEL_DIR}" \ + --run_only \ + "${PYBIND_ARGS[@]}" + echo "::endgroup::" + fi test-muse-glimmer-cuda-e2e: name: test-muse-glimmer-cuda-e2e-${{ matrix.variant }}-${{ matrix.mode }} @@ -621,96 +806,3 @@ jobs: "${{ matrix.variant }}" \ "${RUNNER_TEMP}/muse_glimmer" \ "${{ matrix.mode }}" - - test-cuda-pybind: - name: test-cuda-pybind - # This job downloads models exported by test-model-cuda-e2e and runs them using pybind. - # Explicitly check the producer job so a skipped run (fork PR, - # non-sampled push, or no path match) auto-skips this job too. - needs: [changed-files, test-model-cuda-e2e, run-decision] - if: | - needs.test-model-cuda-e2e.result == 'success' && - ( - contains(needs.changed-files.outputs.changed-files, 'backends/cuda') || - contains(needs.changed-files.outputs.changed-files, 'backends/aoti') || - contains(needs.changed-files.outputs.changed-files, '.github/workflows/cuda.yml') || - contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test-cuda-build.sh') || - contains(needs.changed-files.outputs.changed-files, '.ci/scripts/export_model_artifact.sh') || - contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_model_e2e.sh') || - needs.run-decision.outputs.is-full-run == 'true' - ) - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main - permissions: - id-token: write - contents: read - secrets: inherit - strategy: - fail-fast: false - matrix: - include: - - model: "gemma3-4b" - quantize: "--quantize" - artifact: "google-gemma-3-4b-it-cuda-quantized-int4-tile-packed" - - model: "qwen3-0.6b" - quantize: "" - artifact: "Qwen-Qwen3-0.6B-cuda-non-quantized" - - model: "qwen3-0.6b" - quantize: "--quantize" - artifact: "Qwen-Qwen3-0.6B-cuda-quantized-int4-tile-packed" - with: - timeout: 120 - secrets-env: EXECUTORCH_HF_TOKEN - download-artifact: ${{ matrix.artifact }} - runner: linux.g5.4xlarge.nvidia.gpu - gpu-arch-type: cuda - gpu-arch-version: "13.0" - use-custom-docker-registry: false - submodules: recursive - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} - script: | - set -eux - - # OSDC mounts HF_HOME read-only at /mnt/hf_cache; redirect to a writable dir - # (RUNNER_TEMP, or /tmp when RUNNER_TEMP isn't writable inside the container). - export HF_HOME="${RUNNER_TEMP:-/tmp}/hf_cache" - mkdir -p "${HF_HOME}" 2>/dev/null || export HF_HOME=/tmp/hf_cache - mkdir -p "${HF_HOME}" - - echo "::group::Setup ExecuTorch" - # Disable MKL to avoid duplicate target error when conda has multiple MKL installations - export USE_MKL=OFF - ./install_executorch.sh - echo "::endgroup::" - - echo "::group::Fix libstdc++ GLIBCXX version" - # The embedded .so files in the CUDA blob require GLIBCXX_3.4.29 - # which the default conda libstdc++ doesn't have. Install a newer - # libstdc++ from conda-forge and use it via LD_PRELOAD. - conda install -y -c conda-forge 'libstdcxx-ng>=12' - export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH - # Verify the new libstdc++ has GLIBCXX_3.4.29 - strings /opt/conda/lib/libstdc++.so.6 | grep GLIBCXX_3.4.29 || { - echo "Error: GLIBCXX_3.4.29 not found in /opt/conda/lib/libstdc++.so.6" - exit 1 - } - echo "::endgroup::" - - echo "::group::Setup Huggingface" - pip install -U "huggingface_hub[cli]>=1.2.1,<2.0" - export HF_TOKEN="$(printf '%s' "$SECRET_EXECUTORCH_HF_TOKEN" | tr -d '\r\n')" - echo "::endgroup::" - - echo "::group::Install optimum-executorch" - OPTIMUM_ET_VERSION=$(cat .ci/docker/ci_commit_pins/optimum-executorch.txt) - pip install "optimum~=2.0.0" "transformers==5.0.0rc1" - pip install --no-deps git+https://github.com/huggingface/optimum-executorch.git@${OPTIMUM_ET_VERSION} - echo "::endgroup::" - - echo "::group::Test CUDA Model: ${{ matrix.model }} ${{ matrix.quantize }}" - python .ci/scripts/test_huggingface_optimum_model.py \ - --model ${{ matrix.model }} \ - --recipe cuda \ - --model_dir "${RUNNER_ARTIFACT_DIR}" \ - --run_only \ - ${{ matrix.quantize }} - echo "::endgroup::" diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 3393f2dbfc6..f8a2008ff9e 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -11,7 +11,7 @@ on: workflow_dispatch: concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true jobs: @@ -41,16 +41,24 @@ jobs: with: python-version: '3.11' cache: 'pip' - cache-dependency-path: requirements-lintrunner.txt + cache-dependency-path: | + .github/workflows/lint.yml + requirements-lintrunner.txt + torch_pin.py - name: Install dependencies run: | - pip install torch --index-url https://download.pytorch.org/whl/cpu + TORCH_VERSION=$(python -c "from torch_pin import TORCH_VERSION; print(TORCH_VERSION)") + pip install \ + "torch==${TORCH_VERSION}" \ + torchvision \ + torchaudio \ + --index-url https://download.pytorch.org/whl/cpu pip install lintrunner==0.12.7 lintrunner-adapters==0.14.1 pip install -r requirements-lintrunner.txt USE_CPP=0 pip install --no-build-isolation third-party/ao - pip install pytest numpy parameterized huggingface_hub transformers timm expecttest types-requests - pip install torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu + # Match requirements-examples.txt for the model APIs checked by mypy. + pip install pytest numpy parameterized huggingface_hub "transformers==5.0.0rc1" timm expecttest types-requests - name: Generate mypy stubs for C++ bindings run: | diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index d25a2b89c43..734af8b329f 100644 --- a/.github/workflows/mlx.yml +++ b/.github/workflows/mlx.yml @@ -68,7 +68,42 @@ jobs: echo "::endgroup::" echo "::group::Build test runners" - ${CONDA_RUN} cmake --build cmake-out --target op_test_runner multi_thread_test_runner mlx_mutable_state_test mlx_sequence_cache_test mlx_cell_cache_test -j$(( $(sysctl -n hw.ncpu) - 1 )) + ${CONDA_RUN} cmake --build cmake-out --target op_test_runner multi_thread_test_runner mlx_metallib_path_test mlx_mutable_state_test mlx_sequence_cache_test mlx_cell_cache_test -j$(( $(sysctl -n hw.ncpu) - 1 )) + echo "::endgroup::" + + echo "::group::Check MLX artifact sizes" + size_check_failed=0 + check_artifact_size() { + local path="$1" + local max_size="$2" + local baseline_size="$3" + + if [[ ! -f "${path}" ]]; then + echo "::error file=${path}::Expected MLX artifact not found" + size_check_failed=1 + return + fi + + local actual_size + actual_size=$(wc -c < "${path}" | tr -d '[:space:]') + echo "${path}: ${actual_size} bytes (baseline ${baseline_size}, limit ${max_size})" + if (( actual_size > max_size )); then + echo "::error file=${path}::MLX artifact is ${actual_size} bytes, exceeding the ${max_size}-byte limit (baseline ${baseline_size})" + size_check_failed=1 + fi + } + + # Baselines recorded from the MLX v0.32.2 Release build on 2026-09-02. + check_artifact_size cmake-out/backends/mlx/libmlxdelegate.a 5000000 4159552 + check_artifact_size cmake-out/backends/mlx/mlx/libmlx.a 14000000 11458544 + check_artifact_size cmake-out/backends/mlx/mlx/mlx/backend/metal/kernels/mlx.metallib 2000000 1460840 + echo "::endgroup::" + if (( size_check_failed )); then + exit 1 + fi + + echo "::group::Run SwiftPM metallib path unit test" + ./cmake-out/backends/mlx/test/mlx_metallib_path_test echo "::endgroup::" echo "::group::Run mutable-state (multi-session) unit test" @@ -818,7 +853,7 @@ jobs: # DFlash speculative decoding: target and draft exported as two methods of one # .pte. Also the only coverage that the constants the export publishes - # (get_max_ctx_len / get_prefill_chunk_size / get_max_block_len / + # (get_max_context_len / get_max_seq_len / get_max_block_len / # get_mask_token_id) match what # the runner derives from them -- the runner takes no capacity or block flags. test-mlx-dflash: @@ -884,7 +919,7 @@ jobs: BLOCK_SIZE=$(${CONDA_RUN} python -c "from huggingface_hub import snapshot_download; from executorch.backends.mlx.examples.llm.dflash.model import load_dflash_config; print(load_dflash_config(snapshot_download('${DRAFT_MODEL}', allow_patterns=['*.json'])).block_size)") MASK_TOKEN_ID=$(${CONDA_RUN} python -c "from huggingface_hub import snapshot_download; from executorch.backends.mlx.examples.llm.dflash.model import load_dflash_config; print(load_dflash_config(snapshot_download('${DRAFT_MODEL}', allow_patterns=['*.json'])).mask_token_id)") echo "draft checkpoint block_size: ${BLOCK_SIZE} mask_token_id: ${MASK_TOKEN_ID}" - ${CONDA_RUN} python -c "from executorch.runtime import Runtime, Verification; from executorch.backends.mlx.examples.llm.runtime_meta import read_const_int as r; p = Runtime.get().load_program('${PTE}', verification=Verification.Minimal); got = {n: r(p, n) for n in ['get_max_ctx_len', 'get_prefill_chunk_size', 'get_max_block_len', 'get_mask_token_id']}; want = {'get_max_ctx_len': ${MAX_CTX_LEN}, 'get_prefill_chunk_size': ${PREFILL_CHUNK_SIZE}, 'get_max_block_len': ${BLOCK_SIZE}, 'get_mask_token_id': ${MASK_TOKEN_ID}}; print('published:', got); assert got == want, f'mismatch: {got} != {want}'; assert {'draft', 'target'} <= set(p.method_names), p.method_names; print('Success: constants match and both methods present')" + ${CONDA_RUN} python -c "from executorch.runtime import Runtime, Verification; from executorch.backends.mlx.examples.llm.runtime_meta import read_const_int as r; p = Runtime.get().load_program('${PTE}', verification=Verification.Minimal); got = {n: r(p, n) for n in ['get_max_context_len', 'get_max_seq_len', 'get_max_block_len', 'get_mask_token_id']}; want = {'get_max_context_len': ${MAX_CTX_LEN}, 'get_max_seq_len': ${PREFILL_CHUNK_SIZE}, 'get_max_block_len': ${BLOCK_SIZE}, 'get_mask_token_id': ${MASK_TOKEN_ID}}; print('published:', got); assert got == want, f'mismatch: {got} != {want}'; assert {'draft', 'target'} <= set(p.method_names), p.method_names; print('Success: constants match and both methods present')" echo "::endgroup::" echo "::group::Run DFlash speculative decoding" @@ -1002,10 +1037,13 @@ jobs: ${CONDA_RUN} cmake --build cmake-out/backends/mlx/examples/llm \ -j$(( $(sysctl -n hw.ncpu) - 1 )) RUNNER=cmake-out/backends/mlx/examples/llm/mlx_run_llm_hf - if [ ! -x "${RUNNER}" ]; then - echo "Failed: runner not found at ${RUNNER}" - exit 1 - fi + BATCHED_RUNNER=cmake-out/backends/mlx/examples/llm/mlx_run_llm_batched + for binary in "${RUNNER}" "${BATCHED_RUNNER}"; do + if [ ! -x "${binary}" ]; then + echo "Failed: runner not found at ${binary}" + exit 1 + fi + done echo "::endgroup::" echo "::group::Install LLM requirements" @@ -1026,6 +1064,7 @@ jobs: --model-id "${MODEL_ID}" \ --output /tmp/${MODEL_NAME}_offgraph.pte \ --use-offgraph-cache \ + --logits-to-keep selected \ --max-ctx-len 1024 \ --dtype bf16 \ --qlinear 4w @@ -1050,3 +1089,28 @@ jobs: exit 1 fi echo "::endgroup::" + + echo "::group::Run ${MODEL_NAME} batched off-graph inference" + ${BATCHED_RUNNER} \ + --pte /tmp/${MODEL_NAME}_offgraph.pte \ + --tokenizer "${TOKENIZER}" \ + --chat "${CHAT}" \ + --max-session-tokens 1024 \ + --max-new-tokens 50 \ + --max-decode-sequences 2 \ + --out-prefix /tmp/${MODEL_NAME}_batched \ + "What is the capital of France?" \ + "What color is grass?" + for check in "0:Paris" "1:green"; do + index="${check%%:*}" + expected="${check#*:}" + output="/tmp/${MODEL_NAME}_batched_${index}.txt" + if grep -iq "${expected}" "${output}"; then + echo "Success: '${expected}' found in ${output}" + else + echo "Failed: Expected '${expected}' not found in ${output}" + cat "${output}" + exit 1 + fi + done + echo "::endgroup::" diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 922ed95ab8d..c301b61399e 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -16,6 +16,10 @@ concurrency: cancel-in-progress: true jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + update-pytorch-commit-hash: runs-on: ubuntu-latest environment: ${{ (github.event_name == 'schedule') && 'update-commit-hash' || '' }} @@ -50,8 +54,9 @@ jobs: timeout: 180 test-static-hf-llm-qnn-linux: + needs: docker-image name: test-static-hf-llm-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -60,8 +65,8 @@ jobs: task: [smollm2_135m] fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 diff --git a/.github/workflows/periodic.yml b/.github/workflows/periodic.yml index 01bff087124..a859a9fd779 100644 --- a/.github/workflows/periodic.yml +++ b/.github/workflows/periodic.yml @@ -21,6 +21,10 @@ concurrency: permissions: read-all jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + gather-models: runs-on: ubuntu-22.04 outputs: @@ -42,17 +46,17 @@ jobs: test-models-linux: name: test-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read - needs: gather-models + needs: [docker-image, gather-models] strategy: matrix: ${{ fromJSON(needs.gather-models.outputs.models) }} fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: ${{ matrix.timeout }} diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index a190dc132d1..d17be95b3aa 100644 --- a/.github/workflows/pull.yml +++ b/.github/workflows/pull.yml @@ -6,15 +6,17 @@ on: branches: - main - release/* - tags: - - ciflow/trunk/* workflow_dispatch: concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + # Emits the list of changed files for the current PR or push commit. # On PR: PR diff. On push: diff against `github.event.before`. # On events without a diff base (workflow_dispatch, tag creation, @@ -35,8 +37,9 @@ jobs: uses: ./.github/workflows/_ci-run-decision.yml test-qnn-wheel-packages-linux: + needs: docker-image name: test-qnn-wheel-packages-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -45,8 +48,8 @@ jobs: matrix: python-version: [ "3.10", "3.11", "3.12", "3.13" ] with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -79,7 +82,7 @@ jobs: contents: read test-minimal-wheel-linux: - needs: changed-files + needs: [docker-image, changed-files] if: | github.event_name != 'pull_request' || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_minimal_wheel.sh') || @@ -92,13 +95,13 @@ jobs: contains(needs.changed-files.outputs.changed-files, 'setup.py') || contains(needs.changed-files.outputs.changed-files, 'tools/cmake/') name: test-minimal-wheel-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -109,16 +112,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_minimal_wheel.sh test-setup-linux-gcc: + needs: docker-image name: test-setup-linux-gcc - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-gcc11 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc11-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -134,8 +138,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh "add" "${BUILD_TOOL}" "portable" test-models-linux-basic: + needs: docker-image name: test-models-linux-basic - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -144,23 +149,23 @@ jobs: model: [mv3, vit] backend: [portable, xnnpack-quantization-delegation] build-tool: [cmake, buck2] - runner: [linux.2xlarge, linux.arm64.2xlarge] + runner: [mt-l-x86iavx512-8-64, mt-l-arm64g4-16-62] docker-image: [executorch-ubuntu-22.04-clang12, executorch-ubuntu-22.04-gcc11-aarch64] # Excluding specific runner + docker image combinations that don't make sense: - # - Excluding the ARM64 gcc image on the x86 runner (linux.2xlarge) - # - Excluding the x86 clang image on the ARM64 runner (linux.arm64.2xlarge) + # - Excluding the ARM64 gcc image on the x86 runner + # - Excluding the x86 clang image on the ARM64 runner exclude: - - runner: linux.2xlarge + - runner: mt-l-x86iavx512-8-64 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 - - runner: linux.arm64.2xlarge + - runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-clang12 # TODO: Need to figure out why buck2 doesnt work on Graviton instances. - - runner: linux.arm64.2xlarge + - runner: mt-l-arm64g4-16-62 build-tool: buck2 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:${{ matrix.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:${{ matrix.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -178,8 +183,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" test-models-linux: + needs: docker-image name: test-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -187,33 +193,33 @@ jobs: matrix: model: [linear, add, add_mul, ic3, mv2, resnet18, resnet50, mobilebert, emformer_transcribe] backend: [portable, xnnpack-quantization-delegation] - runner: [linux.2xlarge] + runner: [mt-l-x86iavx512-8-64] include: - model: ic4 backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: ic4 backend: xnnpack-quantization-delegation - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: emformer_join backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: emformer_join backend: xnnpack-quantization-delegation - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: phi_4_mini backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: llama3_2_vision_encoder backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: w2l backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -231,16 +237,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" test-parakeet-xnnpack-linux: + needs: docker-image name: test-parakeet-xnnpack-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.4xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-16-128 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -264,16 +271,17 @@ jobs: echo "::endgroup::" test-voxtral-realtime-xnnpack-linux: + needs: docker-image name: test-voxtral-realtime-xnnpack-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.4xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-16-128 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -300,9 +308,10 @@ jobs: echo "::endgroup::" test-llama-runner-linux: + needs: docker-image # Test Both linux x86 and linux aarch64 name: test-llama-runner-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -310,25 +319,25 @@ jobs: matrix: dtype: [fp32] mode: [xnnpack+custom+qe,xnnpack+custom+quantize_kv,xnnpack+quantize_kv] - runner: [linux.2xlarge, linux.arm64.2xlarge] + runner: [mt-l-x86iavx512-8-64, mt-l-arm64g4-16-62] docker-image: [executorch-ubuntu-22.04-clang12, executorch-ubuntu-22.04-gcc11-aarch64] include: - dtype: bf16 mode: custom - runner: linux.2xlarge + runner: mt-l-x86iavx512-8-64 docker-image: executorch-ubuntu-22.04-clang12 # Excluding specific runner + docker image combinations that don't make sense: - # - Excluding the ARM64 gcc image on the x86 runner (linux.2xlarge) - # - Excluding the x86 clang image on the ARM64 runner (linux.arm64.2xlarge) + # - Excluding the ARM64 gcc image on the x86 runner + # - Excluding the x86 clang image on the ARM64 runner exclude: - - runner: linux.2xlarge + - runner: mt-l-x86iavx512-8-64 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 - - runner: linux.arm64.2xlarge + - runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-clang12 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:${{ matrix.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:${{ matrix.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -351,16 +360,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh -model stories110M -build_tool "${BUILD_TOOL}" -dtype "${DTYPE}" -mode "${MODE}" -upload "${ARTIFACTS_DIR_NAME}" test-llama-runner-linux-android: + needs: docker-image name: test-llama-runner-linux-android - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12-android + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-android-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -376,16 +386,17 @@ jobs: bash .ci/scripts/build_llama_android.sh "${BUILD_TOOL}" test-custom-ops-linux: + needs: docker-image name: test-custom-ops-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -400,16 +411,17 @@ jobs: PYTHON_EXECUTABLE=python bash examples/portable/custom_ops/test_custom_ops.sh "${BUILD_TOOL}" test-selective-build-linux: + needs: docker-image name: test-selective-build-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -424,9 +436,10 @@ jobs: PYTHON_EXECUTABLE=python bash examples/selective_build/test_selective_build.sh "${BUILD_TOOL}" test-multimodal-linux: + needs: docker-image if: ${{ !github.event.pull_request.head.repo.fork }} name: test-multimodal-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -437,8 +450,8 @@ jobs: model: ["gemma3-4b"] # llava gives segfault so not covering. with: secrets-env: EXECUTORCH_HF_TOKEN - runner: linux.24xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-768 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -463,16 +476,17 @@ jobs: echo "::endgroup::" test-moshi-linux: + needs: docker-image name: test-moshi-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -495,16 +509,17 @@ jobs: python -m unittest examples.models.moshi.mimi.test_mimi test-quantized-aot-lib-linux: + needs: docker-image name: test-quantized-aot-lib-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -518,16 +533,17 @@ jobs: PYTHON_EXECUTABLE=python bash examples/xnnpack/quantization/test_quantize.sh "${BUILD_TOOL}" mv2 test-binary-size-linux-gcc: + needs: docker-image name: test-binary-size-linux-gcc - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-gcc9-nopytorch + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc9-nopytorch-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -561,16 +577,17 @@ jobs: fi test-binary-size-linux: + needs: docker-image name: test-binary-size-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -605,8 +622,9 @@ jobs: fi test-arm-cortex-m-size-test: + needs: docker-image name: test-arm-cortex-m-size-test - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -615,8 +633,8 @@ jobs: os: [bare_metal, zephyr-preset] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -698,14 +716,15 @@ jobs: fi test-mcu-cortex-m-backend: + needs: docker-image name: test-mcu-cortex-m-backend - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -765,16 +784,17 @@ jobs: docker-image: ci-image:executorch-ubuntu-22.04-clang12 test-qnn-buck-build-linux: + needs: docker-image name: test-qnn-buck-build-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -799,8 +819,9 @@ jobs: buck2 build //backends/qualcomm/... test-arm-backend-no-driver: + needs: docker-image name: test-arm-backend-no-driver - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -813,8 +834,8 @@ jobs: - test_arm_backend: test_run_tosa fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -834,14 +855,15 @@ jobs: backends/arm/test/test_arm_backend.sh "${ARM_TEST}" test-arm-backend-public-api-backward-compatibility: + needs: docker-image name: test-arm-backend-public-api-backward-compatibility - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-24.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-24.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -861,8 +883,9 @@ jobs: python backends/arm/test/public_api_bc/run_public_api_bc_scenarios.py test-llama-runner-qnn-linux: + needs: docker-image name: test-llama-runner-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -873,8 +896,8 @@ jobs: mode: [qnn] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -900,8 +923,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh -model stories110M -build_tool "${BUILD_TOOL}" -mode "${MODE}" -dtype "${DTYPE}" -pt2e_quantize "${PT2E_QUANTIZE}" test-static-llama-qnn-linux: + needs: docker-image name: test-static-llama-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -910,8 +934,8 @@ jobs: task: [stories_110m, stories_260k_bc] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -934,8 +958,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_qnn_static_llm.sh ${{ matrix.task }} test-sqnr-static-llm-qnn-linux: + needs: docker-image name: test-sqnr-static-llm-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -944,8 +969,8 @@ jobs: task: [smollm2_135m] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -968,8 +993,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_qnn_static_llm.sh ${{ matrix.task }} sqnr test-qnn-models-linux: + needs: docker-image name: test-qnn-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -978,8 +1004,8 @@ jobs: model: [mv2, mv3, dl3] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -993,14 +1019,15 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh ${{ matrix.model }} "cmake" "qnn" test-qnn-direct-build-linux: + needs: docker-image name: test-qnn-direct-build-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 30 @@ -1023,22 +1050,23 @@ jobs: ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 run-linux: true - # No runner-linux, so this takes the memory-optimized default. The suite - # runs one export worker per core, so what matters is memory per core, not - # core count: the previous instance gave each worker about 4 GiB and the - # job was killed. The memory-optimized default gives each worker more. + # No runner-linux, so this takes _test_backend.yml's default. The suite + # runs one export worker per core, so what matters is memory per core: + # the instance this used to run on gave each worker about 4 GiB and the + # job was killed. The default label is a little under 8 GiB per core. test-qnn-passes-linux: + needs: docker-image name: test-qnn-passes-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 30 @@ -1067,16 +1095,21 @@ jobs: pytest -xvs backends/qualcomm/tests/test_import_side_effects.py test-qnn-delegate-linux: + needs: docker-image name: test-qnn-delegate-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + # Intel, unlike the avx512 labels: those are r7a, i.e. AMD EPYC, and the + # QNN backend disables MKLDNN on an AMD host through a non-bracketed + # torch.backends mutation that raises once the tests have frozen the + # flags. Same 8 vCPU / 64Gi, on r7i. + runner: mt-l-x86iamx-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1105,16 +1138,17 @@ jobs: -k "TestQNNFloatingPointOperator or TestQNNQuantizedOperator" test-phi-3-mini-runner-linux: + needs: docker-image name: test-phi-3-mini-runner-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1135,16 +1169,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_phi_3_mini.sh Release test-qnn-python-imports-linux: + needs: docker-image name: test-qnn-python-imports-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 15 @@ -1183,16 +1218,17 @@ jobs: --module-prefix executorch.examples.qualcomm test-eval_llama-wikitext-linux: + needs: docker-image name: test-eval_llama-wikitext-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1212,7 +1248,7 @@ jobs: # TODO(larryliu0820): Fix this issue before reenabling it: https://gist.github.com/larryliu0820/7377ecd0d79dbc06076cec8d9f2b85d2 # test-eval_llama-mmlu-linux: # name: test-eval_llama-mmlu-linux - # uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + # uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main # permissions: # id-token: write # contents: read @@ -1238,16 +1274,17 @@ jobs: # PYTHON_EXECUTABLE=python bash .ci/scripts/test_eval_llama_mmlu.sh test-llama_runner_eager-linux: + needs: docker-image name: test-llama_runner_eager-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1265,16 +1302,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama_runner_eager.sh test-lora-linux: + needs: docker-image name: test-lora-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1292,16 +1330,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_lora.sh test-lora-multimethod-linux: + needs: docker-image name: test-lora-multimethod-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1319,16 +1358,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_lora_multimethod.sh test-mediatek-models-linux: + needs: docker-image name: test-mediatek-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-mediatek-sdk + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-mediatek-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1346,16 +1386,17 @@ jobs: # placeholder for mediatek to add more tests test-openvino-linux: + needs: docker-image name: test-openvino-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-gcc11 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc11-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1368,16 +1409,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_openvino.sh test-build-wasm-linux: + needs: docker-image name: test-build-wasm-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1396,8 +1438,9 @@ jobs: PYTHON_EXECUTABLE=python bash examples/wasm/test_build_wasm.sh unittest-wasm-bindings: + needs: docker-image name: unittest-wasm-bindings - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -1406,8 +1449,8 @@ jobs: enable-etdump: ['', '--enable-etdump'] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1442,13 +1485,14 @@ jobs: pnpm test unittest-nxp-neutron: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 150 @@ -1486,18 +1530,19 @@ jobs: bash backends/nxp/run_unittests.sh test-samsung-quantmodels-linux: + needs: docker-image name: test-samsung-quantmodels-linux # Skip this job if the pull request is from a fork (secrets are not available) if: github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request' - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read secrets: inherit with: secrets-env: SAMSUNG_AI_LITECORE_KEY - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12-android + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-android-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -1524,18 +1569,19 @@ jobs: done test-samsung-models-linux: + needs: docker-image name: test-samsung-models-linux # Skip this job if the pull request is from a fork (secrets are not available) if: github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request' - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read secrets: inherit with: secrets-env: SAMSUNG_AI_LITECORE_KEY - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12-android + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-android-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 360 @@ -1566,14 +1612,15 @@ jobs: python -m unittest discover -s backends/samsung/test/models -p "test_*.py" test-vulkan-models-linux: + needs: docker-image name: test-vulkan-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1607,14 +1654,15 @@ jobs: done test-vulkan-operators-linux: + needs: docker-image name: test-vulkan-operators-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1699,14 +1747,15 @@ jobs: echo "::endgroup::" nxp-build-test: + needs: docker-image name: nxp-build-test - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1727,3 +1776,22 @@ jobs: echo "Neutron backend library not found!" exit 1 fi + + nxp-mcuxpresso-test: + name: nxp-mcuxpresso-test + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + permissions: + id-token: write + contents: read + with: + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + submodules: 'recursive' + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 90 + script: | + # The generic Linux job chooses to use base env, not the one setup by the image + CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]") + conda activate "${CONDA_ENV}" + + ./examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/test_build_from_scratch.sh diff --git a/.github/workflows/qnn-windows-msvc.yml b/.github/workflows/qnn-windows-msvc.yml index 377a2625e60..65dafc3636a 100644 --- a/.github/workflows/qnn-windows-msvc.yml +++ b/.github/workflows/qnn-windows-msvc.yml @@ -56,8 +56,8 @@ permissions: contents: read jobs: - build-qnn-windows-msvc: - name: build-qnn-windows-msvc + build-qnn-windows-x64: + name: build-qnn-windows-x64 uses: pytorch/test-infra/.github/workflows/windows_job.yml@main with: ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} @@ -70,5 +70,35 @@ jobs: Set-PSDebug -Trace 1 \$ErrorActionPreference = 'Stop' \$PSNativeCommandUseErrorActionPreference = \$true - .ci/scripts/build-qnn-windows-msvc.ps1 + .ci/scripts/build-qnn-windows-msvc.ps1 -SkipArm64Windows }" + + build-qnn-windows-arm64: + name: build-qnn-windows-arm64 + runs-on: windows-11-arm + timeout-minutes: 90 + steps: + - name: Enable long paths + shell: cmd + run: | + git config --system --get core.longpaths || echo "core.longpaths is not set, setting it now" + git config --system core.longpaths true + + - name: Checkout ExecuTorch + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + + - name: Initialize submodules + shell: pwsh + run: | + git config --global http.sslBackend openssl + git submodule update --init --recursive + + - name: Build QNN backend (arm64) + shell: pwsh + run: | + Set-PSDebug -Trace 1 + $ErrorActionPreference = 'Stop' + $PSNativeCommandUseErrorActionPreference = $true + .ci/scripts/build-qnn-windows-msvc.ps1 -SkipX86Windows diff --git a/.github/workflows/riscv64.yml b/.github/workflows/riscv64.yml index 244db631021..34f4ab04874 100644 --- a/.github/workflows/riscv64.yml +++ b/.github/workflows/riscv64.yml @@ -10,6 +10,7 @@ on: pull_request: paths: - .github/workflows/riscv64.yml + - .github/workflows/_test_riscv.yml - .ci/scripts/test_riscv_qemu.sh - tools/cmake/preset/riscv64_linux.cmake - examples/riscv/** diff --git a/.github/workflows/rocm.yml b/.github/workflows/rocm.yml index 97d180885f7..154454eb145 100644 --- a/.github/workflows/rocm.yml +++ b/.github/workflows/rocm.yml @@ -169,7 +169,7 @@ jobs: strategy: fail-fast: false matrix: - rocm-version: ["7.1"] + rocm-version: ["7.2"] uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write @@ -206,7 +206,7 @@ jobs: strategy: fail-fast: false matrix: - rocm-version: ["7.1"] + rocm-version: ["7.2"] with: timeout: 180 no-sudo: true @@ -248,7 +248,7 @@ jobs: strategy: fail-fast: false matrix: - rocm-version: ["7.1"] + rocm-version: ["7.2"] with: timeout: 180 no-sudo: true diff --git a/.github/workflows/test-backend-arm.yml b/.github/workflows/test-backend-arm.yml index d71696ee096..7d844748dac 100644 --- a/.github/workflows/test-backend-arm.yml +++ b/.github/workflows/test-backend-arm.yml @@ -19,6 +19,9 @@ concurrency: jobs: test-arm: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: arm flows: >- @@ -35,6 +38,9 @@ jobs: test-arm-vgf: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: arm-vgf flows: >- diff --git a/.github/workflows/test-backend-coreml.yml b/.github/workflows/test-backend-coreml.yml index 86844ffb559..fca10b619ce 100644 --- a/.github/workflows/test-backend-coreml.yml +++ b/.github/workflows/test-backend-coreml.yml @@ -36,6 +36,9 @@ jobs: contains(needs.changed-files.outputs.changed-files, '.github/workflows/test-backend-coreml.yml') || contains(needs.changed-files.outputs.changed-files, '.github/workflows/_test_backend.yml') uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: coreml # The heavier coreml_static_int8 macOS matrix saturates the shared diff --git a/.github/workflows/test-backend-cortex-m.yml b/.github/workflows/test-backend-cortex-m.yml index 9e9b4cdcaa1..b6cd09d3073 100644 --- a/.github/workflows/test-backend-cortex-m.yml +++ b/.github/workflows/test-backend-cortex-m.yml @@ -49,6 +49,9 @@ jobs: contains(needs.changed-files.outputs.changed-files, '.github/workflows/test-backend-cortex-m.yml') || contains(needs.changed-files.outputs.changed-files, '.github/workflows/_test_backend.yml') uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: cortex_m flows: '["cortex_m"]' diff --git a/.github/workflows/test-backend-nxp.yml b/.github/workflows/test-backend-nxp.yml index fed7ab7d19b..0db7d7efeac 100644 --- a/.github/workflows/test-backend-nxp.yml +++ b/.github/workflows/test-backend-nxp.yml @@ -39,6 +39,9 @@ jobs: contains(needs.changed-files.outputs.changed-files, '.github/workflows/test-backend-nxp.yml') || contains(needs.changed-files.outputs.changed-files, '.github/workflows/_test_backend.yml') uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: nxp flows: '["nxp_neutron_imxrt700_int8_ptq"]' diff --git a/.github/workflows/test-backend-openvino.yml b/.github/workflows/test-backend-openvino.yml index aeeb01e3eb1..cff7f8ba6e9 100644 --- a/.github/workflows/test-backend-openvino.yml +++ b/.github/workflows/test-backend-openvino.yml @@ -21,6 +21,9 @@ concurrency: jobs: test-openvino: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: openvino flows: '["openvino"]' diff --git a/.github/workflows/test-backend-qnn.yml b/.github/workflows/test-backend-qnn.yml index 939b7c36aee..51207d139b3 100644 --- a/.github/workflows/test-backend-qnn.yml +++ b/.github/workflows/test-backend-qnn.yml @@ -19,6 +19,9 @@ concurrency: jobs: test-qnn: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: qnn flows: >- @@ -28,4 +31,4 @@ jobs: ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 run-linux: true - runner-linux: linux.8xlarge.memory + runner-linux: mt-l-x86iavx512-32-256 diff --git a/.github/workflows/test-backend-vulkan.yml b/.github/workflows/test-backend-vulkan.yml index 80ac3ee73a1..d3c9829dfe8 100644 --- a/.github/workflows/test-backend-vulkan.yml +++ b/.github/workflows/test-backend-vulkan.yml @@ -21,6 +21,9 @@ jobs: # runners. Runs on every PR and nightly. test-vulkan: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: vulkan flows: >- diff --git a/.github/workflows/test-backend-webgpu.yml b/.github/workflows/test-backend-webgpu.yml index 99dae3ee76d..42e958855ac 100644 --- a/.github/workflows/test-backend-webgpu.yml +++ b/.github/workflows/test-backend-webgpu.yml @@ -19,6 +19,9 @@ concurrency: jobs: test-webgpu: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: webgpu flows: '["webgpu"]' diff --git a/.github/workflows/test-backend-xnnpack.yml b/.github/workflows/test-backend-xnnpack.yml index e9f43608b16..a0906af5bb2 100644 --- a/.github/workflows/test-backend-xnnpack.yml +++ b/.github/workflows/test-backend-xnnpack.yml @@ -11,6 +11,11 @@ on: - ciflow/nightly/* pull_request: workflow_dispatch: + inputs: + save-goldens: + description: 'Write and package the models-suite goldens (Android test fixture)' + type: boolean + default: false concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} @@ -19,6 +24,9 @@ concurrency: jobs: test-xnnpack: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: xnnpack flows: >- @@ -28,3 +36,7 @@ jobs: ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 run-linux: true + # The nightly keeps writing goldens: they are what the Android instrumentation + # test pins (android_test_setup.sh), and a fresh key must exist before the pin + # ages out of S3. Every other run skips them. + save-goldens: ${{ inputs.save-goldens == true || github.event_name == 'schedule' }} diff --git a/.github/workflows/trunk.yml b/.github/workflows/trunk.yml index 25b2df3aa4f..f75e7e76f01 100644 --- a/.github/workflows/trunk.yml +++ b/.github/workflows/trunk.yml @@ -19,6 +19,10 @@ concurrency: cancel-in-progress: true jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + # Emits the list of changed files for the current PR or push commit. # On PR: PR diff. On push: diff against `github.event.before`. # On events without a diff base (workflow_dispatch, tag creation, @@ -81,8 +85,12 @@ jobs: PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" test-arm-backend-zephyr: + needs: docker-image name: test-arm-backend-zephyr - uses: pytorch/test-infra/.github/workflows/linux_job.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read strategy: matrix: include: @@ -93,8 +101,8 @@ jobs: - { readme: zephyr/samples/mv2-ethosu/README.md, target: ethos-u85 } fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-zephyr-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-zephyr-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -119,8 +127,9 @@ jobs: --zephyr-samples-readme-path "${{ matrix.readme }}" test-models-linux-aarch64: + needs: docker-image name: test-models-linux-aarch64 - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -128,30 +137,30 @@ jobs: matrix: model: [linear, add, add_mul, ic3, ic4, mv2, mv3, resnet18, resnet50, vit, w2l, mobilebert, emformer_join, emformer_transcribe] backend: [portable, xnnpack-quantization-delegation] - runner: [linux.arm64.2xlarge] + runner: [mt-l-arm64g4-16-62] include: - model: lstm backend: portable - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 - model: mul backend: portable - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 - model: softmax backend: portable - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 - model: phi_4_mini backend: portable - runner: linux.arm64.m7g.4xlarge + runner: mt-l-arm64g4-16-62 - model: qwen2_5_1_5b backend: portable - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 - model: llama3_2_vision_encoder backend: portable - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:executorch-ubuntu-22.04-gcc11-aarch64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc11-aarch64-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -223,8 +232,9 @@ jobs: PYTHON_EXECUTABLE=python ${CONDA_RUN} bash examples/selective_build/test_selective_build.sh "${BUILD_TOOL}" test-demo-backend-delegation: + needs: docker-image name: test-demo-backend-delegation - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -235,8 +245,8 @@ jobs: - build-tool: cmake fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} script: | @@ -250,8 +260,9 @@ jobs: PYTHON_EXECUTABLE=python bash examples/portable/scripts/test_demo_backend_delegation.sh "${BUILD_TOOL}" test-arm-backend-ethos-u: + needs: docker-image name: test-arm-backend-ethos-u - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -273,8 +284,8 @@ jobs: - test_arm_backend: test_deit_e2e_ethos_u fail-fast: false with: - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -290,7 +301,8 @@ jobs: # Increase number of files user can monitor to bypass buck failures. # Hopefully this is high enough for this setup. - sudo sysctl fs.inotify.max_user_watches=1048576 # 1024 * 1024 + # Node-level, so an unprivileged pod cannot change it. + sudo sysctl fs.inotify.max_user_watches=1048576 2>/dev/null || true ARM_TEST=${{ matrix.test_arm_backend }} @@ -303,8 +315,9 @@ jobs: backends/arm/test/test_arm_backend.sh "${ARM_TEST}" test-arm-backend-vkml: + needs: docker-image name: test-arm-backend-vkml - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -317,8 +330,8 @@ jobs: - test_arm_backend: test_smaller_stories_llama_vkml fail-fast: false with: - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-24.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-24.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -333,7 +346,8 @@ jobs: # Increase number of files user can monitor to bypass buck failures. # Hopefully this is high enough for this setup. - sudo sysctl fs.inotify.max_user_watches=1048576 # 1024 * 1024 + # Node-level, so an unprivileged pod cannot change it. + sudo sysctl fs.inotify.max_user_watches=1048576 2>/dev/null || true ARM_TEST=${{ matrix.test_arm_backend }} @@ -432,9 +446,10 @@ jobs: ${CONDA_RUN} sh .ci/scripts/test_llama_torchao_lowbit.sh test-llama-runner-linux: + needs: docker-image # Test Both linux x86 and linux aarch64 name: test-llama-runner-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -442,33 +457,33 @@ jobs: matrix: dtype: [fp32] mode: [portable, xnnpack+custom] - runner: [linux.2xlarge, linux.arm64.2xlarge] + runner: [mt-l-x86iavx512-8-64, mt-l-arm64g4-16-62] docker-image: [executorch-ubuntu-22.04-clang12, executorch-ubuntu-22.04-gcc11-aarch64] include: - dtype: bf16 mode: portable - runner: linux.2xlarge + runner: mt-l-x86iavx512-8-64 docker-image: executorch-ubuntu-22.04-clang12 - dtype: bf16 mode: portable - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 - dtype: bf16 mode: custom - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 # Excluding specific runner + docker image combinations that don't make sense: - # - Excluding the ARM64 gcc image on the x86 runner (linux.2xlarge) - # - Excluding the x86 clang image on the ARM64 runner (linux.arm64.2xlarge) + # - Excluding the ARM64 gcc image on the x86 runner + # - Excluding the x86 clang image on the ARM64 runner exclude: - - runner: linux.2xlarge + - runner: mt-l-x86iavx512-8-64 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 - - runner: linux.arm64.2xlarge + - runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-clang12 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:${{ matrix.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:${{ matrix.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -543,34 +558,35 @@ jobs: PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_llama.sh -model stories110M -build_tool cmake -dtype "${DTYPE}" -mode "${MODE}" test-torchao-huggingface-checkpoints: + needs: docker-image name: test-torchao-huggingface-checkpoints - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: matrix: model: [qwen3_4b, phi_4_mini, lfm2_5_1_2b] - runner: [linux.2xlarge] + runner: [mt-l-x86iavx512-8-64] docker-image: [executorch-ubuntu-22.04-clang12] backend: [xnnpack] include: - model: qwen3_4b - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 backend: torchao - model: phi_4_mini - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 backend: torchao - model: lfm2_5_1_2b - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 backend: torchao fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:${{ matrix.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:${{ matrix.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -637,8 +653,9 @@ jobs: echo "::endgroup::" test-qnn-model: + needs: docker-image name: test-qnn-model - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -648,8 +665,8 @@ jobs: model: [dl3, mv3, mv2, ic4, ic3, vit, mb, w2l, conv_former] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -663,8 +680,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh ${{ matrix.model }} "cmake" "qnn" test-qnn-optimum-model: + needs: docker-image name: test-qnn-optimum-model - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -674,8 +692,8 @@ jobs: model: [cvt, dit, efficientnet, focalnet, mobilevit_v1, mobilevit_v2, pvt, swin, albert, bert, distilbert, roberta] # eurobert requires transfomer >= 4.48.0, skip for now fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -740,10 +758,11 @@ jobs: PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" test-huggingface-transformers-xnnpack: + needs: docker-image # NB: Don't run this on fork PRs because they won't have access to the secret and would fail anyway if: ${{ !github.event.pull_request.head.repo.fork }} name: test-huggingface-transformers-xnnpack - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -763,8 +782,8 @@ jobs: fail-fast: false with: secrets-env: EXECUTORCH_HF_TOKEN - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -905,8 +924,9 @@ jobs: ${CONDA_RUN} python .ci/scripts/test_huggingface_optimum_model.py --model ${MODEL} --recipe ${RECIPE} ${QUANTIZE} test-llama-runner-qnn-linux: + needs: docker-image name: test-llama-runner-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -917,8 +937,8 @@ jobs: mode: [qnn] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 diff --git a/.github/workflows/windows-msvc.yml b/.github/workflows/windows-msvc.yml index bde38a8288f..269f2cf2381 100644 --- a/.github/workflows/windows-msvc.yml +++ b/.github/workflows/windows-msvc.yml @@ -91,7 +91,7 @@ jobs: - name: Install build dependencies shell: pwsh - run: python -m pip install pyyaml torch==2.13.0 --extra-index-url https://download.pytorch.org/whl/test/cpu + run: python -m pip install pyyaml torch==2.14.0 --extra-index-url https://download.pytorch.org/whl/test/cpu - name: Build ExecuTorch shell: pwsh diff --git a/BUCK b/BUCK new file mode 100644 index 00000000000..fd09ad5ebd3 --- /dev/null +++ b/BUCK @@ -0,0 +1 @@ +oncall("executorch") diff --git a/CMakeLists.txt b/CMakeLists.txt index 93e02c521eb..bf1a34e6f47 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -997,7 +997,9 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/cache) list(APPEND _executorch_extensions extension_llm_cache) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/batching) - list(APPEND _executorch_extensions extension_llm_batching) + list(APPEND _executorch_extensions extension_llm_batching + extension_llm_batching_module + ) endif() if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL) diff --git a/CMakePresets.json b/CMakePresets.json index 34a35123ef6..4989990a962 100644 --- a/CMakePresets.json +++ b/CMakePresets.json @@ -90,6 +90,28 @@ "rhs": "Darwin" } }, + { + "name": "apple-framework-resources", + "hidden": true, + "cacheVariables": { + "EXECUTORCH_MLX_SWIFTPM_RESOURCES": "ON" + } + }, + { + "name": "apple-framework-macos", + "displayName": "Build ExecuTorch frameworks for macOS", + "inherits": ["macos", "apple-framework-resources"] + }, + { + "name": "apple-framework-ios", + "displayName": "Build ExecuTorch frameworks for iOS", + "inherits": ["ios", "apple-framework-resources"] + }, + { + "name": "apple-framework-ios-simulator", + "displayName": "Build ExecuTorch frameworks for iOS Simulator", + "inherits": ["ios-simulator", "apple-framework-resources"] + }, { "name": "linux", "displayName": "Build ExecuTorch for Linux", diff --git a/Package.swift b/Package.swift index e1fac90ad93..356828aaa95 100644 --- a/Package.swift +++ b/Package.swift @@ -66,6 +66,13 @@ let products = deliverables([ ], ], "executorch": [ + "frameworks": [ + "Accelerate", + "CoreGraphics", + "CoreImage", + "CoreVideo", + "Foundation", + ], "libraries": [ "c++", ], diff --git a/README-wheel.md b/README-wheel.md index 8ee58b56f2f..86edce654cb 100644 --- a/README-wheel.md +++ b/README-wheel.md @@ -8,6 +8,16 @@ The `executorch` pip package is in beta. * Supported python versions: 3.10, 3.11, 3.12, 3.13, 3.14 * Compatible systems: Linux x86_64, Linux aarch64, macOS aarch64 +Backend export tools can require optional Python dependencies. For example, +install the dependencies needed for Ethos-U ahead-of-time (AOT) export with: + +```bash +pip install 'executorch[ethos_u]' +``` + +The `ethos_u` extra does not install components needed to build or run a target +runtime. + To build a minimal wheel from source, set `EXECUTORCH_BUILD_MINIMAL=1` when running `pip wheel` or `pip install`. That wheel contains the Python EXIR export path and `flatc` for `.pte` diff --git a/README.md b/README.md index 3e7037640cf..cf35f80e2bc 100644 --- a/README.md +++ b/README.md @@ -64,6 +64,10 @@ Learn more: [How ExecuTorch Works](https://docs.pytorch.org/executorch/main/intr pip install executorch ``` +Backend export tools can require optional dependencies. For example, use +`pip install 'executorch[ethos_u]'` for Ethos-U AOT export. Embedded +toolchains, simulators, and target runtimes are installed separately. + For platform-specific setup (Android, iOS, embedded systems), see the [Quick Start](https://docs.pytorch.org/executorch/main/quick-start-section.html) documentation for additional info. ### Export and Deploy in 3 Steps diff --git a/backends/aoti/BUCK b/backends/aoti/BUCK index 77871de4469..f91c46c0f20 100644 --- a/backends/aoti/BUCK +++ b/backends/aoti/BUCK @@ -1,3 +1,5 @@ load("targets.bzl", "define_common_targets") +oncall("executorch") + define_common_targets() diff --git a/backends/aoti/common_shims.cpp b/backends/aoti/common_shims.cpp index f3a34a09987..e83a9576b6c 100644 --- a/backends/aoti/common_shims.cpp +++ b/backends/aoti/common_shims.cpp @@ -159,6 +159,11 @@ AOTITorchError aoti_torch_get_numel(Tensor* tensor, int64_t* ret_numel) { return Error::Ok; } +AOTITorchError aoti_torch_is_defined(Tensor* tensor, bool* ret_is_defined) { + *ret_is_defined = tensor != nullptr; + return Error::Ok; +} + // Device and layout utility functions int32_t aoti_torch_device_type_cpu() { // Let's say cpu is 0 for ET as well diff --git a/backends/aoti/common_shims.h b/backends/aoti/common_shims.h index d057279e22a..0b85b09ee51 100644 --- a/backends/aoti/common_shims.h +++ b/backends/aoti/common_shims.h @@ -62,6 +62,10 @@ aoti_torch_get_dim(Tensor* tensor, int64_t* ret_dim); AOTI_SHIM_EXPORT AOTITorchError aoti_torch_get_numel(Tensor* tensor, int64_t* ret_numel); +// PyTorch has an undefined-tensor state with no equivalent here: null check. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_is_defined(Tensor* tensor, bool* ret_is_defined); + // Utility functions for device and layout information AOTI_SHIM_EXPORT int32_t aoti_torch_device_type_cpu(); AOTI_SHIM_EXPORT int32_t aoti_torch_layout_strided(); diff --git a/backends/aoti/common_shims_slim.cpp b/backends/aoti/common_shims_slim.cpp index c8c7408aa62..d9e255bf220 100644 --- a/backends/aoti/common_shims_slim.cpp +++ b/backends/aoti/common_shims_slim.cpp @@ -68,6 +68,14 @@ AOTITorchError aoti_torch_get_numel(Tensor* tensor, int64_t* ret_numel) { return Error::Ok; } +AOTITorchError aoti_torch_is_defined(Tensor* tensor, bool* ret_is_defined) { + if (ret_is_defined == nullptr) { + return Error::InvalidArgument; + } + *ret_is_defined = tensor != nullptr && tensor->defined(); + return Error::Ok; +} + int32_t aoti_torch_layout_strided() { // Slimtensor only support strided layout, the return value will always be 0, // a.k.a at::Layout::Strided; diff --git a/backends/aoti/common_shims_slim.h b/backends/aoti/common_shims_slim.h index c5a5cab9413..0d60f578605 100644 --- a/backends/aoti/common_shims_slim.h +++ b/backends/aoti/common_shims_slim.h @@ -51,6 +51,10 @@ aoti_torch_get_dim(Tensor* tensor, int64_t* ret_dim); AOTI_SHIM_EXPORT AOTITorchError aoti_torch_get_numel(Tensor* tensor, int64_t* ret_numel); +// Undefined means either a null handle or a tensor whose storage was released. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_is_defined(Tensor* tensor, bool* ret_is_defined); + AOTI_SHIM_EXPORT int32_t aoti_torch_layout_strided(); // ============================================================ diff --git a/backends/aoti/slim/c10/core/BUCK b/backends/aoti/slim/c10/core/BUCK index 77871de4469..f91c46c0f20 100644 --- a/backends/aoti/slim/c10/core/BUCK +++ b/backends/aoti/slim/c10/core/BUCK @@ -1,3 +1,5 @@ load("targets.bzl", "define_common_targets") +oncall("executorch") + define_common_targets() diff --git a/backends/aoti/slim/c10/core/test/BUCK b/backends/aoti/slim/c10/core/test/BUCK index 981043e51d7..6b124bd857b 100644 --- a/backends/aoti/slim/c10/core/test/BUCK +++ b/backends/aoti/slim/c10/core/test/BUCK @@ -1,4 +1,6 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("targets.bzl", "define_common_targets") +oncall("executorch") + fbcode_target(_kind = define_common_targets,) diff --git a/backends/aoti/slim/c10/macros/BUCK b/backends/aoti/slim/c10/macros/BUCK index 981043e51d7..6b124bd857b 100644 --- a/backends/aoti/slim/c10/macros/BUCK +++ b/backends/aoti/slim/c10/macros/BUCK @@ -1,4 +1,6 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("targets.bzl", "define_common_targets") +oncall("executorch") + fbcode_target(_kind = define_common_targets,) diff --git a/backends/aoti/slim/core/BUCK b/backends/aoti/slim/core/BUCK index 981043e51d7..6b124bd857b 100644 --- a/backends/aoti/slim/core/BUCK +++ b/backends/aoti/slim/core/BUCK @@ -1,4 +1,6 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("targets.bzl", "define_common_targets") +oncall("executorch") + fbcode_target(_kind = define_common_targets,) diff --git a/backends/aoti/slim/core/test/BUCK b/backends/aoti/slim/core/test/BUCK index 981043e51d7..6b124bd857b 100644 --- a/backends/aoti/slim/core/test/BUCK +++ b/backends/aoti/slim/core/test/BUCK @@ -1,4 +1,6 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("targets.bzl", "define_common_targets") +oncall("executorch") + fbcode_target(_kind = define_common_targets,) diff --git a/backends/aoti/slim/factory/BUCK b/backends/aoti/slim/factory/BUCK index 981043e51d7..6b124bd857b 100644 --- a/backends/aoti/slim/factory/BUCK +++ b/backends/aoti/slim/factory/BUCK @@ -1,4 +1,6 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("targets.bzl", "define_common_targets") +oncall("executorch") + fbcode_target(_kind = define_common_targets,) diff --git a/backends/aoti/slim/util/BUCK b/backends/aoti/slim/util/BUCK index 77871de4469..f91c46c0f20 100644 --- a/backends/aoti/slim/util/BUCK +++ b/backends/aoti/slim/util/BUCK @@ -1,3 +1,5 @@ load("targets.bzl", "define_common_targets") +oncall("executorch") + define_common_targets() diff --git a/backends/aoti/slim/util/test/BUCK b/backends/aoti/slim/util/test/BUCK index 77871de4469..f91c46c0f20 100644 --- a/backends/aoti/slim/util/test/BUCK +++ b/backends/aoti/slim/util/test/BUCK @@ -1,3 +1,5 @@ load("targets.bzl", "define_common_targets") +oncall("executorch") + define_common_targets() diff --git a/backends/apple/coreml/scripts/build_tests.sh b/backends/apple/coreml/scripts/build_tests.sh index 0203e5027a2..f14463f032e 100755 --- a/backends/apple/coreml/scripts/build_tests.sh +++ b/backends/apple/coreml/scripts/build_tests.sh @@ -36,7 +36,7 @@ cmake "$EXECUTORCH_ROOT_PATH" -B"$CMAKE_EXECUTORCH_BUILD_DIR_PATH" \ -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=OFF \ -DEXECUTORCH_BUILD_XNNPACK=OFF -cmake --build "$CMAKE_EXECUTORCH_BUILD_DIR_PATH" -j9 -t executorch +cmake --build "$CMAKE_EXECUTORCH_BUILD_DIR_PATH" -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) -t executorch # Build protobuf echo "ExecuTorch: Building libprotobuf-lite" @@ -53,7 +53,7 @@ cmake "$PROTOBUF_DIR_PATH/cmake" -B"$CMAKE_PROTOBUF_BUILD_DIR_PATH" \ -DCMAKE_MACOSX_BUNDLE=OFF \ -DCMAKE_CXX_STANDARD=17 -cmake --build "$CMAKE_PROTOBUF_BUILD_DIR_PATH" -j9 -t libprotobuf-lite +cmake --build "$CMAKE_PROTOBUF_BUILD_DIR_PATH" -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) -t libprotobuf-lite # Copy required libraries echo "ExecuTorch: Copying libraries" diff --git a/backends/apple/metal/CMakeLists.txt b/backends/apple/metal/CMakeLists.txt index 4d242eae235..34477cde833 100644 --- a/backends/apple/metal/CMakeLists.txt +++ b/backends/apple/metal/CMakeLists.txt @@ -14,7 +14,7 @@ # ~~~ # It should also be cmake-lint clean. # -cmake_minimum_required(VERSION 3.29) +cmake_minimum_required(VERSION 3.26) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) diff --git a/backends/apple/metal/metal_backend.py b/backends/apple/metal/metal_backend.py index 57ca0ddf83e..0b9a852343b 100644 --- a/backends/apple/metal/metal_backend.py +++ b/backends/apple/metal/metal_backend.py @@ -32,16 +32,24 @@ def get_device_name(cls) -> str: @classmethod def get_supported_fallback_kernels(cls) -> Dict[str, Any]: return { + # An operator named the way it is registered also needs the name + # Inductor derives for it, so several appear under both. "aoti_torch_mps_addmm_out": None, "aoti_torch_mps_bmm_out": None, "aoti_torch_mps_convolution": None, "aoti_torch_mps_mm_out": None, "at::_ops::_scaled_dot_product_attention_math_for_mps::call": None, + "aoti_torch_mps__scaled_dot_product_attention_math_for_mps": None, "at::_ops::_scaled_dot_product_attention_math_for_mps_v2::call": None, + "aoti_torch_mps__scaled_dot_product_attention_math_for_mps_v2": None, "torchao::_linear_fp_act_4bit_weight": None, + "aoti_torch_mps__linear_fp_act_4bit_weight": None, "at::_ops::topk::call": None, + "aoti_torch_mps_topk": None, "metal::gather_qmv": None, + "aoti_torch_mps_gather_qmv": None, "metal::gated_delta_rule": None, + "aoti_torch_mps_gated_delta_rule": None, } @classmethod diff --git a/backends/arm/BUCK b/backends/arm/BUCK index eb8c6bc7590..f6746f29951 100644 --- a/backends/arm/BUCK +++ b/backends/arm/BUCK @@ -9,6 +9,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "constants", diff --git a/backends/arm/CMakeLists.txt b/backends/arm/CMakeLists.txt index 45640e8f3dd..b7e37748c73 100644 --- a/backends/arm/CMakeLists.txt +++ b/backends/arm/CMakeLists.txt @@ -315,5 +315,26 @@ if(EXECUTORCH_BUILD_VGF) executorch_target_link_options_shared_lib(vgf_backend) + if(EXECUTORCH_BUILD_TESTS) + add_executable( + vgf_vulkan_features_test + ${EXECUTORCH_ROOT}/backends/arm/test/vgf_vulkan_features_test.cpp + ) + target_include_directories( + vgf_vulkan_features_test + PRIVATE ${_common_include_directories} ${VULKAN_HEADERS_PATH} + ${VOLK_HEADERS_PATH} + ) + target_compile_options( + vgf_vulkan_features_test PRIVATE -DUSE_VULKAN_WRAPPER -DUSE_VULKAN_VOLK + ) + if(TARGET GTest::gtest_main) + target_link_libraries(vgf_vulkan_features_test PRIVATE GTest::gtest_main) + else() + target_link_libraries(vgf_vulkan_features_test PRIVATE gtest gtest_main) + endif() + add_test(NAME vgf_vulkan_features_test COMMAND vgf_vulkan_features_test) + endif() + # end config for VGF builds endif() diff --git a/backends/arm/README.md b/backends/arm/README.md index ce0c49919e4..d5986b5ccd1 100644 --- a/backends/arm/README.md +++ b/backends/arm/README.md @@ -156,6 +156,49 @@ compile specs, see: Additional examples are available in `examples/arm`. +#### Export recipes + +An `ExportRecipe` bundles those steps for a target, so a standard export needs +no compile spec, quantizer or partitioner of its own. Each recipe carries the +settings its target expects: + +```python +from executorch.backends.arm.recipes.arm_recipe_types import ArmRecipeType +from executorch.export import export, ExportRecipe + +session = export( + model=model, + example_inputs=[example_inputs], + export_recipe=ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U55_INT8), +) +session.save_to_pte("model") +``` + +The recipe quantizes the model, so pass `example_inputs` that are representative +of real data; they are used to calibrate. + +Available recipes: + +| Recipe | Target | +| --- | --- | +| `ETHOS_U55_INT8`, `ETHOS_U65_INT8`, `ETHOS_U85_INT8` | Ethos-U NPUs, int8 | +| `TOSA_FP`, `TOSA_INT8`, `TOSA_A16W8` | TOSA, for testing without hardware | +| `VGF_FP`, `VGF_INT8` | VGF, for the ML SDK for Vulkan | + +The Ethos-U recipes accept `macs`, `system_config`, `memory_mode`, +`extra_flags` and `config_ini`, matching the corresponding Vela options: + +```python +ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U85_INT8, macs=512) +``` + +`macs` is validated against the accelerator configurations the installed Vela +accepts, so an unsupported count fails at recipe construction rather than during +compilation. + +Reach for the step-by-step flow above when a recipe does not fit -- a custom +quantization scheme, extra passes, or a compile spec the recipe does not expose. + ### Direct Drive (experimental, Ethos-U85 on Linux) workflow Direct Drive enables execution on Ethos-U85 via the Linux driver stack. diff --git a/backends/arm/_passes/BUCK b/backends/arm/_passes/BUCK index 45aca62d95f..49b6ba93287 100644 --- a/backends/arm/_passes/BUCK +++ b/backends/arm/_passes/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "core", @@ -65,7 +67,6 @@ fbcode_target( "//executorch/backends/transforms:fuse_duplicate_users_pass", "//executorch/backends/transforms:propagate_view_copy_permute_pass", "//executorch/backends/transforms:fuse_identical_input_transforms_pass", - "//executorch/backends/transforms:convert_conv1d_to_conv2d_pass", "//executorch/backends/transforms:fuse_view_copy", "//executorch/backends/transforms:remove_getitem_op", "//executorch/backends/transforms:replace_scalar_with_tensor", diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index 2b9fcc2e7eb..7d93aa38621 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -17,9 +17,10 @@ from .canonicalize_gather_pass import CanonicalizeGatherPass # noqa from .canonicalize_view_copy_permute_pass import CanonicalizeViewCopyPermutePass # noqa from .cast_int64_pass import CastInt64BuffersToInt32Pass # noqa +from .cast_int_comparison_inputs_pass import CastIntComparisonInputsPass # noqa from .cast_to_int32_pass import CastToInt32Pass # noqa from .constant_folding_pass import ConstantFoldingPass # noqa -from .conv1d_unsqueeze_pass import Conv1dUnsqueezePass # noqa +from .convert_bool_sum_pass import ConvertBoolSumPass # noqa from .convert_elu_params import ConvertELUParamsPass # noqa from .convert_expand_copy_to_repeat import ConvertExpandCopyToRepeatPass # noqa from .convert_full_like_to_full_pass import ConvertFullLikeToFullPass # noqa @@ -74,6 +75,7 @@ DecomposeIndexTensorToGatherPass, ) from .decompose_int_pow_pass import DecomposeIntPowPass # noqa +from .decompose_isinf_isnan_pass import DecomposeIsInfAndIsNanPass # noqa from .decompose_large_stride_maxpool2d_pass import ( # noqa DecomposeLargeStrideMaxPool2dForU55Pass, ) @@ -95,8 +97,12 @@ from .decompose_quant_nodes import DecomposeQuantNodesPass # noqa from .decompose_remainder_pass import DecomposeRemainderPass # noqa from .decompose_rnn_pass import DecomposeRnnPass # noqa +from .decompose_roll_pass import DecomposeRollPass # noqa from .decompose_round_pass import DecomposeRoundPass # noqa from .decompose_sdpa_pass import DecomposeScaledDotProductAttentionPass # noqa +from .decompose_sdpa_with_regular_softmax_pass import ( # noqa + DecomposeSDPAWithRegularSoftmaxPass, +) from .decompose_select import DecomposeSelectPass # noqa from .decompose_select_scatter_pass import DecomposeSelectScatterPass # noqa from .decompose_sign_pass import DecomposeSignPass # noqa @@ -118,6 +124,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_const_shapes_pass import DeduplicateConstShapesPass # noqa from .deduplicate_get_attr_pass import DeduplicateGetAttrPass # noqa from .ensure_unique_output_nodes_pass import EnsureUniqueOutputNodesPass # noqa from .exir_to_tosa_pass import ExirToTosaPass # noqa @@ -169,6 +176,9 @@ from .normalize_index_put_none_indices_pass import ( # noqa NormalizeIndexPutNoneIndicesPass, ) +from .normalize_max_pool2d_input_rank_pass import ( # noqa + NormalizeMaxPool2dInputRankPass, +) from .normalize_while_initial_args_pass import NormalizeWhileInitialArgsPass # noqa from .promote_bool_operands_pass import PromoteBoolOperandsPass # noqa from .propagate_view_copy_permute_pass import ( # noqa diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index ef30bec9f00..b74860954ba 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -17,11 +17,12 @@ CanonicalizeGatherPass, CanonicalizeViewCopyPermutePass, CastInt64BuffersToInt32Pass, + CastIntComparisonInputsPass, CastToInt32Pass, ComputeConstantOpsAOTPass, ConstantFoldingPass, ControlFlowConstInlinePass, - Conv1dUnsqueezePass, + ConvertBoolSumPass, ConvertEluFamilyToEluPass, ConvertELUParamsPass, ConvertExpandCopyToRepeatPass, @@ -69,6 +70,7 @@ DecomposeIndexSelectToGatherPass, DecomposeIndexTensorToGatherPass, DecomposeIntPowPass, + DecomposeIsInfAndIsNanPass, DecomposeLargeStrideMaxPool2dForU55Pass, DecomposeLayerNormPass, DecomposeLeakyReLUPass, @@ -88,8 +90,10 @@ DecomposeQuantNodesPass, DecomposeRemainderPass, DecomposeRnnPass, + DecomposeRollPass, DecomposeRoundPass, DecomposeScaledDotProductAttentionPass, + DecomposeSDPAWithRegularSoftmaxPass, DecomposeSelectPass, DecomposeSelectScatterPass, DecomposeSignPass, @@ -107,6 +111,7 @@ DecomposeVarPass, DecomposeWhereScalarOtherPass, DecorateFp32toInt32CastingPass, + DeduplicateConstShapesPass, DeduplicateGetAttrPass, EnsureUniqueOutputNodesPass, ExirToTosaPass, @@ -137,6 +142,7 @@ NormalizeDelegateIOLayoutPass, NormalizeIndexPutBoolIndexTensorPass, NormalizeIndexPutNoneIndicesPass, + NormalizeMaxPool2dInputRankPass, NormalizeTransformInputPlaceholdersPass, NormalizeWhileInitialArgsPass, PromoteBoolOperandsPass, @@ -222,6 +228,7 @@ def _graph_pass_name(graph_pass: Callable[[GraphModule], PassResult | None]) -> class _ExportedProgramGraphPassAdapter(ExportedProgramPassBase): def __init__(self, graph_pass: Callable[[GraphModule], PassResult | None]) -> None: self.graph_pass = graph_pass + self.__name__ = _graph_pass_name(graph_pass) def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: graph_pass = cast(Any, self.graph_pass) @@ -310,10 +317,7 @@ def configure_skip_passes(self) -> tuple[type, ...]: skip_set.add(DecomposeLeakyReLUPass) match config.sdpa_safe_softmax_guard: # type: ignore[attr-defined] - case ( - SDPASafeSoftmaxGuardPolicy.PRESERVE - | SDPASafeSoftmaxGuardPolicy.REMOVE_WHEN_PROVEN - ): + case SDPASafeSoftmaxGuardPolicy.PRESERVE | SDPASafeSoftmaxGuardPolicy.AUTO: skip_set.add(RemoveSafeSoftmaxGuardPass) case SDPASafeSoftmaxGuardPolicy.REMOVE: pass @@ -456,6 +460,21 @@ def _tosa_context(self, graph_module: GraphModule) -> TosaLoweringContext: shape_env = _get_shape_env_from_gm(graph_module) return TosaLoweringContext(self.tosa_spec, shape_env) + def transform_for_pre_decomposition_pipeline( + self, exported_program: ExportedProgram + ) -> ExportedProgram: + """Apply Arm passes before default ATen decompositions.""" + config = self.compile_spec._get_pass_pipeline_config() + passes: list[ExportPass] = [] + + if config.sdpa_safe_softmax_guard is SDPASafeSoftmaxGuardPolicy.AUTO: + passes.append(DecomposeSDPAWithRegularSoftmaxPass()) + + if passes: + self.add_passes(passes) + self._transform(exported_program, exported_program.graph_module) + return exported_program + def _transform_graph_module(self, graph_module: GraphModule): # TFA and control-flow submodule paths operate on bare GraphModules # without a standalone ExportedProgram to keep in sync. @@ -559,6 +578,7 @@ def _tosa_pipeline( RemoveGetItemPass(), FuseBatchNorm2dPass(exported_program), DecomposeBatchNormNoStatsPass(), + DecomposeIsInfAndIsNanPass(), DecomposeLogitPass(), DecomposeMaskedFillPass(), DecomposeRoundPass(), @@ -576,26 +596,27 @@ def _tosa_pipeline( DecomposeExpm1Pass(), DecomposeIntPowPass(), DecomposeLog1pPass(), + ConvertBoolSumPass(), PromoteBoolOperandsPass(), DecomposeSinhPass(), DecomposeSignPass(), DecomposeFlipPass(), + DecomposeRollPass(), DecomposeFloorDividePass(), DecomposeGeluPass(), DecomposeAddSubAlphaPass(), DecomposeGroupedConvPass(), DecomposeUnfoldToGatherPass(use_slice=self.tosa_spec.is_U55_subset), DecomposeEmbeddingPass(), - DecomposeIndexSelectToGatherPass(), CastInt64BuffersToInt32Pass(exported_program), + DecomposeIndexSelectToGatherPass(exported_program), DecomposeStridedSliceCopyPass(), DecomposeSliceScatterPass(), AccumulateIndexPutPass(), - DecomposeIndexTensorToGatherPass(), + DecomposeIndexTensorToGatherPass(exported_program), DecomposeAdaptiveAvgPool2dPass(), DecomposeDynamicAdaptiveAvgPool2dPass(), DecomposeAvgPool2dPass(), - Conv1dUnsqueezePass(exported_program), ] ) @@ -603,6 +624,7 @@ def _tosa_pipeline( self.add_passes( [ ReplaceScalarWithTensorByProfilePass(), + CastIntComparisonInputsPass(), RewriteLeLtToGeGtPass(), DecomposeLeakyReLUPass(), # Emits full_like so before ConvertFullLikeToFullPass DecomposePReLUPass(), @@ -635,6 +657,7 @@ def _tosa_pipeline( UnsqueezeBeforeRepeatPass(), DecomposeCumsumPass(exported_program), DecomposeAsStridedCopyPass(), + NormalizeMaxPool2dInputRankPass(), DecomposeMaxPool2dPass(), DecomposeLargeStrideMaxPool2dForU55Pass(), SizeAdjustInputPass(), @@ -707,6 +730,7 @@ def _tosa_pipeline( # fusing generated RESCALE users can corrupt distinct quantized paths. FuseDuplicateUsersPass(), InsertRescalePass(), + DeduplicateConstShapesPass(), EnsureUniqueOutputNodesPass(), ] ) @@ -747,6 +771,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule): DecomposeDynamicFullPass(tfa_pass=True), ConvertInt64ConstOpsToInt32Pass(tfa_pass=True), ConvertInt64OutputOpsToInt32Pass(tfa_pass=True), + ConvertBoolSumPass(tfa_pass=True), InsertInt32CastsAfterInt64PlaceholdersPass(tfa_pass=True), FoldScalarMulIntoConvPass(tfa_pass=True), DecomposeEmbeddingPass(tfa_pass=True), diff --git a/backends/arm/_passes/aten_to_tosa_comparison.py b/backends/arm/_passes/aten_to_tosa_comparison.py new file mode 100644 index 00000000000..d444ce47a96 --- /dev/null +++ b/backends/arm/_passes/aten_to_tosa_comparison.py @@ -0,0 +1,27 @@ +# 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 executorch.backends.transforms.aten_to_dialect_pass import ( + AtenToDialectPass, + DialectNodeSpec, +) +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import Node + + +def rewrite_comparison_operator( + node: Node, pass_: AtenToDialectPass +) -> DialectNodeSpec | None: + match node.target: + case exir_ops.edge.aten.eq.Tensor: + target = exir_ops.backend.tosa.EQUAL.default + case exir_ops.edge.aten.ge.Tensor: + target = exir_ops.backend.tosa.GREATER_EQUAL.default + case exir_ops.edge.aten.gt.Tensor: + target = exir_ops.backend.tosa.GREATER.default + case _: + return None + + return DialectNodeSpec(target, node.args, dict(node.kwargs)) diff --git a/backends/arm/_passes/aten_to_tosa_tensor_operators.py b/backends/arm/_passes/aten_to_tosa_tensor_operators.py index 6705514c606..84afb096c4e 100644 --- a/backends/arm/_passes/aten_to_tosa_tensor_operators.py +++ b/backends/arm/_passes/aten_to_tosa_tensor_operators.py @@ -56,12 +56,6 @@ def rewrite_binary_operator( target = exir_ops.backend.tosa.ARITHMETIC_RIGHT_SHIFT.default case exir_ops.edge.aten.bitwise_xor.Tensor: target = exir_ops.backend.tosa.BITWISE_XOR.default - case exir_ops.edge.aten.eq.Tensor: - target = exir_ops.backend.tosa.EQUAL.default - case exir_ops.edge.aten.ge.Tensor: - target = exir_ops.backend.tosa.GREATER_EQUAL.default - case exir_ops.edge.aten.gt.Tensor: - target = exir_ops.backend.tosa.GREATER.default case exir_ops.edge.aten.logical_and.default: target = exir_ops.backend.tosa.LOGICAL_AND.default case exir_ops.edge.aten.logical_or.default: diff --git a/backends/arm/_passes/cast_int_comparison_inputs_pass.py b/backends/arm/_passes/cast_int_comparison_inputs_pass.py new file mode 100644 index 00000000000..6a0b19b0703 --- /dev/null +++ b/backends/arm/_passes/cast_int_comparison_inputs_pass.py @@ -0,0 +1,61 @@ +# 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 Set, Type + +import torch + +from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass +from executorch.backends.arm.tosa.specification import get_context_spec +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass + + +class CastIntComparisonInputsPass(ArmOpTargetedPass): + """Cast integer comparison inputs to a lossless floating-point type.""" + + _passes_required_after: Set[Type[ExportPass]] = set() + + target_ops = { + exir_ops.edge.aten.eq.Tensor, + exir_ops.edge.aten.ne.Tensor, + exir_ops.edge.aten.ge.Tensor, + exir_ops.edge.aten.gt.Tensor, + exir_ops.edge.aten.le.Tensor, + exir_ops.edge.aten.lt.Tensor, + } + castable_dtypes = {torch.int8, torch.int16} + + def should_run_pass(self, graph_module: torch.fx.GraphModule) -> bool: + tosa_spec = get_context_spec() + return ( + tosa_spec.support_float() + and not tosa_spec.support_integer() + and super().should_run_pass(graph_module) + ) + + def call_operator(self, op, args, kwargs, meta): + if op not in self.target_ops: + return super().call_operator(op, args, kwargs, meta) + + if not all(arg.data.dtype in self.castable_dtypes for arg in args): + return super().call_operator(op, args, kwargs, meta) + + cast_dtype = ( + torch.float16 + if all(arg.data.dtype == torch.int8 for arg in args) + else torch.float32 + ) + casted_args = [] + for arg in args: + casted_args.append( + super().call_operator( + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, + (arg,), + {"dtype": cast_dtype}, + meta, + ) + ) + return super().call_operator(op, tuple(casted_args), kwargs, meta) diff --git a/backends/arm/_passes/conv1d_unsqueeze_pass.py b/backends/arm/_passes/conv1d_unsqueeze_pass.py deleted file mode 100644 index c01fbdb9f60..00000000000 --- a/backends/arm/_passes/conv1d_unsqueeze_pass.py +++ /dev/null @@ -1,41 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# 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. - -from typing import Set, Type - -from executorch.backends.arm._passes import ArmOpTargetedPass -from executorch.backends.arm._passes.convert_squeezes_to_view import ( - ConvertSqueezesToViewPass, -) -from executorch.backends.arm._passes.rewrite_conv_pass import RewriteConvPass -from executorch.backends.arm._passes.size_adjust_input_pass import SizeAdjustInputPass -from executorch.backends.transforms.convert_conv1d_to_conv2d_pass import ( - ConvertConv1dToConv2dPass, -) -from executorch.exir import ExportedProgram -from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.pass_base import ExportPass - - -class Conv1dUnsqueezePass(ConvertConv1dToConv2dPass, ArmOpTargetedPass): - """Arm wrapper for the shared Conv1d-to-Conv2d transform.""" - - _passes_required_after: Set[Type[ExportPass]] = { - ConvertSqueezesToViewPass, - RewriteConvPass, - SizeAdjustInputPass, - } - target_ops = (exir_ops.edge.aten.convolution.default,) - - def __init__(self, exported_program: ExportedProgram) -> None: - # Grouped-convolution decomposition creates one graph-local weight - # slice per group. Allow the shared pass to add the unit-height - # dimension after these producers. - super().__init__( - exported_program, - graph_local_weight_targets={exir_ops.edge.aten.slice_copy.Tensor}, - ) diff --git a/backends/arm/_passes/convert_bool_sum_pass.py b/backends/arm/_passes/convert_bool_sum_pass.py new file mode 100644 index 00000000000..0b77281272c --- /dev/null +++ b/backends/arm/_passes/convert_bool_sum_pass.py @@ -0,0 +1,81 @@ +# 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 math +from typing import Set, Type + +import torch + +from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass +from executorch.backends.arm._passes.decompose_sum_pass import DecomposeSumPass +from executorch.backends.arm.tosa.specification import ( + get_context_shape_env, + get_context_spec, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, NodeMetadata + + +class ConvertBoolSumPass(ArmOpTargetedPass): + """Compute boolean sums with an int32 accumulator.""" + + _passes_required_after: Set[Type[ExportPass]] = {DecomposeSumPass} + _MAX_EXACT_SUM = torch.iinfo(torch.int32).max + target_ops = ( + torch.ops.aten.sum.dim_IntList, + exir_ops.edge.aten.sum.dim_IntList, + ) + check_allowed_to_transform = True + + def call_operator(self, op, args, kwargs, meta): + tosa_spec = get_context_spec() + if ( + op not in self.target_ops + or args[0].data.dtype != torch.bool + or not self.allowed_to_transform(meta) + or not tosa_spec.support_integer() + or tosa_spec.support_float() + ): + return super().call_operator(op, args, kwargs, meta) + + dims = args[1] + if not dims: + dims = range(args[0].data.dim()) + reduced_elements = math.prod(args[0].data.shape[dim] for dim in dims) + if isinstance(reduced_elements, torch.SymInt): + reduced_elements = ( + get_context_shape_env().bound_sympy(reduced_elements.node.expr).upper + ) + if reduced_elements > self._MAX_EXACT_SUM: + return super().call_operator(op, args, kwargs, meta) + + cast_op = ( + exir_ops.edge.dim_order_ops._to_dim_order_copy.default + if op == exir_ops.edge.aten.sum.dim_IntList + else torch.ops.dim_order_ops._to_dim_order_copy.default + ) + accumulator_input = super().call_operator( + cast_op, + (args[0],), + {"dtype": torch.int32}, + NodeMetadata(args[0].node.meta), + updated=True, + ) + sum_kwargs = dict(kwargs) + sum_kwargs["dtype"] = torch.int32 + accumulator_sum = super().call_operator( + op, + (accumulator_input, *args[1:]), + sum_kwargs, + meta, + updated=True, + ) + return super().call_operator( + cast_op, + (accumulator_sum,), + {"dtype": meta["val"].dtype}, + meta, + updated=True, + ) diff --git a/backends/arm/_passes/decompose_grouped_conv_pass.py b/backends/arm/_passes/decompose_grouped_conv_pass.py index 7a8b744d9e3..5ed65e4cd7e 100644 --- a/backends/arm/_passes/decompose_grouped_conv_pass.py +++ b/backends/arm/_passes/decompose_grouped_conv_pass.py @@ -8,8 +8,8 @@ import torch from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass -from executorch.backends.arm._passes.conv1d_unsqueeze_pass import Conv1dUnsqueezePass from executorch.backends.arm._passes.quant_args import QuantArgs +from executorch.backends.arm._passes.rewrite_conv_pass import RewriteConvPass from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass @@ -46,7 +46,7 @@ class DecomposeGroupedConvPass(ArmOpTargetedPass): """ - _passes_required_after: Set[Type[ExportPass]] = {Conv1dUnsqueezePass} + _passes_required_after: Set[Type[ExportPass]] = {RewriteConvPass} target_ops = ( exir_ops.edge.aten.convolution.default, torch.ops.aten.conv_transpose2d.input, diff --git a/backends/arm/_passes/decompose_index_select_to_gather_pass.py b/backends/arm/_passes/decompose_index_select_to_gather_pass.py index be0d4dbb07c..0a356b6c1d4 100644 --- a/backends/arm/_passes/decompose_index_select_to_gather_pass.py +++ b/backends/arm/_passes/decompose_index_select_to_gather_pass.py @@ -9,12 +9,17 @@ import torch from executorch.backends.arm._passes import ArmOpTargetedPass +from executorch.backends.arm._passes.arm_pass_utils import ( + get_param_tensor, + is_param_node, +) from executorch.backends.arm._passes.convert_expand_copy_to_repeat import ( ConvertExpandCopyToRepeatPass, ) from executorch.backends.arm._passes.convert_squeezes_to_view import ( ConvertSqueezesToViewPass, ) +from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass @@ -71,6 +76,12 @@ class DecomposeIndexSelectToGatherPass(ArmOpTargetedPass): exir_ops.edge.aten.index_select.default, } + def __init__( + self, exported_program: ExportedProgram | None = None, *args, **kwargs + ) -> None: + super().__init__(*args, **kwargs) + self.exported_program = exported_program + def call_operator(self, op, args, kwargs, meta): if op not in self.target_ops: return super().call_operator(op, args, kwargs, meta) @@ -81,6 +92,50 @@ def call_operator(self, op, args, kwargs, meta): x_shape, idx_shape = tuple(x_t.shape), tuple(idx_t.shape) x_rank, idx_rank = len(x_shape), len(idx_shape) + if ( + x_rank >= 1 + and idx_rank == 1 + and self.exported_program is not None + and is_param_node(self.exported_program, index.node) + and all(isinstance(size, int) for size in x_shape) + ): + constant_index = get_param_tensor(self.exported_program, index.node) + if constant_index is not None and constant_index.numel() > 0: + indices = constant_index.tolist() + dim_norm = dim % x_rank + dim_size = x_shape[dim_norm] + if any(index < 0 or index >= dim_size for index in indices): + raise RuntimeError( + f"index_select index out of range for dimension of size {dim_size}" + ) + if indices == list(range(indices[0], indices[0] + len(indices))): + return super().call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + (x, dim_norm, indices[0], indices[-1] + 1), + {}, + meta, + updated=True, + ) + + slices = [] + for index_value in indices: + slices.append( + super().call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + (x, dim_norm, index_value, index_value + 1), + {}, + meta, + updated=True, + ) + ) + return super().call_operator( + exir_ops.edge.aten.cat.default, + (slices, dim_norm), + {}, + meta, + updated=True, + ) + assert x_rank >= 1 and idx_rank == 1 and idx_t.dtype == torch.int32, ( f"[{self.__class__.__name__}] unsupported index_select signature: " f"x_rank={x_rank}, index_rank={idx_rank}, index_dtype={idx_t.dtype} " diff --git a/backends/arm/_passes/decompose_index_tensor_to_gather_pass.py b/backends/arm/_passes/decompose_index_tensor_to_gather_pass.py index 93db9f9d434..3413ce87238 100644 --- a/backends/arm/_passes/decompose_index_tensor_to_gather_pass.py +++ b/backends/arm/_passes/decompose_index_tensor_to_gather_pass.py @@ -10,7 +10,11 @@ import torch from executorch.backends.arm._passes import ArmOpTargetedPass -from executorch.backends.arm._passes.arm_pass_utils import meta_without_qparams +from executorch.backends.arm._passes.arm_pass_utils import ( + get_param_tensor, + is_param_node, + meta_without_qparams, +) from executorch.backends.arm._passes.convert_expand_copy_to_repeat import ( ConvertExpandCopyToRepeatPass, ) @@ -20,12 +24,13 @@ from executorch.backends.arm._passes.replace_scalar_with_tensor_pass import ( ReplaceScalarWithTensorByProfilePass, ) +from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass def get_index_tensor_decomposition(op): - """Return the operator overloads used to lower index.Tensor via TOSA gather. + """Return operators used to lower index.Tensor through TOSA gather. Raises: RuntimeError: If the provided operator is not supported by this pass. @@ -48,12 +53,12 @@ def get_index_tensor_decomposition(op): def _broadcast_shape( shapes: Sequence[Sequence[int]], ) -> list[int]: - """Compute the broadcasted shape (PyTorch/Numpy semantics) for a list of - shapes. + """Compute the broadcasted shape using PyTorch/NumPy semantics. Requirements: - static shape only - - shapes are right-aligned; lower-rank shapes are implicitly front-padded with 1s + - shapes are right-aligned; lower-rank shapes are implicitly + front-padded with 1s - per-axis dims must either match exactly or be 1 Raises: @@ -76,86 +81,121 @@ def _broadcast_shape( class DecomposeIndexTensorToGatherPass(ArmOpTargetedPass): - """Decompose edge.aten.index.Tensor into backend TOSA gather (+ basic - arith). + """Decompose edge.aten.index.Tensor into a TOSA gather and arithmetic. Supported subset: - y = x.index([i0, i1, ..., i{m-1}]) + y = x.index([None, ..., None, i0, i1, ..., i{m-1}]) - where each ik is a Tensor index, and m is the number of index tensors. + where each ik is a Tensor index, m is the number of index tensors, and the + optional leading None entries preserve dimensions before the indexed block. Constraints: - - `indices` list contains only Tensor indices (no None/slice/ellipsis) + - `indices` contains an optional leading run of None entries followed by + only Tensor indices - Each index tensor dtype is int32 - - Index tensor shapes are broadcastable to a common shape `S` (per index.Tensor semantics) - - Only prefix indexing is supported: the `m` tensor indices select elements - from the first `m` dimensions of `x`, so `m <= rank(x)`. + - Index tensor shapes are broadcastable to a common shape `S` (per + index.Tensor semantics) + - The `m` tensor indices select one contiguous block of dimensions after + the leading preserved dimensions. - Static shapes are required - - If `x` has more than 2^31 elements, the computed linear index may overflow int32. + - If `x` has more than 2^31 elements, the computed linear index may + overflow int32. Lowering strategy (single gather) --------------------------------- Let: + - `p` be the number of leading None entries - `S` be the broadcasted index shape - `W = prod(S)` (number of indexed positions) - - `K = prod(x.shape[:m])` (flattened size of the indexed prefix) - - `C = prod(x.shape[m:])` (flattened size of the trailing slice per index) - - `trailing = x.shape[m:]` + - `P = prod(x.shape[:p])` (flattened size of the leading preserved + dimensions) + - `K = prod(x.shape[p:p+m])` (flattened size of the indexed block) + - `C = prod(x.shape[p+m:])` (flattened size of the trailing slice per + index) + - `leading = x.shape[:p]` and `trailing = x.shape[p+m:]` Steps: 1) Compute parameters needed to lower index.Tensor - - `S`, `W`, `K`, `C`, `trailing` - - `lin_scales[i] = stride_i // C`, where `stride_i` are the contiguous-style - strides derived from `x.shape` (for dim i). - 2) Reshape x to `[1, K, C]` (`x_1kc`). - 3) Build linear indices (`lin_1w`) by scaling each flattened index and summing: - lin_1w = unsqueeze0( sum_{i=0..m-1} ( idx_flat[i] * lin_scales[i] ) ) - where: - - `m = len(indices)` - - `idx_flat[i]` is the i-th index tensor after broadcast to `S` and flatten to `[W]` - - `lin_1w` has shape `[1, W]` and is used as the `indices` input to `tosa.GATHER` + - `S`, `W`, `P`, `K`, `C`, `leading`, `trailing` + - `lin_scales` as the contiguous strides of the indexed block + `x.shape[p:p+m]`. + 2) Reshape x to `[P, K, C]` (`x_pkc`). + 3) Build linear indices by scaling and accumulating the flattened index + tensors element-wise: + For each tensor index, broadcast it to `S`, then flatten it: + + idx_broadcast[i] = broadcast_to(indices[p + i], S) + idx_flat[i] = reshape(idx_broadcast[i], [W]) + + For each j in [0, W): + + lin_w[j] = + sum_{i=0..m-1} idx_flat[i][j] * lin_scales[i] + + Equivalently, in tensor notation: + + lin_w = + sum_{i=0..m-1} idx_flat[i] * lin_scales[i] # shape [W] + + Then: + + lin_1w = unsqueeze(lin_w, 0) # [1, W] + lin_pw = lin_1w + if P > 1: + lin_pw = expand(lin_1w, [P, W]) # [P, W] 4) Single gather: - `tosa.GATHER(x=x_1kc, indices=lin_1w) -> [1,W,C]` - 5) Reshape result to `[*S, *trailing]`. + `tosa.GATHER(x=x_pkc, indices=lin_pw) -> [P,W,C]` + 5) Reshape result to `[*leading, *S, *trailing]`. - Example - ------- + Example: Consider: - x.shape = [2, 3, 4] - indices = [i0, i1] # m = 2 + x.shape = [2, 3, 4, 5] + indices = [None, i0, i1] # p = 1, m = 2 i0.shape = [2, 1] i1.shape = [1, 2] + This corresponds to ``x[:, i0, i1, :]``: the first dimension is + preserved, the next two dimensions are indexed, and the last dimension is + trailing. + 1) The index shapes broadcast to: S := [2, 2] W := prod(S) = 4 - We index the first m=2 dimensions of x, so: - K := prod(x.shape[:m]) = 2 * 3 = 6 - C := prod(x.shape[m:]) = 4 - trailing := x.shape[m:] = [4] + We preserve p=1 leading dimension and index the next m=2 dimensions: + leading := x.shape[:p] = [2] + trailing := x.shape[p+m:] = [5] + P := prod(leading) = 2 + K := prod(x.shape[p:p+m]) = 3 * 4 = 12 + C := prod(trailing) = 5 - Contiguous strides of x are [12, 4, 1], so: - lin_scales := [stride0 // C, stride1 // C] = [12//4, 4//4] = [3, 1] + The indexed block has shape [3, 4], so its contiguous strides are: + lin_scales := [4, 1] 2) Values are reshaped to: - x_1kc = view(x, [1, K, C]) = [1, 6, 4] + x_pkc = view(x, [P, K, C]) = [2, 12, 5] 3) After broadcasting and flattening the indices to length W: i0_broadcast, i1_broadcast have shape S=[2,2] i0_flat, i1_flat have shape [W]=[4] - Linear indices are computed as: - lin_w - = lin_scales * [i0_flat, i1_flat] - = 3 * i0_flat + 1 * i1_flat # shape [W] - lin_w is reshaped to [1, W] to match tosa.Gather semantics + Linear indices are computed element-wise as: + for each j in [0, W): + + lin_w[j] = 4 * i0_flat[j] + i1_flat[j] + + hence: + + lin_w = 4 * i0_flat + i1_flat # shape [W] + + lin_w is then unsqueezed to [1, W] and expanded so that + lin_pw.shape = [P, W] = [2, 4]. 4) Single Gather: - out_1wc = tosa.GATHER(values=x_1kc, indices=lin_1w) # [1, 4, 4] + out_pwc = tosa.GATHER(values=x_pkc, indices=lin_pw) # [2, 4, 5] 5) Reshape result: - out = view(out_1wc, [*S, *x.shape[m:]]) # [2, 2, 4] + out = view(out_pwc, [*leading, *S, *trailing]) # [2, 2, 2, 5] """ @@ -169,6 +209,12 @@ class DecomposeIndexTensorToGatherPass(ArmOpTargetedPass): exir_ops.edge.aten.index.Tensor, } + def __init__( + self, exported_program: ExportedProgram | None = None, *args, **kwargs + ) -> None: + super().__init__(*args, **kwargs) + self.exported_program = exported_program + @staticmethod def _shape_to_stride( values_shape: Sequence[int], @@ -181,69 +227,140 @@ def _shape_to_stride( return strides @staticmethod - def _validate_tensor_indices(indices): + def _validate_and_split_indices(indices): assert ( isinstance(indices, (list, tuple)) and len(indices) > 0 ), f"index.Tensor expects non-empty indices list/tuple, got {type(indices)}." - for i, idx in enumerate(indices): - assert ( - idx is not None - ), f"index.Tensor: None indices are not supported at the moment (indices[{i}] is None)." + leading_rank = 0 + while leading_rank < len(indices) and indices[leading_rank] is None: + leading_rank += 1 + + tensor_indices = indices[leading_rank:] + assert tensor_indices, "index.Tensor expects at least one tensor index." + for i, idx in enumerate(tensor_indices, start=leading_rank): + assert idx is not None, ( + "index.Tensor supports None entries only before all tensor indices " + f"(indices[{i}] is None)." + ) assert ( idx.data.dtype == torch.int32 ), "index.Tensor requires index dtype must be int32" - def _compute_index_tensor_params(self, x, m, index_shapes): - """Compute shape/stride-derived parameters needed to lower - edge.aten.index.Tensor. + return leading_rank, tensor_indices - Derives the broadcasted index shape and the scale factors used to flatten and - acculumulate multi-dimensional indices into a single gather index, following - the S/W/K/C notation described in the class docstring. + def _compute_index_tensor_params(self, x, leading_rank, m, index_shapes): + """Compute parameters needed to lower edge.aten.index.Tensor. + + Derives the broadcasted index shape and the scale factors used to + flatten and accumulate multi-dimensional indices into a single gather + index, following the S/W/P/K/C notation described in the class + docstring. Args: - x: Values tensor being indexed. - m: Number of tensor indices (i.e., len(indices)). - index_shapes: Shapes corresponding to each tensor index. + x (ProxyValue): Values tensor being indexed. + leading_rank (int): Number of leading dimensions preserved by None + entries. + m (int): Number of tensor indices. + index_shapes (Sequence[Sequence[int]]): Shapes corresponding to + each tensor index. Returns: - (x_data, S, W, K, C, trailing, lin_scales), where: - - x_data is `x.data` (FakeTensor) - - trailing is `x.shape[m:]` as a list of ints - - lin_scales are per-dimension scale factors for linearization + tuple: `(x_data, S, W, P, K, C, leading, trailing, lin_scales)`, + where `x_data` is `x.data`, `leading` and `trailing` contain + the preserved dimensions, and `lin_scales` contains the + indexed-block strides used for linearization. """ - x_data = x.data # FakeTensor x_shape = tuple(x_data.shape) x_rank = len(x_shape) assert x_rank >= 1, f"index.Tensor expects x rank>=1, got {x_shape}." - assert ( - m <= x_rank - ), f"index.Tensor has too many indices ({m}) for x rank {x_rank}." + assert leading_rank + m <= x_rank, ( + "index.Tensor has more preserved and indexed dimensions " + f"({leading_rank + m}) than the input rank ({x_rank})." + ) # Broadcast shape S for indices, and flattened length W S = _broadcast_shape(index_shapes) W = math.prod(S) if S else 1 - # Compute gather factors K and C for leading-dims indexing - leading = list(x_shape[:m]) - trailing = list(x_shape[m:]) - K = math.prod(leading) if leading else 1 + # Compute gather factors for the preserved, indexed, and trailing blocks. + leading = list(x_shape[:leading_rank]) + indexed = list(x_shape[leading_rank : leading_rank + m]) + trailing = list(x_shape[leading_rank + m :]) + P = math.prod(leading) if leading else 1 + K = math.prod(indexed) if indexed else 1 C = math.prod(trailing) if trailing else 1 - # Strides for linearization (contiguous-style) - strides = self._shape_to_stride(x_shape) - - # Stride/C divisibility is guaranteed for contiguous strides and C=prod(trailing). - lin_scales: list[int] = [] - for i in range(m): - stride = strides[i] - lin_scales.append(stride // C) + lin_scales = self._shape_to_stride(indexed) + + return x_data, S, W, P, K, C, leading, trailing, lin_scales + + def _decompose_constant_index(self, x, indices, meta): + tensor_indices = [ + (dim, index) for dim, index in enumerate(indices) if index is not None + ] + if ( + self.exported_program is None + or len(tensor_indices) != 1 + or any(not isinstance(size, int) for size in x.data.shape) + ): + return None + + indexed_dim, index_tensor = tensor_indices[0] + if not is_param_node(self.exported_program, index_tensor.node): + return None + + constant_index = get_param_tensor(self.exported_program, index_tensor.node) + if ( + constant_index is None + or constant_index.dim() != 1 + or constant_index.numel() == 0 + ): + return None + + indexed_dim_size = x.data.shape[indexed_dim] + index_values = [] + for value in constant_index.tolist(): + normalized_value = value if value >= 0 else value + indexed_dim_size + if normalized_value < 0 or normalized_value >= indexed_dim_size: + raise IndexError( + f"index {value} is out of bounds for dimension {indexed_dim} " + f"with size {indexed_dim_size}" + ) + index_values.append(normalized_value) + + if index_values == list( + range(index_values[0], index_values[0] + len(index_values)) + ): + return super().call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + (x, indexed_dim, index_values[0], index_values[-1] + 1), + {}, + meta, + updated=True, + ) - return x_data, S, W, K, C, trailing, lin_scales + slices = [] + for index_value in index_values: + slices.append( + super().call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + (x, indexed_dim, index_value, index_value + 1), + {}, + meta, + updated=True, + ) + ) + return super().call_operator( + exir_ops.edge.aten.cat.default, + (slices, indexed_dim), + {}, + meta, + updated=True, + ) def call_operator(self, op, args, kwargs, meta): if op not in self.target_ops: @@ -255,13 +372,32 @@ def call_operator(self, op, args, kwargs, meta): x, indices = args - self._validate_tensor_indices(indices) + tensor_indices = [index for index in indices if index is not None] + if len(tensor_indices) == 1 and tensor_indices[0].data.dtype in ( + torch.bool, + torch.uint8, + ): + return super().call_operator(op, args, kwargs, meta) + + constant_result = self._decompose_constant_index(x, indices, meta) + if constant_result is not None: + return constant_result + + leading_rank, indices = self._validate_and_split_indices(indices) index_shapes = [idx.data.shape for idx in indices] m = len(indices) - x_data, S, W, K, C, trailing, lin_scales = self._compute_index_tensor_params( - x, m, index_shapes - ) + ( + x_data, + S, + W, + P, + K, + C, + leading, + trailing, + lin_scales, + ) = self._compute_index_tensor_params(x, leading_rank, m, index_shapes) ( view_op, @@ -285,16 +421,16 @@ def call_operator(self, op, args, kwargs, meta): updated=True, ) - # ---- x: [1, K, C] ---- - x_1kc = super().call_operator( + # ---- x: [P, K, C] ---- + x_pkc = super().call_operator( view_op, - (x_for_gather, [1, K, C]), + (x_for_gather, [P, K, C]), {}, meta, updated=True, ) - # Build linear index [1, W] from broadcasted indices + # Build linear index [W] from broadcasted indices lin_w = None plain_meta = meta_without_qparams(meta) for i, idx in enumerate(indices): @@ -341,7 +477,7 @@ def call_operator(self, op, args, kwargs, meta): updated=True, ) - # Accumulate into lin_1w: [1, W] + # Accumulate into lin_w: [W] if lin_w is None: lin_w = idx_scaled else: @@ -355,10 +491,10 @@ def call_operator(self, op, args, kwargs, meta): if lin_w is None: raise RuntimeError( - f"[{self.__class__.__name__}] internal error: lin_1w not constructed." + f"[{self.__class__.__name__}] internal error: lin_w not constructed." ) - # Make indices shape [1, W] for tosa.GATHER + # Make indices shape [P, W] for tosa.GATHER. lin_1w = super().call_operator( unsqueeze_op, (lin_w, 0), @@ -366,22 +502,31 @@ def call_operator(self, op, args, kwargs, meta): plain_meta, updated=True, ) + lin_pw = lin_1w + if P > 1: + lin_pw = super().call_operator( + expand_op, + (lin_1w, [P, W]), + {}, + plain_meta, + updated=True, + ) # ---- backend tosa gather --- - # tosa.GATHER(x=[1,K,C], indices=[1,W]) -> [1,W,C] - gathered_1wc = super().call_operator( + # tosa.GATHER(x=[P,K,C], indices=[P,W]) -> [P,W,C] + gathered_pwc = super().call_operator( tosa_gather_op, - (x_1kc, lin_1w), + (x_pkc, lin_pw), {}, meta, updated=True, ) - # ---- output: [*S, *trailing] ---- - out_shape = list(S) + list(trailing) + # ---- output: [*leading, *S, *trailing] ---- + out_shape = list(leading) + list(S) + list(trailing) out = super().call_operator( view_op, - (gathered_1wc, out_shape), + (gathered_pwc, out_shape), {}, meta, updated=True, diff --git a/backends/arm/_passes/decompose_isinf_isnan_pass.py b/backends/arm/_passes/decompose_isinf_isnan_pass.py new file mode 100644 index 00000000000..7722e41ea40 --- /dev/null +++ b/backends/arm/_passes/decompose_isinf_isnan_pass.py @@ -0,0 +1,48 @@ +# 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 Set, Type + +from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass + + +class DecomposeIsInfAndIsNanPass(ArmOpTargetedPass): + """Decompose ``isinf`` and ``isnan`` into TOSA-supported operations.""" + + _passes_required_after: Set[Type[ExportPass]] = set() + edge_isinf = exir_ops.edge.aten.isinf.default + edge_isnan = exir_ops.edge.aten.isnan.default + target_ops = (edge_isinf, edge_isnan) + check_allowed_to_transform = True + + def call_operator(self, op, args, kwargs, meta): + if op not in self.target_ops or not self.allowed_to_transform(meta): + return super().call_operator(op, args, kwargs, meta) + + (x,) = args + abs_op = exir_ops.edge.aten.abs.default + eq_op = exir_ops.edge.aten.eq.Tensor + logical_not_op = exir_ops.edge.aten.logical_not.default + full_op = exir_ops.edge.aten.full.default + + if op is self.edge_isnan: + equal = super().call_operator(eq_op, (x, x), {}, meta, updated=True) + return super().call_operator( + logical_not_op, (equal,), {}, meta, updated=True + ) + + absolute = super().call_operator(abs_op, (x,), {}, meta, updated=True) + infinity = super().call_operator( + full_op, + (x.data.shape, float("inf")), + {"dtype": x.data.dtype}, + meta, + updated=True, + ) + return super().call_operator( + eq_op, (absolute, infinity), {}, meta, updated=True + ) diff --git a/backends/arm/_passes/decompose_roll_pass.py b/backends/arm/_passes/decompose_roll_pass.py new file mode 100644 index 00000000000..060732d153e --- /dev/null +++ b/backends/arm/_passes/decompose_roll_pass.py @@ -0,0 +1,131 @@ +# 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 collections.abc import Sequence +from typing import Set, Type + +from executorch.backends.arm._passes import ArmOpTargetedPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass + + +RollParameters = tuple[tuple[int, int, int], ...] + + +def get_static_roll_parameters( + input_shape: Sequence[object], shifts: object, dims: object +) -> RollParameters | None: + """Normalize a statically-decomposable roll. + + Args: + input_shape (Sequence[object]): Shape of the roll input. + shifts (object): Roll shifts from the operator arguments. + dims (object): Roll dimensions from the operator arguments. + + Returns: + RollParameters | None: Normalized ``(shift, dim, size)`` tuples, or + ``None`` when the roll cannot be decomposed statically. + + """ + if not input_shape or any( + type(size) is not int or size <= 0 for size in input_shape + ): + return None + if not isinstance(shifts, (list, tuple)) or not isinstance(dims, (list, tuple)): + return None + if not shifts or len(shifts) != len(dims): + return None + if any(type(value) is not int for value in (*shifts, *dims)): + return None + + rank = len(input_shape) + parameters: list[tuple[int, int, int]] = [] + for shift, dim in zip(shifts, dims): + if not -rank <= dim < rank: + return None + normalized_dim = dim % rank + dim_size = int(input_shape[normalized_dim]) + parameters.append((shift % dim_size, normalized_dim, dim_size)) + return tuple(parameters) + + +def can_decompose_roll( + input_shape: Sequence[object], shifts: object, dims: object +) -> bool: + """Return whether a roll can become a nonempty slice/concat graph. + + Args: + input_shape (Sequence[object]): Shape of the roll input. + shifts (object): Roll shifts from the operator arguments. + dims (object): Roll dimensions from the operator arguments. + + Returns: + bool: True when the roll has a supported static decomposition. + + """ + parameters = get_static_roll_parameters(input_shape, shifts, dims) + return parameters is not None and any(shift != 0 for shift, _, _ in parameters) + + +class DecomposeRollPass(ArmOpTargetedPass): + """Decompose a static ``aten.roll`` into slices and concatenation. + + For each nonzero ``(shift, dim)`` pair, normalize the shift and apply: + + shift = shift % dim_size + result = cat( + ( + slice_copy(result, dim, dim_size - shift, dim_size), + slice_copy(result, dim, 0, dim_size - shift), + ), + dim, + ) + + Rewrites are applied sequentially to support multiple and repeated + dimensions. + + """ + + _passes_required_after: Set[Type[ExportPass]] = set() + target_ops = {exir_ops.edge.aten.roll.default} + + def call_operator(self, op, args, kwargs, meta, updated=False): + if op not in self.target_ops: + return super().call_operator(op, args, kwargs, meta, updated) + + input_tensor = args[0] + shifts = args[1] + dims = args[2] if len(args) > 2 else () + parameters = get_static_roll_parameters(input_tensor.data.shape, shifts, dims) + if parameters is None or not any(shift != 0 for shift, _, _ in parameters): + raise ValueError("Expected a nonempty static roll decomposition") + + result = input_tensor + for shift, dim, dim_size in parameters: + if shift == 0: + continue + split = dim_size - shift + suffix = super().call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + (result, dim, split, dim_size, 1), + {}, + meta, + updated=True, + ) + prefix = super().call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + (result, dim, 0, split, 1), + {}, + meta, + updated=True, + ) + result = super().call_operator( + exir_ops.edge.aten.cat.default, + ([suffix, prefix], dim), + {}, + meta, + updated=True, + ) + return result diff --git a/backends/arm/_passes/decompose_round_pass.py b/backends/arm/_passes/decompose_round_pass.py index 48b26f1d027..dcb72d4202c 100644 --- a/backends/arm/_passes/decompose_round_pass.py +++ b/backends/arm/_passes/decompose_round_pass.py @@ -7,47 +7,25 @@ from executorch.backends.arm._passes import ArmOpTargetedPass 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 -from torch._ops import OpOverload -Op = OpOverload | EdgeOpOverload - - -def _get_round_decomposition_ops(op) -> tuple[Op, Op, Op, Op, Op, Op, Op]: - """Returns the (full_op, ge_op, add_op, sub_op, floor_op, ceil_op, where_op) - for the given round operation. - - The ops depend on whether the round op is an aten or edge op. - - """ - if op == exir_ops.edge.aten.round.default: - return ( - exir_ops.edge.aten.full.default, - exir_ops.edge.aten.ge.Tensor, - exir_ops.edge.aten.add.Scalar, - exir_ops.edge.aten.sub.Scalar, - exir_ops.edge.aten.floor.default, - exir_ops.edge.aten.ceil.default, - exir_ops.edge.aten.where.self, - ) - raise RuntimeError(f"Can't get round decomposition ops for op {op}") +class DecomposeRoundPass(ArmOpTargetedPass): + """Decomposes round(x) into round-half-to-even, matching the semantics of + aten.round / torch.round. + x lies between floor(x) and ceil(x), and its distance above floor(x) says + which one is nearer: less than 0.5 takes floor(x), more takes ceil(x), and + exactly 0.5 is a tie that takes whichever of the two is even. -class DecomposeRoundPass(ArmOpTargetedPass): - """ - For inputs >= 0, round(x) is equivalent to floor(x + 0.5), and for inputs < 0, - round(x) is equivalent to ceil(x - 0.5). This pass decomposes the round operation into - a sequence of more primitive operations. Example: - %zero = full((1,), 0.0, dtype=torch.float32) - %is_non_negative = ge(x, %zero) - %plus_half = add(x, 0.5) - %minus_half = sub(x, 0.5) - %floor = floor(%plus_half) - %ceil = ceil(%minus_half) - %result = where(%is_non_negative, %floor, %ceil) + %dist_to_floor = sub(x, floor(x)) + %halved = mul(floor(x), 0.5) + %floor_is_odd = eq(sub(%halved, floor(%halved)), 0.5) + %tie_to_even = logical_and(eq(%dist_to_floor, 0.5), %floor_is_odd) + %take_ceil = logical_or(gt(%dist_to_floor, 0.5), %tie_to_even) + %result = where(%take_ceil, ceil(x), floor(x)) + """ _passes_required_after: Set[Type[ExportPass]] = set() @@ -60,26 +38,30 @@ def call_operator(self, op, args, kwargs, meta, updated=False): if op not in self.target_ops or self._is_quantized_meta(meta): return super().call_operator(op, args, kwargs, meta, updated) x = args[0] - input_dtype = x.node.meta["val"].dtype - full, ge, add, sub, floor, ceil, where = _get_round_decomposition_ops(op) - zero = super().call_operator( - full, - args=((1,), 0.0), - kwargs={"dtype": input_dtype}, - meta=meta, - updated=True, - ) - is_non_negative = super().call_operator( - ge, (x, zero), kwargs, meta, updated=True - ) - plus_half = super().call_operator(add, (x, 0.5), kwargs, meta, updated=True) - minus_half = super().call_operator(sub, (x, 0.5), kwargs, meta, updated=True) - floor = super().call_operator(floor, (plus_half,), kwargs, meta, updated=True) - ceil = super().call_operator(ceil, (minus_half,), kwargs, meta, updated=True) - return super().call_operator( - where, - (is_non_negative, floor, ceil), - kwargs, - meta, - updated=True, - ) + + def call(op, *op_args): + return super(DecomposeRoundPass, self).call_operator( + op, op_args, kwargs, meta, updated=True + ) + + sub = exir_ops.edge.aten.sub.Tensor + mul = exir_ops.edge.aten.mul.Scalar + floor = exir_ops.edge.aten.floor.default + ceil = exir_ops.edge.aten.ceil.default + eq = exir_ops.edge.aten.eq.Scalar + gt = exir_ops.edge.aten.gt.Scalar + logical_and = exir_ops.edge.aten.logical_and.default + logical_or = exir_ops.edge.aten.logical_or.default + where = exir_ops.edge.aten.where.self + + floor_x = call(floor, x) + dist_to_floor = call(sub, x, floor_x) + + # floor_x is odd iff floor_x / 2 has a .5 fractional part + halved = call(mul, floor_x, 0.5) + halved_frac = call(sub, halved, call(floor, halved)) + floor_is_odd = call(eq, halved_frac, 0.5) + + tie_to_even = call(logical_and, call(eq, dist_to_floor, 0.5), floor_is_odd) + take_ceil = call(logical_or, call(gt, dist_to_floor, 0.5), tie_to_even) + return call(where, take_ceil, call(ceil, x), floor_x) diff --git a/backends/arm/_passes/decompose_sdpa_with_regular_softmax_pass.py b/backends/arm/_passes/decompose_sdpa_with_regular_softmax_pass.py new file mode 100644 index 00000000000..9d2e6da7fa9 --- /dev/null +++ b/backends/arm/_passes/decompose_sdpa_with_regular_softmax_pass.py @@ -0,0 +1,111 @@ +# 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 Set, Type + +import torch +from executorch.backends.arm._passes.arm_pass import ArmPass +from executorch.backends.transforms import decompose_sdpa +from executorch.exir.pass_base import ExportPass, PassResult + + +class DecomposeSDPAWithRegularSoftmaxPass( + ArmPass, decompose_sdpa.DecomposeScaledDotProductAttention +): + """Decompose eligible SDPA calls using regular softmax. + + Matches unmasked, noncausal, zero-dropout SDPA calls whose key sequence + length is statically known to be nonzero. The matched form is conceptually:: + + scaled_dot_product_attention( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + scale=scale, + enable_gqa=enable_gqa, + ) + + Other SDPA options, such as ``scale`` and ``enable_gqa``, are preserved. + + The generated subgraph is approximately:: + + scores = (query @ key.transpose(-2, -1)) * scale + output = softmax(scores, dim=-1) @ value + + The standard SDPA decomposition initially generates ``_safe_softmax``. + This pass replaces that operator with regular ``softmax`` only in the + newly generated subgraph. The eligibility checks prevent masks or empty + key sequences from producing all-negative-infinity score rows. They assume + such rows are not produced by nonfinite inputs or numerical overflow. + + """ + + _passes_required_after: Set[Type[ExportPass]] = set() + + def call( + self, graph_module: torch.fx.GraphModule, allow_non_fake_inputs: bool = True + ) -> PassResult: + graph = graph_module.graph + modified = False + for node in list(graph.nodes): + if node.target != torch.ops.aten.scaled_dot_product_attention.default: + continue + if not self._is_auto_guard_removal_candidate(node): + continue + + existing_nodes = set(graph.nodes) + super()._decompose_sdpa_node(graph_module, node, allow_non_fake_inputs) + self._remove_safe_softmax_guard(graph, existing_nodes) + modified = True + + if modified: + graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) + + @classmethod + def _is_auto_guard_removal_candidate(cls, node: torch.fx.Node) -> bool: + """Return true when SDPA meets automatic removal constraints. + + These structural checks exclude fully masked rows. They assume scores + do not become all ``-inf`` through nonfinite inputs or overflow. + + """ + canonical_args, _, _ = cls._canonicalize_sdpa_call(node) + _, key, _, attn_mask, dropout_p, is_causal, _ = canonical_args + + if attn_mask is not None: + return False + if cls._extract_arg_value(is_causal) is not False: + return False + if cls._extract_arg_value(dropout_p) != 0.0: + return False + return cls._has_nonzero_key_sequence_length(key) + + @staticmethod + def _has_nonzero_key_sequence_length(key: object) -> bool: + if not isinstance(key, torch.fx.Node): + return False + + val = key.meta.get("val") + shape = getattr(val, "shape", None) + if shape is None or len(shape) < 2: + return False + + key_sequence_length = shape[-2] + return isinstance(key_sequence_length, int) and key_sequence_length > 0 + + @staticmethod + def _remove_safe_softmax_guard( + graph: torch.fx.Graph, existing_nodes: set[torch.fx.Node] + ) -> None: + for decomposed_node in graph.nodes: + if decomposed_node in existing_nodes: + continue + if decomposed_node.target == torch.ops.aten._safe_softmax.default: + decomposed_node.target = torch.ops.aten.softmax.int diff --git a/backends/arm/_passes/deduplicate_const_shapes_pass.py b/backends/arm/_passes/deduplicate_const_shapes_pass.py new file mode 100644 index 00000000000..61a0083c429 --- /dev/null +++ b/backends/arm/_passes/deduplicate_const_shapes_pass.py @@ -0,0 +1,47 @@ +# 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 Set, Type + +from executorch.backends.arm._passes import ArmPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx import GraphModule, Node + + +class DeduplicateConstShapesPass(ArmPass): + """Reuse the first CONST_SHAPE node with identical static values.""" + + _passes_required_after: Set[Type[ExportPass]] = set() + + def call(self, graph_module: GraphModule) -> PassResult: + representatives: dict[tuple[int, ...], Node] = {} + modified = False + + for node in list(graph_module.graph.nodes): + if node.target != exir_ops.backend.tosa.CONST_SHAPE.default: + continue + + values = node.args[0] + if not isinstance(values, (list, tuple)) or not all( + type(value) is int for value in values + ): + continue + + key = tuple(values) + representative = representatives.get(key) + if representative is None: + representatives[key] = node + continue + + node.replace_all_uses_with(representative) + graph_module.graph.erase_node(node) + modified = True + + if modified: + graph_module.graph.lint() + graph_module.recompile() + + return PassResult(graph_module, modified) diff --git a/backends/arm/_passes/exir_to_tosa_pass.py b/backends/arm/_passes/exir_to_tosa_pass.py index e03b880dfd5..b0d6b11f857 100644 --- a/backends/arm/_passes/exir_to_tosa_pass.py +++ b/backends/arm/_passes/exir_to_tosa_pass.py @@ -9,6 +9,9 @@ from executorch.backends.arm._passes.aten_to_tosa_activation_functions import ( get_activation_replacement, ) +from executorch.backends.arm._passes.aten_to_tosa_comparison import ( + rewrite_comparison_operator, +) from executorch.backends.arm._passes.aten_to_tosa_data_layout import ( rewrite_data_layout_operator, ) @@ -79,9 +82,6 @@ def _get_fft_replacement( exir_ops.edge.aten.bitwise_or.Tensor, exir_ops.edge.aten.bitwise_right_shift.Tensor, exir_ops.edge.aten.bitwise_xor.Tensor, - exir_ops.edge.aten.eq.Tensor, - exir_ops.edge.aten.ge.Tensor, - exir_ops.edge.aten.gt.Tensor, exir_ops.edge.aten.logical_and.default, exir_ops.edge.aten.logical_or.default, exir_ops.edge.aten.logical_xor.default, @@ -97,6 +97,17 @@ def _get_binary_operator_replacement( return rewrite_binary_operator(node, pass_) +@register_dialect_substitutions( + exir_ops.edge.aten.eq.Tensor, + exir_ops.edge.aten.ge.Tensor, + exir_ops.edge.aten.gt.Tensor, +) +def _get_comparison_operator_replacement( + node: Node, pass_: AtenToDialectPass +) -> DialectNodeSpec | None: + return rewrite_comparison_operator(node, pass_) + + @register_dialect_substitutions( exir_ops.edge.aten.abs.default, exir_ops.edge.aten.bitwise_not.default, diff --git a/backends/arm/_passes/fuse_batch_norm2d_pass.py b/backends/arm/_passes/fuse_batch_norm2d_pass.py index a13ed9da922..69cd52c40b9 100644 --- a/backends/arm/_passes/fuse_batch_norm2d_pass.py +++ b/backends/arm/_passes/fuse_batch_norm2d_pass.py @@ -12,6 +12,9 @@ create_node, get_first_fake_tensor, ) +from executorch.backends.arm._passes.decompose_grouped_conv_pass import ( + DecomposeGroupedConvPass, +) from executorch.backends.arm.common.debug import get_node_debug_info from executorch.backends.transforms.utils import ( create_constant_placeholder, @@ -27,11 +30,13 @@ class FuseBatchNorm2dPass(ArmPass): - """Fuses the pattern convolution -> batchnorm by updating the weights and - bias of the convolution and removing the batchnorm. + """Fuse convolution followed by BatchNorm. + + Update the convolution weights and bias and remove the BatchNorm operation. + """ - _passes_required_after: Set[Type[ExportPass]] = set() + _passes_required_after: Set[Type[ExportPass]] = {DecomposeGroupedConvPass} def __init__(self, exported_program: ExportedProgram, *args, **kwargs): super().__init__(*args, **kwargs) @@ -45,6 +50,79 @@ def get_bias_name(self, weight_node: Node, bias_node: Node | None) -> str: else: return weight_node.name + "_bias_fused_bn" + @staticmethod + def _fuse_grouped_transposed_conv_bn_weights( + conv_weight: torch.Tensor, + conv_bias: torch.Tensor | None, + bn_mean: torch.Tensor, + bn_var: torch.Tensor, + bn_epsilon: float, + bn_weight: torch.Tensor | None, + bn_bias: torch.Tensor | None, + groups: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse BatchNorm into grouped transposed-convolution parameters. + + This helper runs before ``DecomposeGroupedConvPass`` and transforms:: + + grouped ConvTranspose -> BatchNorm + + into a grouped ConvTranspose with fused weights and bias. A transposed + convolution weight has layout ``[Cin, Cout/groups, ...]``. The weight + is split on its input-channel dimension, while the bias and BatchNorm + parameters are split on their output-channel dimension. Each group is + fused independently before the original grouped layout is restored. + + Args: + conv_weight (torch.Tensor): Grouped transposed-convolution weight. + conv_bias (torch.Tensor | None): Convolution bias. + bn_mean (torch.Tensor): BatchNorm running mean. + bn_var (torch.Tensor): BatchNorm running variance. + bn_epsilon (float): BatchNorm numerical-stability constant. + bn_weight (torch.Tensor | None): BatchNorm weight. + bn_bias (torch.Tensor | None): BatchNorm bias. + groups (int): Number of convolution groups. + + Returns: + tuple[torch.Tensor, torch.Tensor]: Fused weight and bias in the + original grouped layout. + + Raises: + RuntimeError: If the grouped channel dimensions are inconsistent. + + """ + if conv_weight.size(0) % groups != 0 or bn_mean.numel() % groups != 0: + raise RuntimeError("Grouped transposed convolution has invalid channels") + + input_channels_per_group = conv_weight.size(0) // groups + output_channels_per_group = bn_mean.numel() // groups + if conv_weight.size(1) != output_channels_per_group: + raise RuntimeError("BatchNorm channels do not match convolution output") + + fused_weights: list[torch.Tensor] = [] + fused_biases: list[torch.Tensor] = [] + for group in range(groups): + input_start = group * input_channels_per_group + input_end = input_start + input_channels_per_group + output_start = group * output_channels_per_group + output_end = output_start + output_channels_per_group + output_slice = slice(output_start, output_end) + + fused_weight, fused_bias = fuse_conv_bn_weights( + conv_weight[input_start:input_end], + conv_bias[output_slice] if conv_bias is not None else None, + bn_mean[output_slice], + bn_var[output_slice], + bn_epsilon, + bn_weight[output_slice] if bn_weight is not None else None, + bn_bias[output_slice] if bn_bias is not None else None, + transpose=True, + ) + fused_weights.append(fused_weight) + fused_biases.append(fused_bias) + + return torch.cat(fused_weights, dim=0), torch.cat(fused_biases, dim=0) + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 modified = False constant_placeholders_to_delete = set() @@ -176,15 +254,32 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 ) # Fuse bn weights/bias with input weights/bias - fused_weight, fused_bias = fuse_conv_bn_weights( - input_weight_tensor, - input_bias_tensor, - bn_mean_tensor, - bn_var_tensor, - epsilon, - bn_weight_tensor, - bn_bias_tensor, - ) + transposed = bool(input_node.args[6]) + groups = int(input_node.args[8]) + if transposed and groups > 1: + fused_weight, fused_bias = ( + self._fuse_grouped_transposed_conv_bn_weights( + input_weight_tensor, + input_bias_tensor, + bn_mean_tensor, + bn_var_tensor, + epsilon, + bn_weight_tensor, + bn_bias_tensor, + groups, + ) + ) + else: + fused_weight, fused_bias = fuse_conv_bn_weights( + input_weight_tensor, + input_bias_tensor, + bn_mean_tensor, + bn_var_tensor, + epsilon, + bn_weight_tensor, + bn_bias_tensor, + transpose=transposed, + ) # Create fused weights and bias to conv and replace conv args with graph_module.graph.inserting_before(input_weight_node): diff --git a/backends/arm/_passes/insert_rescales_pass.py b/backends/arm/_passes/insert_rescales_pass.py index 2798d5d17f1..76c1c2995f8 100644 --- a/backends/arm/_passes/insert_rescales_pass.py +++ b/backends/arm/_passes/insert_rescales_pass.py @@ -408,6 +408,12 @@ def call(self, graph_module: GraphModule) -> PassResult: if node.op != "call_function" or node.target not in self.included_targets: continue + has_preserved_qdq = any( + input_node.target in DQ_OPS for input_node in node.all_input_nodes + ) or any(user.target in Q_OPS for user in node.users) + if has_preserved_qdq: + continue + if "input_qparams" not in node.meta or len(node.meta["input_qparams"]) == 0: continue input_qparams = node.meta["input_qparams"] diff --git a/backends/arm/_passes/normalize_max_pool2d_input_rank_pass.py b/backends/arm/_passes/normalize_max_pool2d_input_rank_pass.py new file mode 100644 index 00000000000..98d417f6f56 --- /dev/null +++ b/backends/arm/_passes/normalize_max_pool2d_input_rank_pass.py @@ -0,0 +1,97 @@ +# 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 Set, Type + +import torch +from executorch.backends.arm._passes import ArmOpTargetedPass +from executorch.backends.arm._passes.arm_pass_utils import ( + create_node, + get_first_fake_tensor, +) +from executorch.backends.arm._passes.convert_squeezes_to_view import ( + ConvertSqueezesToViewPass, +) +from executorch.backends.arm._passes.decompose_maxpool2d_with_dilation_pass import ( + DecomposeMaxPool2dPass, +) +from executorch.backends.arm._passes.rewrite_max_pool2d_pass import RewriteMaxPool2dPass +from executorch.backends.arm._passes.size_adjust_input_pass import SizeAdjustInputPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx import Node + + +class NormalizeMaxPool2dInputRankPass(ArmOpTargetedPass): + """Normalize unbatched rank-3 max_pool2d inputs to rank 4. + + Unsqueeze inputs from ``[C, H, W]`` to ``[1, C, H, W]``. + + Squeeze the leading dimension after pooling to restore the rank-3 output. + + The complete shape transformation is:: + + [C, H, W] + -> unsqueeze(0) -> [1, C, H, W] + -> max_pool2d -> [1, C, H_out, W_out] + -> squeeze(0) -> [C, H_out, W_out] + + """ + + target_ops = (exir_ops.edge.aten.max_pool2d.default,) + _passes_required_after: Set[Type[ExportPass]] = { + ConvertSqueezesToViewPass, + DecomposeMaxPool2dPass, + RewriteMaxPool2dPass, + SizeAdjustInputPass, + } + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + graph = graph_module.graph + modified = False + pool_nodes = graph.find_nodes(op="call_function", target=self.target_ops[0]) + + for pool_node in pool_nodes: + input_node = pool_node.args[0] + if not isinstance(input_node, Node): + raise RuntimeError("Expected max_pool2d input to be a node") + + input_fake = get_first_fake_tensor(input_node) + if input_fake.dim() != 3: + continue + + output_fake = get_first_fake_tensor(pool_node) + with graph.inserting_before(pool_node): + unsqueeze = create_node( + graph, + exir_ops.edge.aten.unsqueeze_copy.default, + args=(input_node, 0), + from_node=pool_node, + inherit_qparams=False, + ) + unsqueeze.meta["val"] = input_fake.unsqueeze(0) + pool_node.replace_input_with(input_node, unsqueeze) + + pool_node.meta["val"] = output_fake.unsqueeze(0) + original_users = list(pool_node.users) + with graph.inserting_after(pool_node): + squeeze = create_node( + graph, + exir_ops.edge.aten.squeeze_copy.dims, + args=(pool_node, [0]), + from_node=pool_node, + inherit_qparams=False, + ) + squeeze.meta["val"] = output_fake + for user in original_users: + user.replace_input_with(pool_node, squeeze) + + modified = True + + if modified: + graph.lint() + graph_module.recompile() + + return PassResult(graph_module, modified) diff --git a/backends/arm/_passes/rewrite_conv_pass.py b/backends/arm/_passes/rewrite_conv_pass.py index 69d4d76c9cf..1ed60228582 100644 --- a/backends/arm/_passes/rewrite_conv_pass.py +++ b/backends/arm/_passes/rewrite_conv_pass.py @@ -37,11 +37,20 @@ TOSA_CONTROL_FLOW_SOURCE_NODE_META, TosaSpecialDtype, ) -from executorch.backends.arm.tosa.specification import get_context_shape_env +from executorch.backends.arm.tosa.specification import ( + get_context_shape_env, + get_context_spec, +) +from executorch.backends.transforms.fuse_duplicate_users_pass import ( + build_node_signature, + DO_NOT_FUSE_DUPLICATE_META_KEY, +) from executorch.backends.transforms.utils import create_constant_placeholder 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 torch._ops import OpOverload from torch._subclasses.fake_tensor import FakeTensor from torch.export.graph_signature import InputKind @@ -94,7 +103,7 @@ def _adjust_pad_if_needed( pass instead. """ - mod_remainder = ( + mod_remainder: int | torch.SymInt = ( input_len + 2 * pad - dilation * (input_weight - 1) - 1 ) % stride @@ -121,14 +130,14 @@ def _adjust_pad_if_needed( return pad - mod_remainder - def _is_depthwise_conv2d(self, node: torch.fx.Node) -> bool: + def _is_depthwise_conv(self, node: torch.fx.Node) -> bool: if ( node.op != "call_function" or node.target != exir_ops.edge.aten.convolution.default ): return False input_tensor = get_first_fake_tensor(node.all_input_nodes[0]) - if len(input_tensor.shape) != 4: + if len(input_tensor.shape) not in (3, 4): return False groups = node.args[-1] in_channels = input_tensor.shape[1] @@ -524,11 +533,11 @@ def _combine_rescale_scales( @staticmethod def _is_direct_int32_rescale(node: torch.fx.Node) -> bool: """Return whether a node directly rescales its input to INT32.""" - return ( + return bool( node.op == "call_function" and node.target == exir_ops.backend.tosa.RESCALE.default and len(node.args) > 1 - and node.args[1] == torch.int32 + and node.args[1] is torch.int32 ) def _get_direct_int32_rescale_users( @@ -571,6 +580,88 @@ def _insert_layout_permute( output.meta["val"] = output_fake_tensor return output, output_fake_tensor + @classmethod + def _deduplicate_a16w8_output_rescales( + cls, + graph_module: torch.fx.GraphModule, + tosa_op: torch.fx.Node, + node_order: dict[torch.fx.Node, int], + ) -> list[torch.fx.Node] | None: + """Merge only complete, canonical RESCALE-to-PERMUTE heads.""" + if any(user not in node_order for user in tosa_op.users): + return None + rescale_users = sorted(tosa_op.users, key=node_order.__getitem__) + if any( + user.target != exir_ops.backend.tosa.RESCALE.default + for user in rescale_users + ): + # RewriteConvPass creates only RESCALE users for this accumulator; + # preserve an unfamiliar future shape instead of partially rewriting it. + return None + + unique_rescales: dict[tuple[Any, ...], torch.fx.Node] = {} + deduplicated_rescales: list[torch.fx.Node] = [] + for rescale in rescale_users: + rescale_outputs = list(rescale.users) + if ( + len(rescale_outputs) != 1 + or rescale_outputs[0].target != exir_ops.edge.aten.permute_copy.default + ): + deduplicated_rescales.append(rescale) + continue + layout_permute = rescale_outputs[0] + rescale_signature = build_node_signature(rescale, positional_arg_start=1) + permute_signature = build_node_signature( + layout_permute, positional_arg_start=1 + ) + if rescale_signature is None or permute_signature is None: + deduplicated_rescales.append(rescale) + continue + signature = ( + rescale_signature, + permute_signature, + ) + canonical_permute = unique_rescales.get(signature) + if canonical_permute is not None: + # Layout permutes are inserted directly after their RESCALE, + # so the earliest RESCALE also provides a dominating permute. + layout_permute.replace_all_uses_with(canonical_permute) + graph_module.graph.erase_node(layout_permute) + graph_module.graph.erase_node(rescale) + else: + unique_rescales[signature] = layout_permute + deduplicated_rescales.append(rescale) + + return deduplicated_rescales + + def _separate_u55_a16w8_output_rescales( + self, + graph_module: torch.fx.GraphModule, + tosa_op: torch.fx.Node, + node_order: dict[torch.fx.Node, int], + ) -> None: + if len(tosa_op.users) < 2: + return + + rescale_users = self._deduplicate_a16w8_output_rescales( + graph_module, tosa_op, node_order + ) + if rescale_users is None or len(rescale_users) < 2: + return + tosa_op.meta[DO_NOT_FUSE_DUPLICATE_META_KEY] = True + for rescale in rescale_users[1:]: + with graph_module.graph.inserting_before(rescale): + cloned_tosa_op = create_node( + graph=graph_module.graph, + op_target=cast(OpOverload | EdgeOpOverload, tosa_op.target), + args=tosa_op.args, + kwargs=tosa_op.kwargs, + from_node=tosa_op, + inherit_qparams=True, + ) + cloned_tosa_op.meta[DO_NOT_FUSE_DUPLICATE_META_KEY] = True + rescale.replace_input_with(tosa_op, cloned_tosa_op) + def _insert_a16w8_output_branches( self, graph_module: torch.fx.GraphModule, @@ -787,6 +878,7 @@ def _insert_output_conversion( def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 modified = False + a16w8_tosa_ops: list[torch.fx.Node] = [] for node in graph_module.graph.nodes: if ( node.op != "call_function" @@ -911,7 +1003,70 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 dilation = tuple(dilation_list) pad = pad_attr - if self._is_conv3d(len(input_shape), group): + if spatial_rank == 1: + target_op = ( + exir_ops.backend.tosa.DEPTHWISE_CONV2D.default + if self._is_depthwise_conv(node) + else exir_ops.backend.tosa.CONV2D.default + ) + pre_permute_dims = (0, 2, 1) + post_permute_dims = (0, 2, 1) + with graph_module.graph.inserting_before(node): + x = create_node( + graph=graph_module.graph, + op_target=exir_ops.edge.aten.permute_copy.default, + args=(x, list(pre_permute_dims)), + from_node=node, + ) + permuted_input_fake = permute_fake_tensor_metadata( + input_fake_tensor, pre_permute_dims + ) + x.meta["val"] = permuted_input_fake + input_tensor_for_tosa_fake = permuted_input_fake.unsqueeze(1) + x = create_node( + graph=graph_module.graph, + op_target=exir_ops.edge.aten.view_copy.default, + args=(x, list(input_tensor_for_tosa_fake.shape)), + from_node=node, + ) + x.meta["val"] = input_tensor_for_tosa_fake + + kernel_width = weight_shape[2] + if target_op == exir_ops.backend.tosa.DEPTHWISE_CONV2D.default: + in_channels = input_fake_tensor.shape[1] + channel_multiplier = weight_shape[0] // in_channels + weight = self._rewrite_weight( + graph_module, + weight, + node, + permute_dims=(1, 2, 0), + name_suffix="hwicm", + reshape_dims=( + 1, + kernel_width, + in_channels, + channel_multiplier, + ), + ) + else: + weight = self._rewrite_weight( + graph_module, + weight, + node, + permute_dims=(0, 2, 1), + name_suffix="ohwi", + reshape_dims=( + weight_shape[0], + 1, + kernel_width, + weight_shape[1], + ), + ) + weight_fake_tensor = get_first_fake_tensor(weight) + stride = (1, stride[0]) + dilation = (1, dilation[0]) + pad = [0, 0, pad[0], pad[1]] + elif self._is_conv3d(len(input_shape), group): target_op = exir_ops.backend.tosa.CONV3D.default pre_permute_dims = ODHWI_ORDER post_permute_dims = ODHWI_INVERSE_ORDER @@ -934,7 +1089,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 name_suffix="odhwi", ) weight_fake_tensor = get_first_fake_tensor(weight) - elif self._is_depthwise_conv2d(node): + elif self._is_depthwise_conv(node): target_op = exir_ops.backend.tosa.DEPTHWISE_CONV2D.default pre_permute_dims = NHWC_ORDER post_permute_dims = NHWC_INVERSE_ORDER @@ -1039,7 +1194,28 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 if post_permute_dims is None: raise RuntimeError("Expected post permute dims for explicit layout") + output_conversion_node = node_replacement post_permute_input = node_replacement + squeeze_view: torch.fx.Node | None = None + if spatial_rank == 1: + squeezed_output_fake = cast( + FakeTensor, node_replacement_fake_tensor.squeeze(1) + ) + special_dtype = node_replacement.meta.get(TosaSpecialDtype.meta_key()) + with graph_module.graph.inserting_after(node_replacement): + node_replacement = create_node( + graph=graph_module.graph, + op_target=exir_ops.edge.aten.view_copy.default, + args=(node_replacement, list(squeezed_output_fake.shape)), + from_node=node, + ) + node_replacement.meta["val"] = squeezed_output_fake + if special_dtype: + node_replacement.meta[TosaSpecialDtype.meta_key()] = special_dtype + squeeze_view = node_replacement + post_permute_input = node_replacement + node_replacement_fake_tensor = squeezed_output_fake + with graph_module.graph.inserting_after(node_replacement): node_replacement = create_node( graph=graph_module.graph, @@ -1059,16 +1235,23 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 tosa_node_fake_tensor.dtype == torch.int32 and input_fake_tensor.dtype == torch.int16 ) - if is_a16w8_conv: + if is_a16w8_conv and spatial_rank != 1: # Keep values in INT32 whenever a consumer supports it, even # though the declared output is INT16, by branching from the # accumulator before narrowing. + # + # Rank-three convolutions are excluded. The legacy Conv1d + # expansion placed a rank-changing view between the convolution + # and its INT32 consumers, so the convolution narrowed to its + # exported output domain instead of forking. Forking here would + # give each branch its own boundary rescale and permute, which + # Vela materialises as a second full transpose of the output. self._insert_a16w8_output_branches( graph_module, node, tosa_op, tosa_node_fake_tensor, - post_permute_input, + output_conversion_node, post_permute_dims, ) # Only users not moved to widened branches remain on the @@ -1078,12 +1261,24 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 node.replace_all_uses_with(node_replacement) else: graph_module.graph.erase_node(node_replacement) - graph_module.graph.erase_node(post_permute_input) + if squeeze_view is not None: + graph_module.graph.erase_node(squeeze_view) + graph_module.graph.erase_node(output_conversion_node) + a16w8_tosa_ops.append(tosa_op) else: node.replace_all_uses_with(node_replacement) graph_module.graph.erase_node(node) + if a16w8_tosa_ops and get_context_spec().is_U55_subset: + node_order = { + node: index for index, node in enumerate(graph_module.graph.nodes) + } + for tosa_op in a16w8_tosa_ops: + self._separate_u55_a16w8_output_rescales( + graph_module, tosa_op, node_order + ) + if modified: graph_module.recompile() graph_module = super().call(graph_module).graph_module diff --git a/backends/arm/_passes/symbolic_value_range.py b/backends/arm/_passes/symbolic_value_range.py index 609a84edc54..36bff1ed47a 100644 --- a/backends/arm/_passes/symbolic_value_range.py +++ b/backends/arm/_passes/symbolic_value_range.py @@ -145,12 +145,20 @@ def mod(lhs: _ExactValues, rhs: _ExactValues) -> _ExactValues: return None return _combine_values(lhs, rhs, lambda a, b: sympy.Mod(a, b)) + @staticmethod + def python_mod(lhs: _ExactValues, rhs: _ExactValues) -> _ExactValues: + return _ExactValueAnalysis.mod(lhs, rhs) + @staticmethod def floordiv(lhs: _ExactValues, rhs: _ExactValues) -> _ExactValues: if rhs is None or any(value == 0 for value in rhs): return None return _combine_values(lhs, rhs, lambda a, b: sympy.floor(a / b)) + @staticmethod + def python_floordiv(lhs: _ExactValues, rhs: _ExactValues) -> _ExactValues: + return _ExactValueAnalysis.floordiv(lhs, rhs) + @staticmethod def pow(lhs: _ExactValues, rhs: _ExactValues) -> _ExactValues: return _combine_values(lhs, rhs, lambda a, b: a**b) @@ -181,7 +189,6 @@ def evaluate_symbolic_expr_values( """ root_expr = expr.node.expr if isinstance(expr, torch.SymInt) else expr - constant_values = _constant_expr_values(root_expr) if constant_values is not None: return constant_values diff --git a/backends/arm/cmake/ArmRunnerUtilsInternal.cmake b/backends/arm/cmake/ArmRunnerUtilsInternal.cmake index 04c979b04ce..d7d66c45450 100644 --- a/backends/arm/cmake/ArmRunnerUtilsInternal.cmake +++ b/backends/arm/cmake/ArmRunnerUtilsInternal.cmake @@ -328,8 +328,10 @@ function(arm_runner_configure_ethos_u_platform) arm_ensure_ethos_u_content( "${ARG_SDK_PATH}" "${EXECUTORCH_ROOT}" ${FETCH_ETHOS_U_CONTENT} ) - add_corstone_subdirectory(${ARG_SYSTEM_CONFIG} ${ARG_SDK_PATH}) - configure_timing_adapters(${ARG_SYSTEM_CONFIG} ${ARG_MEMORY_MODE}) + add_corstone_subdirectory( + "${ARG_SYSTEM_CONFIG}" "${ARG_SDK_PATH}" "${ARG_MEMORY_MODE}" + ) + configure_timing_adapters("${ARG_SYSTEM_CONFIG}" "${ARG_MEMORY_MODE}") foreach(_platform_variable TARGET_BOARD ETHOSU_MODEL ETHOSU_ARENA) if(DEFINED ${_platform_variable}) set(${_platform_variable} diff --git a/backends/arm/common/pipeline_config.py b/backends/arm/common/pipeline_config.py index 3784849556b..4d1563cc328 100644 --- a/backends/arm/common/pipeline_config.py +++ b/backends/arm/common/pipeline_config.py @@ -24,11 +24,18 @@ class LeakyReLULoweringConfig(Enum): class SDPASafeSoftmaxGuardPolicy(Enum): - """Options for preserving or removing SDPA safe-softmax guards.""" + """Options for preserving or removing SDPA safe-softmax guards. + + ``AUTO`` removes guards only for structurally eligible, unmasked, + noncausal, zero-dropout SDPA calls with a nonempty key sequence. ``AUTO`` + and ``REMOVE`` assume attention scores do not become all ``-inf`` through + nonfinite inputs or overflow. + + """ PRESERVE = auto() # Preserve safe-softmax all--inf row guards REMOVE = auto() # Remove exact expanded safe-softmax guards - REMOVE_WHEN_PROVEN = auto() # Preserve unless a proof is available + AUTO = auto() # Remove eligible guards; preserve uncertain cases @dataclass diff --git a/backends/arm/debug/BUCK b/backends/arm/debug/BUCK index 8374fab364e..e87e419763f 100644 --- a/backends/arm/debug/BUCK +++ b/backends/arm/debug/BUCK @@ -2,6 +2,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "schema", diff --git a/backends/arm/ethosu/partitioner.py b/backends/arm/ethosu/partitioner.py index a36584a975a..eb04e69c8bf 100644 --- a/backends/arm/ethosu/partitioner.py +++ b/backends/arm/ethosu/partitioner.py @@ -34,6 +34,7 @@ def __init__( self.delegation_spec = DelegationSpec( EthosUBackend.__name__, compile_spec._to_list() ) + self.compile_spec = compile_spec self.additional_checks = additional_checks self.tosa_spec = compile_spec.tosa_spec self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec) diff --git a/backends/arm/operator_support/BUCK b/backends/arm/operator_support/BUCK index 9152adf2b5b..aee02389128 100644 --- a/backends/arm/operator_support/BUCK +++ b/backends/arm/operator_support/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "operator_support", diff --git a/backends/arm/operator_support/ethos_u55_support.py b/backends/arm/operator_support/ethos_u55_support.py index 35ce0e51df9..fc8182062c3 100644 --- a/backends/arm/operator_support/ethos_u55_support.py +++ b/backends/arm/operator_support/ethos_u55_support.py @@ -14,8 +14,12 @@ import torch import torch.fx as fx -from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor +from executorch.backends.arm._passes.arm_pass_utils import ( + get_first_fake_tensor, + is_param_node, +) from executorch.backends.arm._passes.insert_table_ops import TableOps +from executorch.exir import ExportedProgram from executorch.exir.backend.utils import WhyNoPartitionReporter from executorch.exir.dialects._ops import ops as exir_ops from torch.fx.passes.operator_support import OperatorSupportBase @@ -199,8 +203,6 @@ class EthosU55NotSupported(OperatorSupportBase): exir_ops.edge.aten.ne.Scalar, exir_ops.edge.aten.gather.default, # GATHER exir_ops.edge.aten.grid_sampler_2d, # GATHER - exir_ops.edge.aten.index.Tensor, # GATHER - exir_ops.edge.aten.index_select.default, # GATHER exir_ops.edge.aten.index_put.default, # SCATTER exir_ops.edge.aten.scatter.src, exir_ops.edge.aten.scatter.value, @@ -208,9 +210,6 @@ class EthosU55NotSupported(OperatorSupportBase): exir_ops.edge.aten.scatter_reduce.two, exir_ops.edge.aten.scatter_add.default, exir_ops.edge.aten.upsample_bilinear2d.vec, # RESIZE - exir_ops.edge.aten.reflection_pad1d.default, # REVERSE - exir_ops.edge.aten.reflection_pad2d.default, # REVERSE - exir_ops.edge.aten.reflection_pad3d.default, # REVERSE exir_ops.edge.aten.where.self, # SELECT ] @@ -355,6 +354,35 @@ def is_node_supported( ) return False + reflection_pad_constraints = { + exir_ops.edge.aten.reflection_pad1d.default: ((2, 3), (2,)), + exir_ops.edge.aten.reflection_pad2d.default: ((3, 4), (2, 4)), + exir_ops.edge.aten.reflection_pad3d.default: ((4, 5), (6,)), + } + if node.target in reflection_pad_constraints: + input_shape = get_first_fake_tensor(node.all_input_nodes[0]).shape + padding = typing.cast(typing.Sequence[int], node.args[1]) + supported_ranks, supported_padding_lengths = reflection_pad_constraints[ + node.target + ] + if ( + len(input_shape) in supported_ranks + and len(padding) in supported_padding_lengths + ): + spatial_sizes = tuple(reversed(input_shape[-(len(padding) // 2) :])) + pad_pairs = tuple(zip(padding[::2], padding[1::2])) + if all( + isinstance(size, int) and 0 <= before < size and 0 <= after < size + for (before, after), size in zip(pad_pairs, spatial_sizes) + ): + return True + self.reporter.report_reject( + node, + "U55 reflection padding requires a supported static input rank " + "and nonnegative padding smaller than its spatial dimension.", + ) + return False + return True @@ -402,6 +430,92 @@ def is_node_supported( return True +class EthosU55IndexTensorCheck(OperatorSupportBase): + """Accept single constant index.Tensor cases that lower to slices.""" + + def __init__( + self, exported_program: ExportedProgram, reporter: WhyNoPartitionReporter + ): + self.exported_program = exported_program + self.reporter = reporter + + def is_node_supported( + self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node + ) -> bool: + del submodules + if node.target != exir_ops.edge.aten.index.Tensor: + return True + + input_arg, indices_arg = node.args + input_node = typing.cast(fx.Node, input_arg) + indices = typing.cast(typing.Sequence[fx.Node | None], indices_arg) + input_shape = get_first_fake_tensor(input_node).shape + tensor_indices = [index for index in indices if index is not None] + if len(tensor_indices) != 1: + self.reporter.report_reject( + node, + "U55 index.Tensor only supports indexing along one dimension but got " + f"{len(tensor_indices)}.", + ) + return False + + index_node = tensor_indices[0] + index_shape = get_first_fake_tensor(index_node).shape + if ( + not is_param_node(self.exported_program, index_node) + or len(index_shape) != 1 + or index_shape[0] == 0 + or any(not isinstance(size, int) for size in input_shape) + ): + self.reporter.report_reject( + node, + "U55 index.Tensor requires static input shape and a nonempty " + "constant rank-1 index.", + ) + return False + + return True + + +class EthosU55IndexSelectCheck(OperatorSupportBase): + """Accept constant contiguous index_select cases that lower to a slice.""" + + def __init__( + self, exported_program: ExportedProgram, reporter: WhyNoPartitionReporter + ): + self.exported_program = exported_program + self.reporter = reporter + + def is_node_supported( + self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node + ) -> bool: + del submodules + if node.target != exir_ops.edge.aten.index_select.default: + return True + + input_arg, dim, index_arg = node.args + input_node = typing.cast(fx.Node, input_arg) + index_node = typing.cast(fx.Node, index_arg) + input_shape = get_first_fake_tensor(input_node).shape + index_shape = get_first_fake_tensor(index_node).shape + if ( + not isinstance(dim, int) + or len(input_shape) == 0 + or not is_param_node(self.exported_program, index_node) + or len(index_shape) != 1 + or index_shape[0] == 0 + or any(not isinstance(size, int) for size in input_shape) + ): + self.reporter.report_reject( + node, + "U55 index_select requires static input shape and nonempty " + "constant indices.", + ) + return False + + return True + + class EthosU55CastCheck(OperatorSupportBase): """Reject unsupported casts on U55. diff --git a/backends/arm/operator_support/index_tensor_support.py b/backends/arm/operator_support/index_tensor_support.py index 29134fe964d..b14a30e0e94 100644 --- a/backends/arm/operator_support/index_tensor_support.py +++ b/backends/arm/operator_support/index_tensor_support.py @@ -4,12 +4,13 @@ # LICENSE file in the root directory of this source tree. """Provide TOSA support checks for ``aten.index.Tensor``. -Reject unsupported patterns such as front-positioned slice/ellipsis/None -markers and cases that exceed ``int32`` element limits. +Reject unsupported indexing layouts, zero-sized tensors, and cases that exceed +``int32`` element limits. """ import math +from typing import cast, Sequence import torch import torch.fx as fx @@ -23,6 +24,15 @@ from executorch.exir.dialects._ops import ops as exir_ops +def _has_leading_full_slices_only(indices) -> bool: + found_tensor_index = False + for index in indices: + if index is None and found_tensor_index: + return False + found_tensor_index |= index is not None + return found_tensor_index + + @register_tosa_support_check class IndexTensorSupported(SupportedTOSAOperatorCheck): """Prevent partitioning of unsupported ``index.Tensor`` usages. @@ -30,10 +40,10 @@ class IndexTensorSupported(SupportedTOSAOperatorCheck): This support check is intended to prevent the partitioning of currently unsupported usages of the index.Tensor operator. - 1. Usages where slice, ellipsis or None are present before an indexing tensor: - t[{start}:{end}, indexTensor] - slicing - t[None, indexTensor] - unsqueeze - t[..., indexTensor] - ellipsis + 1. Usages where a slice, ellipsis, or None separates indexing tensors: + t[indexTensor, {start}:{end}, indexTensor] - slicing + t[indexTensor, None, indexTensor] - unsqueeze + t[indexTensor, ..., indexTensor] - ellipsis 2. Usages where the value tensor contains more than int32.max elements This is due to int32 TOSA limitation and the fact that we flatten out @@ -41,25 +51,21 @@ class IndexTensorSupported(SupportedTOSAOperatorCheck): As such to avoid overflow we reject lowering of this operator if it is possible for indices to go over the int32 limit. - Extra information regarding #2: + 3. Usages where the value or an index tensor is zero-sized, because TOSA + requires every tensor dimension to be at least one. + + Extra information regarding #1: Pytorch decomposes slice and None usages before they reach aten. In the case of Slicing and Unsqueeze, Pytorch will add the relevant operation just before the index.Tensor op. In the case of Ellipsis no extra operation is added. - In all three cases Pytorch will insert "None"(s) in the index list - only if the above operations are done on a dimension BEFORE one being indexed. - - When slicing, unsqueeze and ellipsis are done on dimensions after - the ones being indexed, then they do not affect the final output - values, only the shape. Thus None is not passed to the index.Tensor op. - The purpose of None is to signify to index.Tensor that a dimension should not be indexed. - In such cases the logic behaves similar to batching along that dimension. - For the sake of simplicity we have not implemented this behavior yet - and thus have put this support check in place to prevent the partitioning - of index.Tensor ops which include None. + A leading run of None entries behaves like batching along those + dimensions and is supported. None entries after the first tensor index + remain unsupported because they interleave preserved and indexed + dimensions. Examples: #1 - Slice ----------------------------------------------------- @@ -87,12 +93,6 @@ class IndexTensorSupported(SupportedTOSAOperatorCheck): out = ...edge__ops_aten_index_Tensor(unsqueeze_res, [torch.arange(3)]) NB. - With the current implementation of flattening tensors and indices out, - supporting None (Unsqueeze) is simply a matter of ignoring the - None dimension. - This is not the case for Slice and Ellipsis operators, where - the size of the new dimension can be > 1. - Note that slice ops interleaved between indexes such as: t[1:3, torch.arange(5), 2:3, torch.arange(3).reshape(3,1)] are also possible and can result in some unintuitive behaviors @@ -108,27 +108,47 @@ def is_node_tosa_supported( """Return True if ``aten.index.Tensor`` usage fits supported patterns. Enforces the following constraints: - - No ``None`` (unsqueeze), slice, or ellipsis before an indexing tensor. + - ``None`` entries may only form a leading run before all tensor indices. + - At least one tensor index is present. + - Value and index tensors must not be zero-sized. + - Boolean and byte mask indices are not supported. - The value tensor element count fits in ``int32``. """ - indices = node.args[1] - for index in indices: # type: ignore[union-attr] - # Usage 1 guard - if index is None: - self.reporter.report_reject( - node, - ( - "None (from slice/unsqueeze/ellipsis) before an indexing tensor" - " is not supported." - ), - ) - return False + indices = cast(Sequence[fx.Node | None], node.args[1]) + if not _has_leading_full_slices_only(indices): + self.reporter.report_reject( + node, + "Only leading None entries followed by tensor indices are supported.", + ) + return False + + if any( + get_first_fake_tensor(ensure_type(fx.Node, index)).dtype + in (torch.bool, torch.uint8) + for index in indices + if index is not None + ): + self.reporter.report_reject( + node, "Boolean and byte mask indices are not supported." + ) + return False - # Usage 2 guard input_node = ensure_type(torch.fx.Node, node.args[0]) input_val = get_first_fake_tensor(input_node) total_vals = math.prod(input_val.shape) + has_zero_sized_index = any( + math.prod(get_first_fake_tensor(ensure_type(fx.Node, index)).shape) == 0 + for index in indices + if index is not None + ) + if total_vals == 0 or has_zero_sized_index: + self.reporter.report_reject( + node, + "Zero-sized value or index tensors are not supported by TOSA.", + ) + return False + if total_vals > torch.iinfo(torch.int32).max: self.reporter.report_reject( node, diff --git a/backends/arm/operator_support/pool_2d_support.py b/backends/arm/operator_support/pool_2d_support.py index a022ed942fd..03b52fbb85f 100644 --- a/backends/arm/operator_support/pool_2d_support.py +++ b/backends/arm/operator_support/pool_2d_support.py @@ -228,6 +228,14 @@ def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification): """ shape = cast(torch.Tensor, node.all_input_nodes[0].meta["val"]).shape + if len(shape) == 3: + shape = torch.Size((1, *shape)) + elif len(shape) != 4: + self.reporter.report_reject( + node, f"Maxpool2d needs rank 3 or 4 input, got shape {list(shape)}" + ) + return False + kernel = cast(tuple[int, int], node.args[1]) stride = cast(tuple[int, int], node.args[2]) padding = cast(tuple[int, int], node.args[3]) if len(node.args) >= 4 else (0, 0) diff --git a/backends/arm/operator_support/tosa_profile_supported_op_lists.py b/backends/arm/operator_support/tosa_profile_supported_op_lists.py index 78e9c617b13..a9a7c2413b0 100644 --- a/backends/arm/operator_support/tosa_profile_supported_op_lists.py +++ b/backends/arm/operator_support/tosa_profile_supported_op_lists.py @@ -186,6 +186,8 @@ exir_ops.edge.aten.expm1.default, exir_ops.edge.aten.log1p.default, exir_ops.edge.aten.log.default, + exir_ops.edge.aten.isnan.default, + exir_ops.edge.aten.isinf.default, exir_ops.edge.aten.linear.default, exir_ops.edge.aten.split_with_sizes_copy.default, exir_ops.edge.aten.split_copy.Tensor, diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index c1fa7015623..868a4b1d23f 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -39,6 +39,8 @@ from executorch.backends.arm.operator_support.ethos_u55_support import ( EthosU55CastCheck, EthosU55DtypeSupport, + EthosU55IndexSelectCheck, + EthosU55IndexTensorCheck, EthosU55NotSupported, EthosU55ResizeCheck, EthosU55ReverseCheck, @@ -222,7 +224,7 @@ def _floating_profile_negative_checks( ) -> list[OperatorSupportBase]: checks: list[OperatorSupportBase] = [CheckMixedFloatingInputs(reporter)] if not tosa_spec.support_integer(): - checks.append(CheckInt32ComparisonInputs(reporter)) + checks.append(CheckFPComparisonInputs(reporter)) return checks @@ -352,7 +354,7 @@ def _positive_checks( def _disallowed_dtypes(tosa_spec: TosaSpecification) -> list[torch.dtype]: - dtypes = [torch.float64] + dtypes = [torch.float64, torch.complex32, torch.complex64, torch.complex128] if not tosa_spec.support_extension("bf16"): dtypes.append(torch.bfloat16) if not ( @@ -413,6 +415,8 @@ def _negative_checks( checks.append(EthosU55ResizeCheck(reporter)) checks.append(EthosU55ReverseCheck(reporter)) checks.append(EthosU55UnfoldCopyCheck(reporter)) + checks.append(EthosU55IndexTensorCheck(exported_program, reporter)) + checks.append(EthosU55IndexSelectCheck(exported_program, reporter)) checks.append(EthosU55DtypeSupport(reporter)) checks.append(EthosU55CastCheck(reporter)) @@ -762,6 +766,15 @@ def is_node_supported( input_node = node.all_input_nodes[0] input_quantized = FuseQuantizedActivationPass._is_fuseable_input(input_node) + if any( + isinstance(input_node.meta["val"], torch.SymInt) + for input_node in node.all_input_nodes + ): + self.reporter.report_reject( + node, "Symbolic scalar inputs cannot be delegated." + ) + return False + input_quantized = input_quantized or all( (input_node.target in DQ_OPS) or _is_integer_dtype(get_first_fake_tensor(input_node).dtype) @@ -1137,12 +1150,14 @@ def is_node_supported( return True -class CheckInt32ComparisonInputs(OperatorSupportBase): - """Reject int32 comparisons under the FP profile.""" +class CheckFPComparisonInputs(OperatorSupportBase): + """Reject unsupported comparison inputs under the FP profile.""" - target_ops = { + comparison_ops = { exir_ops.edge.aten.eq.Tensor, exir_ops.edge.aten.eq.Scalar, + exir_ops.edge.aten.ne.Tensor, + exir_ops.edge.aten.ne.Scalar, exir_ops.edge.aten.ge.Tensor, exir_ops.edge.aten.ge.Scalar, exir_ops.edge.aten.gt.Tensor, @@ -1152,6 +1167,12 @@ class CheckInt32ComparisonInputs(OperatorSupportBase): exir_ops.edge.aten.lt.Tensor, exir_ops.edge.aten.lt.Scalar, } + target_ops = comparison_ops | { + exir_ops.edge.aten.isinf.default, + exir_ops.edge.aten.isnan.default, + } + supported_dtypes = {torch.float16, torch.float32, torch.bfloat16} + castable_comparison_dtypes = {torch.int8, torch.int16} def __init__(self, reporter: WhyNoPartitionReporter) -> None: self.reporter = reporter @@ -1163,19 +1184,27 @@ def is_node_supported( if node.target not in self.target_ops: return True - for input_node in ( - input_node + input_dtypes = [ + get_first_fake_tensor(input_node).dtype for input_node in node.all_input_nodes if input_node.op != "get_attr" + ] + if all(dtype in self.supported_dtypes for dtype in input_dtypes): + return True + + if node.target in self.comparison_ops and all( + dtype in self.castable_comparison_dtypes for dtype in input_dtypes ): - if get_first_fake_tensor(input_node).dtype == torch.int32: - self.reporter.report_reject( - node, - "FP profile does not support int32 comparison inputs.", - ) - return False + return True - return True + unsupported_dtype = next( + dtype for dtype in input_dtypes if dtype not in self.supported_dtypes + ) + self.reporter.report_reject( + node, + f"FP profile does not support {unsupported_dtype} comparison inputs.", + ) + return False class CheckScalarReductionInputs(OperatorSupportBase): diff --git a/backends/arm/operators/BUCK b/backends/arm/operators/BUCK index 2555aa78035..b775d7337c5 100644 --- a/backends/arm/operators/BUCK +++ b/backends/arm/operators/BUCK @@ -2,6 +2,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "node_visitor", @@ -31,6 +33,7 @@ fbcode_target( "fbsource//third-party/tosa_tools:tosa", ":node_visitor", ":operator_validation_utils", + "//executorch/backends/arm/tosa:constant_pool", "//executorch/backends/arm/tosa:mapping", "//executorch/backends/arm/tosa:utils", "//executorch/backends/arm/_passes:passes", diff --git a/backends/arm/operators/op_tosa_matmul.py b/backends/arm/operators/op_tosa_matmul.py index eb6a26fcc18..e46795c361c 100644 --- a/backends/arm/operators/op_tosa_matmul.py +++ b/backends/arm/operators/op_tosa_matmul.py @@ -76,8 +76,12 @@ def define_node( input_A_ZP_name = f"{output.name}_A_ZP" input_B_ZP_name = f"{output.name}_B_ZP" - tosa_graph.addConst([1], inputs[0].dtype, [input0_zp], name=input_A_ZP_name) - tosa_graph.addConst([1], inputs[1].dtype, [input1_zp], name=input_B_ZP_name) + input_A_ZP = tosa_graph.addConst( + [1], inputs[0].dtype, [input0_zp], name=input_A_ZP_name + ) + input_B_ZP = tosa_graph.addConst( + [1], inputs[1].dtype, [input1_zp], name=input_B_ZP_name + ) # Add the MATMUL to the TOSA graph. attr = ts.TosaSerializerAttribute() @@ -90,8 +94,8 @@ def define_node( [ inputs[0].name, inputs[1].name, - input_A_ZP_name, - input_B_ZP_name, + input_A_ZP.name, + input_B_ZP.name, ], [output.name], attr, diff --git a/backends/arm/operators/op_tosa_mul.py b/backends/arm/operators/op_tosa_mul.py index 65521eab88d..38ee2a6e9b0 100644 --- a/backends/arm/operators/op_tosa_mul.py +++ b/backends/arm/operators/op_tosa_mul.py @@ -47,14 +47,14 @@ def define_node( self.tosa_spec, ) - tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{output.name}_shift") + shift = tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{output.name}_shift") attr = ts.TosaSerializerAttribute() attr.MulAttribute() self._serialize_operator( node, tosa_graph, ts.Op.MUL, - [inputs[0].name, inputs[1].name, f"{output.name}_shift"], + [inputs[0].name, inputs[1].name, shift.name], [output.name], attr, ) diff --git a/backends/arm/operators/op_tosa_shapes.py b/backends/arm/operators/op_tosa_shapes.py index b7480d78a4d..f5831f2e440 100644 --- a/backends/arm/operators/op_tosa_shapes.py +++ b/backends/arm/operators/op_tosa_shapes.py @@ -14,6 +14,7 @@ register_node_visitor, ) from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.constant_pool import TosaSerializerWithConstantPool from executorch.backends.arm.tosa.mapping import TosaArg from executorch.backends.arm.tosa.utils import normalize_symint @@ -32,8 +33,10 @@ def define_node( shape_input = inputs[0].special rank = len(shape_input) vals = normalize_symint(node.meta["val"]) - tosa_graph = cast(ts.TosaSerializer, tosa_graph) - tosa_graph.addConst( + tosa_graph = cast(TosaSerializerWithConstantPool, tosa_graph) + # Downstream visitors reference this FX output by name. Pooling it with a + # serializer-generated constant could leave output.name undefined. + tosa_graph.addUnpooledConst( [ rank, ], diff --git a/backends/arm/operators/ops_quant_utils.py b/backends/arm/operators/ops_quant_utils.py index 1b5bf4caea9..fef9e3c56c1 100644 --- a/backends/arm/operators/ops_quant_utils.py +++ b/backends/arm/operators/ops_quant_utils.py @@ -25,12 +25,11 @@ def add_input_weight_zp_consts(tosa_graph, node, inputs, output_name): input_zp_name = f"{output_name}_input_zp" weight_zp_name = f"{output_name}_weight_zp" - tosa_graph.addConst([1], inputs[0].dtype, [input_zp], name=input_zp_name) - tosa_graph.addConst( - [1], - inputs[1].dtype, - weight_zp, - name=weight_zp_name, + input_zp_tensor = tosa_graph.addConst( + [1], inputs[0].dtype, [input_zp], name=input_zp_name + ) + weight_zp_tensor = tosa_graph.addConst( + [1], inputs[1].dtype, weight_zp, name=weight_zp_name ) - return input_zp_name, weight_zp_name + return input_zp_tensor.name, weight_zp_tensor.name diff --git a/backends/arm/process_node.py b/backends/arm/process_node.py index a0c2dbeb1fb..129413cabe7 100644 --- a/backends/arm/process_node.py +++ b/backends/arm/process_node.py @@ -106,12 +106,18 @@ def _add_const( tosa_arg: TosaArg, name: str, ) -> None: - """Add a constant, preserving packed FP4 storage when required.""" + """Add a graph-owned constant under its exact name. + + Parameters, buffers, and lifted constants are referenced by their FX names, + so pooling them could leave those names undefined. Preserve packed FP4 + storage when required. + + """ if _is_packed_fp4_const(values, tosa_arg): # TOSA FP4 tensors have logical FP4 shape, but constants are stored as # packed bytes (two values per byte). Add the raw bytes as INT8 first # then set TOSA dtype and shape correctly on the tensor metadata. - tosa_graph.addConst( + tosa_graph.addUnpooledConst( normalize_symint(values.shape), ts.DType.INT8, values, @@ -124,7 +130,7 @@ def _add_const( return prepared_values = _prepare_const_values_for_tosa_dtype(values, tosa_arg) - tosa_graph.addConst( + tosa_graph.addUnpooledConst( _get_const_shape(prepared_values, tosa_arg), tosa_arg.dtype, prepared_values, diff --git a/backends/arm/public_api_manifests/api_manifest_1_5.toml b/backends/arm/public_api_manifests/api_manifest_1_5.toml new file mode 100644 index 00000000000..dc87ac9ee96 --- /dev/null +++ b/backends/arm/public_api_manifests/api_manifest_1_5.toml @@ -0,0 +1,287 @@ +# 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. +# +# This file is generated by +# backends/arm/scripts/public_api_manifest/generate_public_api_manifest.py + +[python] + +[python.EthosUBackend] +kind = "class" +signature = "EthosUBackend()" + +[python.EthosUBackend.preprocess] +kind = "function" +signature = "EthosUBackend.preprocess(edge_program: torch.export.exported_program.ExportedProgram, compile_specs: List[executorch.exir.backend.compile_spec_schema.CompileSpec]) -> executorch.exir.backend.backend_details.PreprocessResult" + +[python.EthosUCompileSpec] +kind = "class" +signature = "EthosUCompileSpec(target: str, system_config: str | None = None, memory_mode: str | None = None, extra_flags: list[str] | None = None, config_ini: str | None = 'Arm/vela.ini', external_block_placements: executorch.backends.arm.ethosu.compile_spec.VelaExternalBlockPlacements | None = None)" + +[python.EthosUCompileSpec.DebugMode] +kind = "enum" +signature = "EthosUCompileSpec.DebugMode(*values)" + +[python.EthosUCompileSpec.__eq__] +kind = "function" +signature = "EthosUCompileSpec.__eq__(self, other)" + +[python.EthosUCompileSpec.__repr__] +kind = "function" +signature = "EthosUCompileSpec.__repr__(self)" + +[python.EthosUCompileSpec.dump_debug_info] +kind = "function" +signature = "EthosUCompileSpec.dump_debug_info(self, debug_mode: executorch.backends.arm.common.arm_compile_spec.ArmCompileSpec.DebugMode | None)" + +[python.EthosUCompileSpec.dump_intermediate_artifacts_to] +kind = "function" +signature = "EthosUCompileSpec.dump_intermediate_artifacts_to(self, output_path: str | None)" + +[python.EthosUCompileSpec.set_pass_pipeline_config] +kind = "function" +signature = "EthosUCompileSpec.set_pass_pipeline_config(self, config: executorch.backends.arm.common.pipeline_config.ArmPassPipelineConfig) -> None" + +[python.EthosUPartitioner] +kind = "class" +signature = "EthosUPartitioner(compile_spec: executorch.backends.arm.ethosu.compile_spec.EthosUCompileSpec, additional_checks: Optional[Sequence[torch.fx.passes.operator_support.OperatorSupportBase]] = None) -> None" + +[python.EthosUPartitioner.ops_to_not_decompose] +kind = "function" +signature = "EthosUPartitioner.ops_to_not_decompose(self, ep: torch.export.exported_program.ExportedProgram) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.node.Node], bool]]]" + +[python.EthosUPartitioner.partition] +kind = "function" +signature = "EthosUPartitioner.partition(self, exported_program: torch.export.exported_program.ExportedProgram) -> executorch.exir.backend.partitioner.PartitionResult" + +[python.EthosUPartitioner.register_custom_partition_op] +kind = "function" +signature = "EthosUPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None" + +[python.EthosUPartitioner.transform_for_pre_decomposition] +kind = "function" +signature = "EthosUPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram" + +[python.EthosUQuantizer] +kind = "class" +signature = "EthosUQuantizer(compile_spec: 'EthosUCompileSpec', use_composable_quantizer: 'bool' = True) -> 'None'" + +[python.EthosUQuantizer.annotate] +kind = "function" +signature = "EthosUQuantizer.annotate(self, model: 'GraphModule') -> 'GraphModule'" + +[python.EthosUQuantizer.set_global] +kind = "function" +signature = "EthosUQuantizer.set_global(self, quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.EthosUQuantizer.set_io] +kind = "function" +signature = "EthosUQuantizer.set_io(self, quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.EthosUQuantizer.set_module_name] +kind = "function" +signature = "EthosUQuantizer.set_module_name(self, module_name: 'str', quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.EthosUQuantizer.set_module_type] +kind = "function" +signature = "EthosUQuantizer.set_module_type(self, module_type: 'Callable', quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.EthosUQuantizer.transform_for_annotation] +kind = "function" +signature = "EthosUQuantizer.transform_for_annotation(self, model: 'GraphModule') -> 'GraphModule'" + +[python.EthosUQuantizer.validate] +kind = "function" +signature = "EthosUQuantizer.validate(self, model: 'GraphModule') -> 'None'" + +[python.VelaExternalBlockPlacements] +kind = "class" +signature = "VelaExternalBlockPlacements(cmd_data: str | None = None, weight_data: str | None = None) -> None" + +[python.VelaExternalBlockPlacements.__delattr__] +kind = "function" +signature = "VelaExternalBlockPlacements.__delattr__(self, name)" + +[python.VelaExternalBlockPlacements.__eq__] +kind = "function" +signature = "VelaExternalBlockPlacements.__eq__(self, other)" + +[python.VelaExternalBlockPlacements.__hash__] +kind = "function" +signature = "VelaExternalBlockPlacements.__hash__(self)" + +[python.VelaExternalBlockPlacements.__post_init__] +kind = "function" +signature = "VelaExternalBlockPlacements.__post_init__(self) -> None" + +[python.VelaExternalBlockPlacements.__repr__] +kind = "function" +signature = "VelaExternalBlockPlacements.__repr__(self)" + +[python.VelaExternalBlockPlacements.__setattr__] +kind = "function" +signature = "VelaExternalBlockPlacements.__setattr__(self, name, value)" + +[python.VelaExternalBlockPlacements.to_block_placements] +kind = "function" +signature = "VelaExternalBlockPlacements.to_block_placements(self) -> dict[str, str]" + +[python.VgfBackend] +kind = "class" +signature = "VgfBackend()" + +[python.VgfBackend.preprocess] +kind = "function" +signature = "VgfBackend.preprocess(edge_program: torch.export.exported_program.ExportedProgram, compile_specs: List[executorch.exir.backend.compile_spec_schema.CompileSpec]) -> executorch.exir.backend.backend_details.PreprocessResult" + +[python.VgfCompileSpec] +kind = "class" +signature = "VgfCompileSpec(tosa_spec: executorch.backends.arm.tosa.specification.TosaSpecification | str | None = None, compiler_flags: list[str] | None = None)" + +[python.VgfCompileSpec.DebugMode] +kind = "enum" +signature = "VgfCompileSpec.DebugMode(*values)" + +[python.VgfCompileSpec.__eq__] +kind = "function" +signature = "VgfCompileSpec.__eq__(self, other)" + +[python.VgfCompileSpec.__repr__] +kind = "function" +signature = "VgfCompileSpec.__repr__(self)" + +[python.VgfCompileSpec.dump_debug_info] +kind = "function" +signature = "VgfCompileSpec.dump_debug_info(self, debug_mode: executorch.backends.arm.common.arm_compile_spec.ArmCompileSpec.DebugMode | None)" + +[python.VgfCompileSpec.dump_intermediate_artifacts_to] +kind = "function" +signature = "VgfCompileSpec.dump_intermediate_artifacts_to(self, output_path: str | None)" + +[python.VgfCompileSpec.set_pass_pipeline_config] +kind = "function" +signature = "VgfCompileSpec.set_pass_pipeline_config(self, config: executorch.backends.arm.common.pipeline_config.ArmPassPipelineConfig) -> None" + +[python.VgfCompileSpec.validate_environment] +kind = "function" +signature = "VgfCompileSpec.validate_environment(self, build_dir: str | None = None, *, require_runtime_build: bool = False) -> 'VgfEnvironmentReport'" + +[python.VgfPartitioner] +kind = "class" +signature = "VgfPartitioner(compile_spec: executorch.backends.arm.vgf.compile_spec.VgfCompileSpec, additional_checks: Optional[Sequence[torch.fx.passes.operator_support.OperatorSupportBase]] = None) -> None" + +[python.VgfPartitioner.ops_to_not_decompose] +kind = "function" +signature = "VgfPartitioner.ops_to_not_decompose(self, ep: torch.export.exported_program.ExportedProgram) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.node.Node], bool]]]" + +[python.VgfPartitioner.partition] +kind = "function" +signature = "VgfPartitioner.partition(self, exported_program: torch.export.exported_program.ExportedProgram) -> executorch.exir.backend.partitioner.PartitionResult" + +[python.VgfPartitioner.register_custom_partition_op] +kind = "function" +signature = "VgfPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None" + +[python.VgfPartitioner.transform_for_pre_decomposition] +kind = "function" +signature = "VgfPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram" + +[python.VgfQuantizer] +kind = "class" +signature = "VgfQuantizer(compile_spec: 'VgfCompileSpec', use_composable_quantizer: 'bool' = True) -> 'None'" + +[python.VgfQuantizer.annotate] +kind = "function" +signature = "VgfQuantizer.annotate(self, model: 'GraphModule') -> 'GraphModule'" + +[python.VgfQuantizer.set_global] +kind = "function" +signature = "VgfQuantizer.set_global(self, quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.VgfQuantizer.set_io] +kind = "function" +signature = "VgfQuantizer.set_io(self, quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.VgfQuantizer.set_module_name] +kind = "function" +signature = "VgfQuantizer.set_module_name(self, module_name: 'str', quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.VgfQuantizer.set_module_type] +kind = "function" +signature = "VgfQuantizer.set_module_type(self, module_type: 'Callable', quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.VgfQuantizer.transform_for_annotation] +kind = "function" +signature = "VgfQuantizer.transform_for_annotation(self, model: 'GraphModule') -> 'GraphModule'" + +[python.VgfQuantizer.validate] +kind = "function" +signature = "VgfQuantizer.validate(self, model: 'GraphModule') -> 'None'" + +[python.get_symmetric_a16w8_quantization_config] +kind = "function" +signature = "get_symmetric_a16w8_quantization_config(is_per_channel: 'bool' = True, is_qat: 'bool' = False, is_dynamic: 'bool' = False, weight_qmin: 'int' = -127, weight_qmax: 'int' = 127, epsilon: 'float' = 0.000244140625) -> 'QuantizationConfig'" + +[python.get_symmetric_quantization_config] +kind = "function" +signature = "get_symmetric_quantization_config(is_per_channel: 'bool' = True, is_qat: 'bool' = False, is_dynamic: 'bool' = False, act_qmin: 'int' = -128, act_qmax: 'int' = 127, weight_qmin: 'int' = -127, weight_qmax: 'int' = 127, eps: 'float' = 1.52587890625e-05) -> 'QuantizationConfig'" + +[cmake] + +[cmake.arm_runner_add_minimal_executable] +kind = "function" +signature = "arm_runner_add_minimal_executable(*, TARGET, SOURCE, OPS_PREFIX, COMPILE_DEFINITIONS=())" + +[cmake.arm_runner_add_standalone_executorch] +kind = "macro" +signature = "arm_runner_add_standalone_executorch()" + +[cmake.arm_runner_configure_ethos_u_platform] +kind = "function" +signature = "arm_runner_configure_ethos_u_platform(*, SDK_PATH, SYSTEM_CONFIG, MEMORY_MODE)" + +[cmake.arm_runner_configure_linker_script] +kind = "function" +signature = "arm_runner_configure_linker_script(*, TARGET, SYSTEM_CONFIG, OUTPUT_NAME=None)" + +[cmake.arm_runner_configure_model] +kind = "function" +signature = "arm_runner_configure_model(*, TARGET, PTE_FILE=None, MODEL_PTE_ADDR=None, MODEL_PTE_SIZE=None, PUBLIC=False)" + +[cmake.arm_runner_configure_runtime_output] +kind = "function" +signature = "arm_runner_configure_runtime_output(TARGET_NAME, FALLBACK_DIR)" + +[cmake.arm_runner_create_default_selected_ops_libs] +kind = "function" +signature = "arm_runner_create_default_selected_ops_libs(*, PREFIX, SUFFIX=None, OP_LIST=None, OPS_FROM_MODEL=None, DTYPE_SELECTIVE_BUILD=None, OUT_LIBS=None, DEPS=())" + +[cmake.arm_runner_create_selected_ops_lib] +kind = "function" +signature = "arm_runner_create_selected_ops_lib(*, LIB_NAME, FUNCTIONS_YAML=None, CUSTOM_OPS_YAML=None, OP_LIST=None, OPS_FROM_MODEL=None, DTYPE_SELECTIVE_BUILD=None, KERNEL_LIBS=(), DEPS=(), INCLUDE_ALL_OPS=False, PRIM_OPS=False)" + +[cmake.arm_runner_define_cache_options] +kind = "function" +signature = "arm_runner_define_cache_options(*, METHOD_ALLOCATOR_SIZE=None)" + +[cmake.arm_runner_link_minimal_specs] +kind = "function" +signature = "arm_runner_link_minimal_specs(TARGET_NAME)" + +[cmake.arm_runner_link_registration_libraries] +kind = "function" +signature = "arm_runner_link_registration_libraries(*, TARGET, SCOPE=None, BASE_LIBS=(), REGISTRATION_LIBS=(), NORMAL_LIBS=(), SUPPRESS_LIBS=())" + +[cmake.arm_runner_require_baremetal_targets] +kind = "function" +signature = "arm_runner_require_baremetal_targets()" + +[cmake.arm_runner_require_python] +kind = "macro" +signature = "arm_runner_require_python()" + +[cmake.arm_runner_validate_model_source] +kind = "function" +signature = "arm_runner_validate_model_source(*, ALLOW_SEMIHOSTING=False)" diff --git a/backends/arm/public_api_manifests/api_manifest_running.toml b/backends/arm/public_api_manifests/api_manifest_running.toml index 9ff031c41bf..dc87ac9ee96 100644 --- a/backends/arm/public_api_manifests/api_manifest_running.toml +++ b/backends/arm/public_api_manifests/api_manifest_running.toml @@ -60,6 +60,10 @@ signature = "EthosUPartitioner.partition(self, exported_program: torch.export.ex kind = "function" signature = "EthosUPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None" +[python.EthosUPartitioner.transform_for_pre_decomposition] +kind = "function" +signature = "EthosUPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram" + [python.EthosUQuantizer] kind = "class" signature = "EthosUQuantizer(compile_spec: 'EthosUCompileSpec', use_composable_quantizer: 'bool' = True) -> 'None'" @@ -180,6 +184,10 @@ signature = "VgfPartitioner.partition(self, exported_program: torch.export.expor kind = "function" signature = "VgfPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None" +[python.VgfPartitioner.transform_for_pre_decomposition] +kind = "function" +signature = "VgfPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram" + [python.VgfQuantizer] kind = "class" signature = "VgfQuantizer(compile_spec: 'VgfCompileSpec', use_composable_quantizer: 'bool' = True) -> 'None'" diff --git a/backends/arm/quantizer/BUCK b/backends/arm/quantizer/BUCK index 632a38523e2..c9c2fcdb203 100644 --- a/backends/arm/quantizer/BUCK +++ b/backends/arm/quantizer/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + # Exposed through __init__.py fbcode_target( _kind = runtime.python_library, diff --git a/backends/arm/quantizer/arm_quantizer.py b/backends/arm/quantizer/arm_quantizer.py index 121db2e0c90..8a1171421fc 100644 --- a/backends/arm/quantizer/arm_quantizer.py +++ b/backends/arm/quantizer/arm_quantizer.py @@ -14,7 +14,7 @@ import functools import logging from contextlib import contextmanager -from typing import Any, Callable, Dict, List, Optional +from typing import Any, Callable, Dict, Iterable, List, Optional import torch from executorch.backends.arm._passes import ArmPassManager @@ -959,7 +959,7 @@ def validate(self, model: GraphModule) -> None: def _quantize_with_submodules( self, model: GraphModule, - calibration_samples: list[tuple], + calibration_samples: Iterable[tuple], is_qat: bool = False, fold_quantize: bool = True, ): @@ -971,7 +971,7 @@ def _quantize_with_submodules( Args: model (GraphModule): The model to quantize. - calibration_samples (list[tuple]): A list of inputs to used to + calibration_samples (Iterable[tuple]): Inputs used to calibrate the model during quantization. To properly calibrate a model with submodules, at least one sample per code path is needed. diff --git a/backends/arm/quantizer/arm_quantizer_utils.py b/backends/arm/quantizer/arm_quantizer_utils.py index 8689fbdadec..daef262ccf6 100644 --- a/backends/arm/quantizer/arm_quantizer_utils.py +++ b/backends/arm/quantizer/arm_quantizer_utils.py @@ -451,6 +451,7 @@ class SharedQspecQuantizer(Quantizer, QuantizerReporterUser): torch.ops.aten.split_copy.Tensor, torch.ops.aten.tile.default, torch.ops.aten.flip.default, + torch.ops.aten.roll.default, torch.ops.aten.index_select.default, torch.ops.aten.index_put.default, torch.ops.aten.index_put_.default, diff --git a/backends/arm/quantizer/quantization_annotator.py b/backends/arm/quantizer/quantization_annotator.py index be0c7b7b453..5cc4d46fff4 100644 --- a/backends/arm/quantizer/quantization_annotator.py +++ b/backends/arm/quantizer/quantization_annotator.py @@ -611,6 +611,7 @@ def _get_fixed_qparams_qspec( torch.ops.aten.t_copy.default, torch.ops.aten.tile.default, torch.ops.aten.flip.default, + torch.ops.aten.roll.default, torch.ops.aten.chunk.default, torch.ops.aten.contiguous.default, torch.ops.aten.upsample_bilinear2d.vec, diff --git a/backends/arm/recipes/BUCK b/backends/arm/recipes/BUCK new file mode 100644 index 00000000000..df3c7874f9e --- /dev/null +++ b/backends/arm/recipes/BUCK @@ -0,0 +1,60 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +# 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. + +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +fbcode_target( + _kind = runtime.python_library, + name = "recipes", + srcs = [ + "__init__.py", + ], + visibility = ["PUBLIC"], + deps = [ + ":arm_recipe_provider", + ":arm_recipe_types", + "//executorch/export:recipe_registry", + ], +) + +fbcode_target( + _kind = runtime.python_library, + name = "arm_recipe_provider", + srcs = [ + "arm_recipe_provider.py", + ], + visibility = ["PUBLIC"], + deps = [ + ":arm_recipe_types", + # Imported lazily for the accelerator-config list; declared so the dep + # does not rely on reaching vela through :ethosu. + "fbsource//third-party/pypi/ethos-u-vela:ethos-u-vela", + "//executorch/backends/arm:_factory", + "//executorch/backends/arm:arm_compile_spec", + "//executorch/backends/arm:ethosu", + "//executorch/backends/arm:vgf", + "//executorch/backends/arm/quantizer:lib", + "//executorch/backends/arm/tosa:compile_spec", + "//executorch/backends/cortex_m/passes:replace_quant_nodes_pass", + "//executorch/exir:lib", + "//executorch/export:lib", + ], +) + +fbcode_target( + _kind = runtime.python_library, + name = "arm_recipe_types", + srcs = [ + "arm_recipe_types.py", + ], + visibility = ["PUBLIC"], + deps = [ + "//executorch/export:recipe", + ], +) diff --git a/backends/arm/recipes/__init__.py b/backends/arm/recipes/__init__.py new file mode 100644 index 00000000000..2b751645d68 --- /dev/null +++ b/backends/arm/recipes/__init__.py @@ -0,0 +1,15 @@ +# 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. + +from executorch.export import recipe_registry + +from .arm_recipe_provider import ArmRecipeProvider +from .arm_recipe_types import ArmRecipeType + +recipe_registry.register_backend_recipe_provider(ArmRecipeProvider()) + + +__all__ = ["ArmRecipeProvider", "ArmRecipeType"] diff --git a/backends/arm/recipes/arm_recipe_provider.py b/backends/arm/recipes/arm_recipe_provider.py new file mode 100644 index 00000000000..35f1370856a --- /dev/null +++ b/backends/arm/recipes/arm_recipe_provider.py @@ -0,0 +1,287 @@ +# 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. + +# pyre-strict + +import logging +from dataclasses import dataclass +from typing import Any, Callable, Optional, Sequence + +from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec +from executorch.backends.arm.ethosu import EthosUCompileSpec +from executorch.backends.arm.quantizer import ( + get_symmetric_a16w8_quantization_config, + get_symmetric_quantization_config, +) +from executorch.backends.arm.recipes.arm_recipe_types import ARM_BACKEND, ArmRecipeType +from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec +from executorch.backends.arm.util._factory import create_partitioner, create_quantizer +from executorch.backends.arm.vgf import VgfCompileSpec +from executorch.exir.capture import EdgeCompileConfig, ExecutorchBackendConfig +from executorch.exir.pass_manager import PassType +from executorch.exir.program import EdgeProgramManager +from executorch.export import ( + BackendRecipeProvider, + ExportRecipe, + LoweringRecipe, + QuantizationRecipe, + RecipeType, +) + + +logger: logging.Logger = logging.getLogger(__name__) + +# (target prefix, default MAC count). Which counts are *accepted* is Vela's to +# say, so it is asked at build time rather than restated here. +_ETHOS_U_FAMILIES: dict[ArmRecipeType, tuple[str, int]] = { + ArmRecipeType.ETHOS_U55_INT8: ("ethos-u55", 128), + ArmRecipeType.ETHOS_U65_INT8: ("ethos-u65", 256), + ArmRecipeType.ETHOS_U85_INT8: ("ethos-u85", 256), +} + +_ETHOS_U_KWARGS: frozenset[str] = frozenset( + {"macs", "system_config", "memory_mode", "extra_flags", "config_ini"} +) + +# Prepended to any caller-supplied Vela flags, matching `_get_compile_spec` in +# aot_arm_compiler.py. +_VELA_DEFAULT_FLAGS: tuple[str, ...] = ( + "--verbose-operators", + "--verbose-cycle-estimate", +) + + +@dataclass(frozen=True) +class _TosaVersionedTarget: + """A target with no caller-tunable options: class plus TOSA version.""" + + compile_spec: Callable[[str], ArmCompileSpec] + tosa_spec: str + quant_mode: Optional[str] + replace_quant_nodes: bool + + +# VGF keeps the quantized_decomposed QDQ ops it is given; see +# `_apply_replace_quant_nodes` in aot_arm_compiler.py. +_TOSA_VERSIONED_TARGETS: dict[ArmRecipeType, _TosaVersionedTarget] = { + ArmRecipeType.TOSA_FP: _TosaVersionedTarget( + TosaCompileSpec, "TOSA-1.0+FP", None, False + ), + ArmRecipeType.TOSA_INT8: _TosaVersionedTarget( + TosaCompileSpec, "TOSA-1.0+INT", "INT8", True + ), + ArmRecipeType.TOSA_A16W8: _TosaVersionedTarget( + TosaCompileSpec, "TOSA-1.0+INT+int16", "A16W8", True + ), + ArmRecipeType.VGF_FP: _TosaVersionedTarget( + VgfCompileSpec, "TOSA-1.0+FP", None, False + ), + ArmRecipeType.VGF_INT8: _TosaVersionedTarget( + VgfCompileSpec, "TOSA-1.0+INT", "INT8", False + ), +} + + +def _reject_unsupported_accelerator( + recipe_type: ArmRecipeType, family: str, target: str, macs: int +) -> None: + """Ask Vela which accelerator configurations it accepts. + + A local copy would go stale on the next Vela bump. Without Vela there is + nothing to check against and the compile spec still has to build, so the + import is guarded the way `arm_vela` guards its own. + + """ + try: + from ethosu.vela.architecture_features import Accelerator # type: ignore + except ImportError: + logger.debug("ethos-u-vela is not installed; macs=%s unvalidated", macs) + return + + supported = {accelerator.value for accelerator in Accelerator} + if target not in supported: + allowed = sorted( + int(name.rsplit("-", 1)[1]) + for name in supported + if name.startswith(f"{family}-") + ) + raise ValueError( + f"Recipe '{recipe_type.value}' does not support macs={macs}. " + f"Allowed: {allowed}" + ) + + +def _replace_quant_nodes( + edge_program_manager: EdgeProgramManager, +) -> list[PassType]: + """Rewrite the QDQ ops left outside the delegate into cortex_m kernels. + + Matches `_apply_replace_quant_nodes` in aot_arm_compiler.py, which applies + the same pass to the whole manager once the partitioner has run. Without it + the boundary quantize/dequantize keep their `quantized_decomposed` targets, + which have no out variants, and `to_executorch` refuses to emit them. + + """ + # Function-local: an FP recipe must not pull in the cortex_m operator + # library, which registers its whole op set on import. + from executorch.backends.cortex_m.passes.replace_quant_nodes_pass import ( + ReplaceQuantNodesPass, + ) + + return [ReplaceQuantNodesPass()] + + +class ArmRecipeProvider(BackendRecipeProvider): + """Builds ExportRecipes for the delegated Arm targets: Ethos-U, TOSA, VGF. + + Each recipe is built to reproduce the default + ``backends/arm/scripts/aot_arm_compiler.py`` invocation for its target: the + same compile spec, quantizer, pass pipeline and backend config. The CLI + options with no recipe equivalent, debug mode and direct drive, are the + exceptions. + """ + + @property + def backend_name(self) -> str: + return ARM_BACKEND + + def get_supported_recipes(self) -> Sequence[RecipeType]: + return list(_ETHOS_U_FAMILIES) + list(_TOSA_VERSIONED_TARGETS) + + def create_recipe( + self, recipe_type: RecipeType, **kwargs: Any + ) -> Optional[ExportRecipe]: + if not isinstance(recipe_type, ArmRecipeType): + return None + + if recipe_type in _ETHOS_U_FAMILIES: + self._warn_unknown_kwargs(recipe_type, kwargs, _ETHOS_U_KWARGS) + return self._build_recipe( + recipe_type, + self._ethos_u_compile_spec(recipe_type, kwargs), + quant_mode="INT8", + replace_quant_nodes=True, + ) + + target = _TOSA_VERSIONED_TARGETS.get(recipe_type) + if target is None: + return None + + self._warn_unknown_kwargs(recipe_type, kwargs, frozenset()) + return self._build_recipe( + recipe_type, + target.compile_spec(target.tosa_spec), + quant_mode=target.quant_mode, + replace_quant_nodes=target.replace_quant_nodes, + ) + + @staticmethod + def _ethos_u_compile_spec( + recipe_type: ArmRecipeType, kwargs: dict[str, Any] + ) -> EthosUCompileSpec: + family, default_macs = _ETHOS_U_FAMILIES[recipe_type] + macs = kwargs.get("macs", default_macs) + if not isinstance(macs, int): + raise ValueError(f"macs must be an int, got {macs!r}") + + extra_flags = kwargs.get("extra_flags") or [] + # The list check comes first: a bare string would be iterated into one + # flag per character, and anything not iterable would raise TypeError + # out of `all` rather than reaching this message. + if not isinstance(extra_flags, list) or not all( + isinstance(flag, str) for flag in extra_flags + ): + raise ValueError( + f"extra_flags must be a list of strings, got {extra_flags!r}" + ) + + target = f"{family}-{macs}" + _reject_unsupported_accelerator(recipe_type, family, target, macs) + + return EthosUCompileSpec( + target=target, + system_config=kwargs.get("system_config"), + memory_mode=kwargs.get("memory_mode"), + extra_flags=list(_VELA_DEFAULT_FLAGS) + list(extra_flags), + # EthosUCompileSpec owns the default. + config_ini=kwargs.get("config_ini"), + ) + + @classmethod + def _build_recipe( + cls, + recipe_type: ArmRecipeType, + compile_spec: ArmCompileSpec, + quant_mode: Optional[str], + replace_quant_nodes: bool, + ) -> ExportRecipe: + # The partitioner snapshots the compile spec and the pipeline config is + # materialised on first read, which the CLI gets for free by quantizing + # before it partitions. + compile_spec.set_pass_pipeline_config(compile_spec._get_pass_pipeline_config()) + + return ExportRecipe( + name=recipe_type.value, + quantization_recipe=cls._build_quantization_recipe( + compile_spec, quant_mode + ), + lowering_recipe=LoweringRecipe( + partitioners=[create_partitioner(compile_spec)], + # The CLI disables edge verification on every Arm path. + edge_compile_config=EdgeCompileConfig(_check_ir_validity=False), + edge_manager_transform_passes=( + [_replace_quant_nodes] if replace_quant_nodes else None + ), + ), + # The Arm runtime expects the delegate payload inline rather than + # in its own segment, as every other Arm AOT path asks for. + executorch_backend_config=ExecutorchBackendConfig( + extract_delegate_segments=False + ), + ) + + @staticmethod + def _build_quantization_recipe( + compile_spec: ArmCompileSpec, quant_mode: Optional[str] + ) -> Optional[QuantizationRecipe]: + if quant_mode is None: + return None + + if quant_mode == "INT8": + operator_config = get_symmetric_quantization_config(is_per_channel=True) + elif quant_mode == "A16W8": + if not compile_spec.tosa_spec.support_extension("int16"): + raise ValueError( + f"TOSA spec {compile_spec.tosa_spec} does not support int16 " + "(required for A16W8)" + ) + operator_config = get_symmetric_a16w8_quantization_config( + is_per_channel=True + ) + else: + raise ValueError(f"Unsupported quant_mode: {quant_mode}") + + quantizer = create_quantizer(compile_spec) + quantizer.set_global(operator_config) + return QuantizationRecipe(quantizers=[quantizer]) + + @staticmethod + def _warn_unknown_kwargs( + recipe_type: ArmRecipeType, + kwargs: dict[str, Any], + expected: frozenset[str], + ) -> None: + # Warn, as XNNPACK and QNN do: `_create_target_recipe` hands every + # recipe in a combination the same kwargs. + unexpected = set(kwargs.keys()) - expected + if unexpected: + allowed = sorted(expected) if expected else "none" + logger.warning( + "Arm recipe '%s' ignoring unexpected parameters: %s. Allowed: %s", + recipe_type.value, + sorted(unexpected), + allowed, + ) diff --git a/backends/arm/recipes/arm_recipe_types.py b/backends/arm/recipes/arm_recipe_types.py new file mode 100644 index 00000000000..91904a4d981 --- /dev/null +++ b/backends/arm/recipes/arm_recipe_types.py @@ -0,0 +1,52 @@ +# 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. + +# pyre-strict + +from executorch.export import RecipeType + + +ARM_BACKEND: str = "arm" + + +class ArmRecipeType(RecipeType): + """Arm-specific recipe types. + + Covers the delegated targets of ``backends/arm/scripts/aot_arm_compiler.py``. + Its non-delegated Cortex-M/CMSIS-NN path is a separate backend and is not + reachable from these recipes. + + Ethos-U recipes accept the following kwargs: + macs (int): MAC count for the family, validated against the accelerator + configurations the installed Vela accepts -- today 32/64/128/256 for + U55, 256/512 for U65 and 128/256/512/1024/2048 for U85. Defaults to + 128 for U55 and 256 for U65 and U85. + system_config (str): Vela system config name. Defaults from + ``EthosUCompileSpec`` apply when omitted. + memory_mode (str): Vela memory mode. Defaults from + ``EthosUCompileSpec`` apply when omitted. + extra_flags (list[str]): Vela compiler flags, appended to the + ``--verbose-operators --verbose-cycle-estimate`` the CLI always + passes rather than replacing them. + config_ini (str): Path to a Vela .ini configuration file. Defaults to + ``"Arm/vela.ini"``. + + """ + + ETHOS_U55_INT8 = "arm_ethos_u55_int8" + ETHOS_U65_INT8 = "arm_ethos_u65_int8" + ETHOS_U85_INT8 = "arm_ethos_u85_int8" + + TOSA_FP = "arm_tosa_fp" + TOSA_INT8 = "arm_tosa_int8" + TOSA_A16W8 = "arm_tosa_a16w8" + + VGF_FP = "arm_vgf_fp" + VGF_INT8 = "arm_vgf_int8" + + @classmethod + def get_backend_name(cls) -> str: + return ARM_BACKEND diff --git a/backends/arm/requirements-arm-models-test.txt b/backends/arm/requirements-arm-models-test.txt index c6a1d94aef2..f59928bae53 100644 --- a/backends/arm/requirements-arm-models-test.txt +++ b/backends/arm/requirements-arm-models-test.txt @@ -8,4 +8,5 @@ diffusers[torch] @ git+https://github.com/huggingface/diffusers.git@a7cb14efbe4b pydantic slangtorch rich +torcheval==0.0.7 setuptools==80.10.2 diff --git a/backends/arm/runtime/EthosUBackend.cpp b/backends/arm/runtime/EthosUBackend.cpp index 1305c5b4995..185623d8504 100644 --- a/backends/arm/runtime/EthosUBackend.cpp +++ b/backends/arm/runtime/EthosUBackend.cpp @@ -56,16 +56,34 @@ namespace arm { extern "C" { void __attribute__((weak)) EthosUBackend_execute_begin() {} void __attribute__((weak)) EthosUBackend_execute_end() {} +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) +void __attribute__((weak)) EthosUBackend_delegate_begin(const void*) {} +void __attribute__((weak)) EthosUBackend_delegate_end() {} +#endif +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) +void __attribute__((weak)) EthosUBackend_input_memcpy(size_t) {} +void __attribute__((weak)) EthosUBackend_output_memcpy(size_t) {} +#endif __attribute__((weak)) unsigned char* ethosu_fast_scratch = nullptr; __attribute__((weak)) size_t ethosu_fast_scratch_size = 0; } class EthosUBackendExecuteCallbacks { public: +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + explicit EthosUBackendExecuteCallbacks(const void* handle) { + EthosUBackend_execute_begin(); + EthosUBackend_delegate_begin(handle); + } +#else EthosUBackendExecuteCallbacks() { EthosUBackend_execute_begin(); } +#endif ~EthosUBackendExecuteCallbacks() { +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + EthosUBackend_delegate_end(); +#endif EthosUBackend_execute_end(); } }; @@ -107,10 +125,11 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { } MemoryAllocator* allocator = context.get_runtime_allocator(); - ExecutionHandle* handle = new (std::nothrow) ExecutionHandle(); + ExecutionHandle* handle = allocator->allocateInstance(); if (handle == nullptr) { return Error::MemoryAllocationFailed; } + new (handle) ExecutionHandle(); EXECUTORCH_PROF_START( event_tracer, @@ -120,11 +139,16 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { data, size, context.get_named_data_map(), &handle->handles); EXECUTORCH_PROF_END(event_tracer, event_tracer_local_scope); if (read_status != Error::Ok) { - delete handle; + handle->~ExecutionHandle(); return read_status; } - handle->platform_state = platform_init(compile_specs, allocator); + const Error platform_status = + platform_init(compile_specs, allocator, handle); + if (platform_status != Error::Ok) { + delete handle; + return platform_status; + } // Return the same buffer we were passed - this data will be // executed directly @@ -152,7 +176,11 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { // and EthosUBackend_execute_end() is called while CollectArm_CPU_Cycles is // in scope. e.g. We meassure from now until we exit this metod (in any way // we might do it). +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + EthosUBackendExecuteCallbacks CollectArm_CPU_Cycles(input_handle); +#else EthosUBackendExecuteCallbacks CollectArm_CPU_Cycles; +#endif ExecutionHandle* execution_handle = static_cast(input_handle); @@ -250,6 +278,9 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { event_tracer, "+EthosUBackend::execute()handles.input.memcpy()"); // Sizes match and elt size matches so memcpy. // Routed through arm_ethos_io_memcpy so firmware can DMA-accelerate. +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + EthosUBackend_input_memcpy(tensor_in.nbytes()); +#endif arm_ethos_io_memcpy( scratch_addr, tensor_in.mutable_data_ptr(), @@ -299,7 +330,7 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { platform_destroy(exec_handle->platform_state); } - delete exec_handle; + exec_handle->~ExecutionHandle(); } private: @@ -404,6 +435,9 @@ Error copy_with_layout_adjustment( const char* src_bytes = src; for (size_t chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) { // Routed through arm_ethos_io_memcpy so firmware can DMA-accelerate. +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + EthosUBackend_output_memcpy(chunk_size); +#endif arm_ethos_io_memcpy(dest, src_bytes, chunk_size); src_bytes += vela_chunk_size; dest += chunk_size; diff --git a/backends/arm/runtime/EthosUBackend_Cortex_A.cpp b/backends/arm/runtime/EthosUBackend_Cortex_A.cpp index 41c1bca97bf..ef5009c25fb 100644 --- a/backends/arm/runtime/EthosUBackend_Cortex_A.cpp +++ b/backends/arm/runtime/EthosUBackend_Cortex_A.cpp @@ -49,6 +49,11 @@ struct LinuxDriverOptions { struct PlatformState { LinuxDriverOptions options; + std::shared_ptr network; + std::shared_ptr constant_buffer; + std::shared_ptr intermediate_buffer; + std::vector> ifm_buffers; + std::vector> ofm_buffers; }; namespace { @@ -181,35 +186,9 @@ Error invoke_linux_driver( const std::vector& output_ptrs, const std::vector& input_copy_sizes, const std::vector& output_copy_sizes, - const LinuxDriverOptions& options) { - if (handles.outputs == nullptr) { - ET_LOG(Error, "Ethos-U backend missing output metadata"); - return Error::InvalidProgram; - } + const PlatformState& state) { + const LinuxDriverOptions& options = state.options; try { - EthosU::Device& device = get_linux_device_cache().get(options.device_path); - auto network = std::make_shared( - device, - reinterpret_cast(handles.cmd_data), - handles.cmd_data_size); - - std::shared_ptr constant_buffer = - std::make_shared(); - if (handles.weight_data_size > 0) { - auto constant_buffers = device.createBuffers({handles.weight_data_size}); - constant_buffer = constant_buffers.front(); - constant_buffer->write( - const_cast(handles.weight_data), handles.weight_data_size); - } - - std::shared_ptr intermediate_buffer = - std::make_shared(); - if (handles.scratch_data_size > 0) { - auto scratch_buffers = device.createBuffers({handles.scratch_data_size}); - intermediate_buffer = scratch_buffers.front(); - } - - std::vector> ifm_buffers; if (handles.inputs != nullptr && handles.inputs->count > 0) { if (input_copy_sizes.size() != static_cast(handles.inputs->count)) { @@ -228,7 +207,6 @@ Error invoke_linux_driver( input_copy_sizes.size()); return Error::InvalidState; } - ifm_buffers = device.createBuffers(input_copy_sizes); for (int i = 0; i < handles.inputs->count; ++i) { const size_t copy_size = input_copy_sizes[i]; if (copy_size == 0) { @@ -239,7 +217,7 @@ Error invoke_linux_driver( ET_LOG(Error, "Missing input buffer for index %d", i); return Error::InvalidState; } - ifm_buffers[i]->write(const_cast(src), copy_size); + state.ifm_buffers[i]->write(const_cast(src), copy_size); } } @@ -260,16 +238,14 @@ Error invoke_linux_driver( output_copy_sizes.size()); return Error::InvalidState; } - auto ofm_buffers = device.createBuffers(output_copy_sizes); - auto inference = std::make_unique( - network, - ifm_buffers.begin(), - ifm_buffers.end(), - ofm_buffers.begin(), - ofm_buffers.end(), - intermediate_buffer, - constant_buffer, + state.network, + state.ifm_buffers.begin(), + state.ifm_buffers.end(), + state.ofm_buffers.begin(), + state.ofm_buffers.end(), + state.intermediate_buffer, + state.constant_buffer, options.pmu_events, options.enable_cycle_counter); @@ -311,7 +287,7 @@ Error invoke_linux_driver( ET_LOG(Error, "Missing output buffer for index %d", i); return Error::InvalidState; } - ofm_buffers[i]->read(dst, copy_size); + state.ofm_buffers[i]->read(dst, copy_size); } } catch (const std::exception& e) { ET_LOG(Error, "Ethos-U Linux driver invocation failed: %s", e.what()); @@ -320,20 +296,85 @@ Error invoke_linux_driver( return Error::Ok; } + +// Get the byte size of an IO tensor from its Vela descriptor. +size_t vela_io_bytes(const VelaIO& io) { + size_t count = 1; + for (int i = 0; i < shapeDim; i++) { + count *= static_cast(io.shape[i]); + } + return count * static_cast(io.elem_size); +} + +// Created once in platform_init(), reused by every invoke_linux_driver(). +Error create_driver_objects(const VelaHandles& handles, PlatformState* state) { + if (handles.outputs == nullptr) { + ET_LOG(Error, "Ethos-U backend missing output metadata"); + return Error::InvalidProgram; + } + const LinuxDriverOptions& options = state->options; + try { + EthosU::Device& device = get_linux_device_cache().get(options.device_path); + state->network = std::make_shared( + device, + reinterpret_cast(handles.cmd_data), + handles.cmd_data_size); + + state->constant_buffer = std::make_shared(); + if (handles.weight_data_size > 0) { + auto constant_buffers = device.createBuffers({handles.weight_data_size}); + state->constant_buffer = constant_buffers.front(); + state->constant_buffer->write( + const_cast(handles.weight_data), handles.weight_data_size); + } + + state->intermediate_buffer = std::make_shared(); + if (handles.scratch_data_size > 0) { + auto scratch_buffers = device.createBuffers({handles.scratch_data_size}); + state->intermediate_buffer = scratch_buffers.front(); + } + + if (handles.inputs != nullptr && handles.inputs->count > 0) { + std::vector ifm_sizes; + for (int i = 0; i < handles.inputs->count; ++i) { + ifm_sizes.push_back(vela_io_bytes(handles.inputs->io[i])); + } + state->ifm_buffers = device.createBuffers(ifm_sizes); + } + + std::vector ofm_sizes; + for (int i = 0; i < handles.outputs->count; ++i) { + ofm_sizes.push_back(vela_io_bytes(handles.outputs->io[i])); + } + state->ofm_buffers = device.createBuffers(ofm_sizes); + } catch (const std::exception& e) { + ET_LOG(Error, "Ethos-U Linux driver setup failed: %s", e.what()); + return Error::InvalidState; + } + + return Error::Ok; +} } // namespace // Used by EthosUBackend.cpp through EthosUBackend_Internal.h. // cppcheck-suppress unusedFunction -PlatformState* platform_init( +Error platform_init( ArrayRef specs, - MemoryAllocator* allocator) { + MemoryAllocator* allocator, + ExecutionHandle* handle) { (void)allocator; PlatformState* state = new (std::nothrow) PlatformState(); if (state == nullptr) { - return nullptr; + return Error::MemoryAllocationFailed; } state->options = parse_linux_options(specs); - return state; + const Error status = create_driver_objects(handle->handles, state); + if (status != Error::Ok) { + delete state; + return status; + } + handle->platform_state = state; + return Error::Ok; } // Used by EthosUBackend.cpp through EthosUBackend_Internal.h. @@ -401,7 +442,7 @@ Error platform_execute( linux_output_ptrs, input_copy_sizes, output_io_bytes, - state->options); + *state); if (status != Error::Ok) { return status; } diff --git a/backends/arm/runtime/EthosUBackend_Cortex_M.cpp b/backends/arm/runtime/EthosUBackend_Cortex_M.cpp index 82cbe99afad..0697b783a36 100644 --- a/backends/arm/runtime/EthosUBackend_Cortex_M.cpp +++ b/backends/arm/runtime/EthosUBackend_Cortex_M.cpp @@ -13,7 +13,6 @@ #include #include #include -#include #include @@ -54,15 +53,14 @@ namespace arm { struct PlatformState {}; -PlatformState* platform_init( +executorch::runtime::Error platform_init( executorch::runtime::ArrayRef /*specs*/, - executorch::runtime::MemoryAllocator* /*allocator*/) { - return nullptr; + executorch::runtime::MemoryAllocator* /*allocator*/, + ExecutionHandle* /*handle*/) { + return executorch::runtime::Error::Ok; } -void platform_destroy(PlatformState* state) { - delete state; -} +void platform_destroy(PlatformState* /*state*/) {} bool needs_scratch_allocation() { return true; @@ -152,6 +150,9 @@ Error platform_execute( io_bytes_total += tensor_bytes; } else { // Routed through arm_ethos_io_memcpy so firmware can DMA-accelerate. +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + EthosUBackend_output_memcpy(tensor_bytes); +#endif arm_ethos_io_memcpy( tensor_out.mutable_data_ptr(), static_cast(output_addr), diff --git a/backends/arm/runtime/EthosUBackend_Internal.h b/backends/arm/runtime/EthosUBackend_Internal.h index 48fc4aa3a79..a62926676b3 100644 --- a/backends/arm/runtime/EthosUBackend_Internal.h +++ b/backends/arm/runtime/EthosUBackend_Internal.h @@ -74,13 +74,22 @@ struct ExecutionHandle { extern "C" { void EthosUBackend_execute_begin(); void EthosUBackend_execute_end(); +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) +void EthosUBackend_delegate_begin(const void* handle); +void EthosUBackend_delegate_end(); +#endif +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) +void EthosUBackend_input_memcpy(size_t size); +void EthosUBackend_output_memcpy(size_t size); +#endif extern unsigned char* ethosu_fast_scratch; extern size_t ethosu_fast_scratch_size; } -PlatformState* platform_init( +executorch::runtime::Error platform_init( executorch::runtime::ArrayRef specs, - executorch::runtime::MemoryAllocator* allocator); + executorch::runtime::MemoryAllocator* allocator, + ExecutionHandle* handle); void platform_destroy(PlatformState* state); diff --git a/backends/arm/runtime/VGFBackend.cpp b/backends/arm/runtime/VGFBackend.cpp index c7b735376a4..3365cd259cf 100644 --- a/backends/arm/runtime/VGFBackend.cpp +++ b/backends/arm/runtime/VGFBackend.cpp @@ -58,6 +58,8 @@ using executorch::runtime::EventTracerEntry; // We use the platform and runtime environment provided by the Vulkan delegate #include +#include + // Dependencies for processing VGF files into Vulkan calls #include #include @@ -953,25 +955,31 @@ VkResult vkml_allocate_basics( }; // Query features + VkPhysicalDeviceShaderBfloat16FeaturesKHR available_bfloat16{ + .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_BFLOAT16_FEATURES_KHR, + .pNext = nullptr, + }; VkPhysicalDeviceVulkan12Features available_12 = { .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, - .pNext = NULL, + .pNext = &available_bfloat16, }; VkPhysicalDeviceVulkan11Features available_11 = { .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, .pNext = &available_12, }; + VkPhysicalDeviceDataGraphFeaturesARM available_graph{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DATA_GRAPH_FEATURES_ARM, &available_11}; #if defined(VK_ARM_data_graph_neural_accelerator_statistics) VkPhysicalDeviceDataGraphNeuralAcceleratorStatisticsFeaturesARM available_neural_statistics{ .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DATA_GRAPH_NEURAL_ACCELERATOR_STATISTICS_FEATURES_ARM, - .pNext = &available_11, + .pNext = &available_graph, .dataGraphNeuralAcceleratorStatistics = VK_FALSE, }; void* available_features_pnext = &available_neural_statistics; #else - void* available_features_pnext = &available_11; + void* available_features_pnext = &available_graph; #endif VkPhysicalDeviceFeatures2 available_2 = { .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2, @@ -979,7 +987,23 @@ VkResult vkml_allocate_basics( }; vkGetPhysicalDeviceFeatures2(*physical_device, &available_2); + if (!vgf_data_graph_features_supported(available_graph)) { + ET_LOG( + Error, + "VGF requires VK_ARM_data_graph features dataGraph and " + "dataGraphShaderModule (reported dataGraph=%u, " + "dataGraphShaderModule=%u)", + available_graph.dataGraph, + available_graph.dataGraphShaderModule); + return VK_ERROR_FEATURE_NOT_PRESENT; + } + // Select features + VkPhysicalDeviceShaderBfloat16FeaturesKHR features_bfloat16{ + .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_BFLOAT16_FEATURES_KHR, + .pNext = nullptr, + .shaderBFloat16Type = VK_FALSE, + }; VkPhysicalDeviceShaderReplicatedCompositesFeaturesEXT features_c{ VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_REPLICATED_COMPOSITES_FEATURES_EXT, nullptr}; @@ -1013,10 +1037,8 @@ VkResult vkml_allocate_basics( features_tensor.shaderTensorAccess = true; features_tensor.tensors = true; features_tensor.pNext = &features_11; - VkPhysicalDeviceDataGraphFeaturesARM features_graph{ - VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DATA_GRAPH_FEATURES_ARM, nullptr}; - features_graph.dataGraph = true; - features_graph.pNext = &features_tensor; + VkPhysicalDeviceDataGraphFeaturesARM features_graph = + make_vgf_data_graph_features(&features_tensor); #if defined(VK_ARM_data_graph_neural_accelerator_statistics) VkPhysicalDeviceDataGraphNeuralAcceleratorStatisticsFeaturesARM features_neural_statistics{ @@ -1049,6 +1071,34 @@ VkResult vkml_allocate_basics( vector requested_exts; + const bool bfloat16_extension_available = std::any_of( + available.begin(), available.end(), [](const auto& ext_avail) { + return std::strcmp( + VK_KHR_SHADER_BFLOAT16_EXTENSION_NAME, + ext_avail.extensionName) == 0; + }); + const bool bfloat16_feature_available = + available_bfloat16.shaderBFloat16Type == VK_TRUE; + + if (bfloat16_extension_available && bfloat16_feature_available) { + requested_exts.push_back(VK_KHR_SHADER_BFLOAT16_EXTENSION_NAME); + features_bfloat16.shaderBFloat16Type = VK_TRUE; + features_c.pNext = &features_bfloat16; + ET_LOG( + Info, + "Enabled %s with shaderBFloat16Type", + VK_KHR_SHADER_BFLOAT16_EXTENSION_NAME); + } else if (!bfloat16_extension_available) { + ET_LOG( + Info, + "VGF BF16 shaders are unavailable: Vulkan device does not expose %s", + VK_KHR_SHADER_BFLOAT16_EXTENSION_NAME); + } else { + ET_LOG( + Info, + "VGF BF16 shaders are unavailable: shaderBFloat16Type is not supported"); + } + #if defined(VK_ARM_data_graph_neural_accelerator_statistics) const bool neural_statistics_extension_available = std::any_of( available.begin(), available.end(), [](const auto& ext_avail) { diff --git a/backends/arm/runtime/VGFSetup.cpp b/backends/arm/runtime/VGFSetup.cpp index 9fca73d3551..4adcb406666 100644 --- a/backends/arm/runtime/VGFSetup.cpp +++ b/backends/arm/runtime/VGFSetup.cpp @@ -2668,9 +2668,44 @@ bool VgfRepr::process_vgf( } } + // Keep the pipeline cache local to this VgfRepr. This preserves reuse across + // the many segment pipelines in one VGF without sharing mutable cache state + // or cache lifetime across independently initialized delegate handles. + if (vk_pipeline_cache == VK_NULL_HANDLE) { + VkPipelineCacheCreateInfo pipeline_cache_info{ + .sType = VK_STRUCTURE_TYPE_PIPELINE_CACHE_CREATE_INFO, + .pNext = nullptr, + // Default mode: Vulkan internally synchronizes concurrent + // pipeline-cache access. The cache is per-VgfRepr anyway, so + // independent delegates do not contend on it. + .flags = 0, + .initialDataSize = 0, + .pInitialData = nullptr, + }; + + { + VGF_PROFILE_SCOPE(event_tracer, "VGF_INIT_CREATE_PIPELINE_CACHE"); + result = vkCreatePipelineCache( + vk_device, &pipeline_cache_info, nullptr, &vk_pipeline_cache); + } + + if (result != VK_SUCCESS) { + ET_LOG( + Info, + "Failed to create optional per-VgfRepr Vulkan pipeline cache, " + "error 0x%08X; continuing without pipeline caching", + result); + vk_pipeline_cache = VK_NULL_HANDLE; + } else { + ET_LOG(Info, "VGF per-VgfRepr Vulkan pipeline cache enabled"); + } + } + // Build per-segment pipelines and descriptor sets. segments.clear(); segments.reserve(segment_count); + size_t graph_segment_count = 0; + size_t compute_segment_count = 0; { VGF_PROFILE_SCOPE(event_tracer, "VGF_INIT_BUILD_SEGMENTS"); @@ -2683,6 +2718,12 @@ bool VgfRepr::process_vgf( return false; } + if (segment_type == vgflib::ModuleType::GRAPH) { + ++graph_segment_count; + } else { + ++compute_segment_count; + } + SegmentState segment; segment.segment_id = segment_id; segment.use_data_graph_pipeline = @@ -3137,7 +3178,7 @@ bool VgfRepr::process_vgf( result = vkCreateDataGraphPipelinesARM( vk_device, VK_NULL_HANDLE, - VK_NULL_HANDLE, + vk_pipeline_cache, 1, &graph_pipeline_info, nullptr, @@ -3429,7 +3470,7 @@ bool VgfRepr::process_vgf( VGF_PROFILE_SCOPE(event_tracer, "VGF_INIT_CREATE_COMPUTE_PIPELINE"); result = vkCreateComputePipelines( vk_device, - VK_NULL_HANDLE, + vk_pipeline_cache, 1, &compute_info, nullptr, @@ -3445,6 +3486,14 @@ bool VgfRepr::process_vgf( } } + ET_LOG( + Info, + "VGF segment counts: total=%d graph=%zu compute=%zu pipeline_cache=%s", + segment_count, + graph_segment_count, + compute_segment_count, + vk_pipeline_cache != VK_NULL_HANDLE ? "enabled" : "disabled"); + // Map model sequence inputs/outputs to IO indices auto input_handle = sequence_decoder->getModelSequenceInputBindingSlotsHandle(); diff --git a/backends/arm/runtime/VGFSetup.h b/backends/arm/runtime/VGFSetup.h index e4b60bbbf94..b87f648962c 100644 --- a/backends/arm/runtime/VGFSetup.h +++ b/backends/arm/runtime/VGFSetup.h @@ -177,6 +177,12 @@ class VgfRepr { ~VgfRepr() { free_vgf(); + if (vk_pipeline_cache != VK_NULL_HANDLE) { + // The cache is private to this VgfRepr, so no other delegate instance can + // be accessing it while this object is being destroyed. + vkDestroyPipelineCache(vk_device, vk_pipeline_cache, nullptr); + vk_pipeline_cache = VK_NULL_HANDLE; + } } private: @@ -188,6 +194,12 @@ class VgfRepr { VkCommandPool vk_command_pool; uint32_t vk_queue_family_index = UINT32_MAX; + // Owned by this VgfRepr. One cache is reused across all graph and compute + // segments in this loaded VGF, but is not shared with independent VgfRepr + // instances. flags=0 uses Vulkan's default internally synchronized cache + // mode. + VkPipelineCache vk_pipeline_cache = VK_NULL_HANDLE; + bool neural_statistics_requested_ = false; bool neural_statistics_device_enabled_ = false; int neural_statistics_mode_index_ = 1; diff --git a/backends/arm/runtime/VGFVulkanFeatures.h b/backends/arm/runtime/VGFVulkanFeatures.h new file mode 100644 index 00000000000..8e88afc8f32 --- /dev/null +++ b/backends/arm/runtime/VGFVulkanFeatures.h @@ -0,0 +1,34 @@ +/* + * 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. + */ + +#pragma once + +#include + +namespace executorch { +namespace backends { +namespace vgf { + +inline VkPhysicalDeviceDataGraphFeaturesARM make_vgf_data_graph_features( + void* p_next) { + VkPhysicalDeviceDataGraphFeaturesARM features{}; + features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DATA_GRAPH_FEATURES_ARM; + features.pNext = p_next; + features.dataGraph = VK_TRUE; + features.dataGraphShaderModule = VK_TRUE; + return features; +} + +inline bool vgf_data_graph_features_supported( + const VkPhysicalDeviceDataGraphFeaturesARM& features) { + return features.dataGraph == VK_TRUE && + features.dataGraphShaderModule == VK_TRUE; +} + +} // namespace vgf +} // namespace backends +} // namespace executorch diff --git a/backends/arm/runtime/targets.bzl b/backends/arm/runtime/targets.bzl index 0ba2dec3994..230e95e8fba 100644 --- a/backends/arm/runtime/targets.bzl +++ b/backends/arm/runtime/targets.bzl @@ -51,6 +51,7 @@ def define_common_targets(): exported_headers = [ "VGFNeuralStatistics.h", "VGFSetup.h", + "VGFVulkanFeatures.h", ], # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole) link_whole = True, diff --git a/backends/arm/scripts/aot_arm_compiler.py b/backends/arm/scripts/aot_arm_compiler.py index 250606f9a4b..5c89c79a960 100644 --- a/backends/arm/scripts/aot_arm_compiler.py +++ b/backends/arm/scripts/aot_arm_compiler.py @@ -626,6 +626,14 @@ def _get_args(): choices=TARGETS, help=f"Target backend. For delegated models: Ethos-U/VGF/TOSA variants. For non-delegated: cortex-m (CMSIS-NN portable kernels). Valid targets: {TARGETS}", ) + parser.add_argument( + "--cortex-m-explicit-layout", + action="store_true", + help=( + "Use explicit NCHW/NHWC permutes for Cortex-M instead of dim-order " + "operators. This is an experimental Cortex-M-only option." + ), + ) # TODO: Remove --evaluate and --evaluate_config completely after a suitable time. # They are deprecated and no longer functional in this script. parser.add_argument( @@ -923,9 +931,16 @@ def _to_edge_cortex_m( """Cortex-M/CMSIS-NN compilation path with no delegation.""" logging.info( f"Using Cortex-M/CMSIS-NN compilation path for cpu={target_config.cpu.name} " - f"backend={target_config.backend.name}" + f"backend={target_config.backend.name} " + f"layout={'explicit' if args.cortex_m_explicit_layout else 'dim-order'}" ) + if args.cortex_m_explicit_layout and not args.quantize: + raise RuntimeError( + "--cortex-m-explicit-layout requires --quantize; explicit layout " + "does not fall back to portable float spatial operators." + ) + def _to_channels_last(x): if isinstance(x, torch.Tensor): if x.dim() == 4: @@ -949,17 +964,24 @@ def _to_channels_last(x): ) model_quant = None else: - model = model.to(memory_format=torch.channels_last) # type: ignore[call-overload] - example_inputs = tuple(_to_channels_last(x) for x in example_inputs) + if not args.cortex_m_explicit_layout: + model = model.to(memory_format=torch.channels_last) # type: ignore[call-overload] + example_inputs = tuple(_to_channels_last(x) for x in example_inputs) + # Refresh fake-tensor strides after changing the captured module's + # memory format so the legacy quantizer sees channels-last inputs. + model = torch.export.export( + model, example_inputs, strict=args.strict_export + ).module() + + quantizer = CortexMQuantizer(use_explicit_layout=args.cortex_m_explicit_layout) - quantizer = CortexMQuantizer() prepared = prepare_pt2e(model, quantizer) if calibration_samples is None: calibration_samples = [example_inputs] for sample in calibration_samples: - prepared(*tuple(_to_channels_last(x) for x in sample)) + prepared(*sample) model_quant = convert_pt2e(prepared) @@ -973,7 +995,9 @@ def _to_channels_last(x): ) pass_manager = CortexMPassManager( - edge.exported_program(), target_config=target_config + edge.exported_program(), + target_config=target_config, + use_explicit_layout=args.cortex_m_explicit_layout, ) edge._edge_programs["forward"] = pass_manager.transform() diff --git a/backends/arm/scripts/corstone_utils.cmake b/backends/arm/scripts/corstone_utils.cmake index 72a0c27a8c4..2dba2052555 100644 --- a/backends/arm/scripts/corstone_utils.cmake +++ b/backends/arm/scripts/corstone_utils.cmake @@ -32,14 +32,9 @@ function(fetch_ethos_u_content ETHOS_SDK_PATH ET_DIR_PATH) GIT_REPOSITORY https://git.gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u.git GIT_TAG ${ethos_u_base_tag} - SOURCE_DIR - ${ETHOS_SDK_PATH} - BINARY_DIR - ${ETHOS_SDK_PATH} - SUBBUILD_DIR - ${ETHOS_SDK_PATH}/../ethos_u-subbuild - SOURCE_SUBDIR - none + SOURCE_DIR ${ETHOS_SDK_PATH} BINARY_DIR ${ETHOS_SDK_PATH} + # Keep the generator-specific population project local to this build. + SOURCE_SUBDIR none ) FetchContent_MakeAvailable(ethos_u) # Patch manifest to remove unused projects. @@ -126,7 +121,7 @@ function(get_corstone_linker_script OUT_VAR SYSTEM_CONFIG) ) endfunction() -function(add_corstone_subdirectory SYSTEM_CONFIG ETHOS_SDK_PATH) +function(add_corstone_subdirectory SYSTEM_CONFIG ETHOS_SDK_PATH MEMORY_MODE) if(MEMORY_MODE MATCHES "^Dedicated_Sram($|_)") # Both model and scratch in DRAM. set(MEMORY_MODEL dram) diff --git a/backends/arm/scripts/docgen/generate_vgf_op_support.py b/backends/arm/scripts/docgen/generate_vgf_op_support.py index b2068d1d157..a4df1ba55d3 100644 --- a/backends/arm/scripts/docgen/generate_vgf_op_support.py +++ b/backends/arm/scripts/docgen/generate_vgf_op_support.py @@ -54,6 +54,10 @@ BACKEND_PIPELINE_CLASS_NAMES = frozenset({"VgfPipeline"}) BACKEND_PIPELINE_LABEL = "VgfPipeline" BACKEND_TOSA_SPEC = "TOSA-1.0+FP+INT+int4+int16" +BACKEND_PROFILE_TOSA_SPECS = { + "FP": "TOSA-1.0+FP", + "INT": "TOSA-1.0+INT", +} GENERATOR_PATH = Path("backends/arm/scripts/docgen/generate_vgf_op_support.py") GENERATOR_COMMAND = f"python {GENERATOR_PATH}" @@ -2197,6 +2201,37 @@ def _profiles_for_checker( return profiles +def _collect_backend_custom_partition_ops( + backend_tosa_spec: TosaSpecificationLike, +) -> dict[str, set[object]]: + """Collect VGF custom partition ops for each enabled support profile. + + Instantiate the partitioner with a single-profile compile spec so custom + registrations that are conditional on the compile spec are attributed only + to the profiles for which they are actually registered. + + """ + from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner + + enabled_profiles = { + "FP": backend_tosa_spec.support_float(), + "INT": backend_tosa_spec.support_integer(), + } + custom_ops_by_profile: dict[str, set[object]] = {} + + for profile, enabled in enabled_profiles.items(): + if not enabled: + continue + partitioner = VgfPartitioner( + VgfCompileSpec(BACKEND_PROFILE_TOSA_SPECS[profile]) + ) + custom_ops_by_profile[profile] = set( + getattr(partitioner, "_custom_partition_ops", ()) + ) + + return custom_ops_by_profile + + def _collect_backend_supported_ops( # noqa: C901 repo_root: Path, ) -> dict[str, SupportedOperatorEvidence]: @@ -2248,6 +2283,10 @@ def add(target: object, profile: str, evidence: str) -> None: for profile in _profiles_for_checker(checker, tosa_spec): add(target, profile, checker_evidence) + for profile, targets in _collect_backend_custom_partition_ops(tosa_spec).items(): + for target in targets: + add(target, profile, "VgfPartitioner.register_custom_partition_op") + # Lowering visitors are not the source of partitioner support, but they are # useful evidence when the exported op name matches a registered visitor # target directly. @@ -2669,7 +2708,7 @@ def run_check(repo_root: Path, *, strict_ast: bool = False) -> int: # noqa: C90 print(f"| `{op}` | {profile} | {classification} | {sat} | {test_cell} |") print() - if unresolved: + if strict_ast and unresolved: _print_unresolved(unresolved) if diagnostics: print("AST normalisation diagnostics:") diff --git a/backends/arm/scripts/fvp_utils.sh b/backends/arm/scripts/fvp_utils.sh index 73f67112efd..7c66845d908 100644 --- a/backends/arm/scripts/fvp_utils.sh +++ b/backends/arm/scripts/fvp_utils.sh @@ -138,8 +138,12 @@ function setup_path_fvp() { # Fixup for Corstone-320 python dependency append_env_in_setup_path LD_LIBRARY_PATH "${root_dir}/FVP-corstone320/python/lib/" - echo "hash FVP_Corstone_SSE-300_Ethos-U55" >> ${setup_path_script}.sh - echo "hash FVP_Corstone_SSE-300_Ethos-U65" >> ${setup_path_script}.sh - echo "hash FVP_Corstone_SSE-320" >> ${setup_path_script}.sh - echo "hash FVP_Corstone-1000-A320" >> ${setup_path_script}.sh + local fvp_command + for fvp_command in \ + FVP_Corstone_SSE-300_Ethos-U55 \ + FVP_Corstone_SSE-300_Ethos-U65 \ + FVP_Corstone_SSE-320 \ + FVP_Corstone-1000-A320; do + echo "hash ${fvp_command} 2>/dev/null || true" >> "${setup_path_script}.sh" + done } diff --git a/backends/arm/scripts/generate_neural_graphics_test_data.py b/backends/arm/scripts/generate_neural_graphics_test_data.py new file mode 100644 index 00000000000..67bfa2a8754 --- /dev/null +++ b/backends/arm/scripts/generate_neural_graphics_test_data.py @@ -0,0 +1,435 @@ +# 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. +"""Generate NSS autoencoder calibration and verification data.""" + +from __future__ import annotations + +import json +import os +import shutil +from importlib.resources import files +from pathlib import Path + +import torch +import torch.nn.functional as F +from executorch.backends.arm.scripts.neural_graphics_test_data import ( + nss_input_shape, + nss_test_calibration_path, + nss_test_data_root, + nss_test_verification_path, +) +from safetensors.torch import save_file + +os.environ.setdefault("HF_HUB_DISABLE_XET", "1") + +from huggingface_hub import snapshot_download +from ng_model_gym.core.config.config_model import ( # type: ignore[import-not-found,import-untyped] + ConfigModel, +) +from ng_model_gym.core.data.data_utils import ( # type: ignore[import-not-found,import-untyped] + DataLoaderMode, + DatasetType, + tonemap_forward, + ToneMapperMode, +) +from ng_model_gym.usecases.nss.data.dataset import ( # type: ignore[import-not-found,import-untyped] + NSSDataset, +) + + +_DATASET_REPO_ID = "Arm/neural-graphics-dataset" +_CALIBRATION_SOURCE_ALLOW_PATTERNS = [ + "train/**/*.safetensors", + "nss/train/**/*.safetensors", +] +_EVALUATION_SOURCE_ALLOW_PATTERNS = [ + "test/test_full_resolution_sample.safetensors", + "nss/test/test_full_resolution_sample.safetensors", +] + +EPS = 1e-7 +NSS_V1_SPATIAL_MULTIPLE = 8 + + +def _luminance(rgb: torch.Tensor) -> torch.Tensor: + weights = torch.tensor( + [0.2126, 0.7152, 0.0722], + dtype=rgb.dtype, + device=rgb.device, + ).view(1, 3, 1, 1) + return torch.sum(rgb * weights, dim=1, keepdim=True) + + +def _motion_detector( + motion_lr: torch.Tensor, render_size: torch.Tensor +) -> torch.Tensor: + # render_size is stored as [height, width], matching the dataset writer. + size = render_size.to(dtype=torch.float32).view(-1, 2, 1, 1) + motion_norm = motion_lr.to(dtype=torch.float32) / torch.clamp(size, min=1.0) + motion_length = torch.linalg.vector_norm(motion_norm, dim=1, keepdim=True) + + pix_min = torch.linalg.vector_norm( + 1.0 / torch.clamp(size, min=1.0), dim=1 + ).unsqueeze(1) + pix_max = torch.linalg.vector_norm( + 200.0 / torch.clamp(size, min=1.0), dim=1 + ).unsqueeze(1) + detector = (torch.clamp(motion_length, pix_min, pix_max) - pix_min) / torch.clamp( + pix_max - pix_min, min=EPS + ) + return torch.sqrt(torch.clamp(detector, min=0.0)) + + +def _depth_edge(depth: torch.Tensor) -> torch.Tensor: + dx = F.pad(torch.abs(depth[..., :, 1:] - depth[..., :, :-1]), (0, 1, 0, 0)) + dy = F.pad(torch.abs(depth[..., 1:, :] - depth[..., :-1, :]), (0, 0, 0, 1)) + return torch.clamp((dx + dy) * 100.0, 0.0, 1.0) + + +def _reflect_pad_to_multiple( + tensor: torch.Tensor, + multiple: int = NSS_V1_SPATIAL_MULTIPLE, +) -> torch.Tensor: + h, w = tensor.shape[-2:] + pad_h = (multiple - (h % multiple)) % multiple + pad_w = (multiple - (w % multiple)) % multiple + if pad_h == 0 and pad_w == 0: + return tensor + return F.pad(tensor, (0, pad_w, 0, pad_h), mode="reflect") + + +def _model_gym_dataset(src: Path) -> NSSDataset: + config_path = files("ng_model_gym.usecases.nss.configs").joinpath( + "nss_v1_template.json" + ) + config = json.loads(config_path.read_text(encoding="utf-8")) + for split in ("train", "validation", "test"): + config["dataset"]["path"][split] = str(src) + config["dataset"].update( + exposure=None, + tonemapper=ToneMapperMode.KARIS.value, + gt_augmentation=False, + ) + params = ConfigModel.model_validate(config) + return NSSDataset(params, DataLoaderMode.TEST, DatasetType.SAFETENSOR) + + +def _make_autoencoder_input( + current: dict[str, torch.Tensor], + previous: dict[str, torch.Tensor] | None = None, +) -> torch.Tensor: + colour_tm = current["colour"] + exposure = current["exposure"] + _, _, h, w = colour_tm.shape + same_sequence = previous is not None and torch.equal( + current["seq"], previous["seq"] + ) + + history_linear = torch.zeros_like(current["colour_linear"]) + if previous is not None and same_sequence: + history_linear = previous["ground_truth_linear"] + if history_linear.shape[-2:] != (h, w): + history_linear = F.interpolate( + history_linear, + size=(h, w), + mode="bilinear", + align_corners=False, + ) + history_tm = tonemap_forward(history_linear * exposure, mode=ToneMapperMode.KARIS) + + motion_signal = _motion_detector(current["motion_lr"], current["render_size"]) + luma = _luminance(colour_tm) + previous_luma = torch.zeros_like(luma) + feedback = torch.zeros((1, 4, h, w), dtype=torch.float32) + if previous is not None and same_sequence: + previous_luma = _luminance(previous["colour"]) + previous_luma_derivative = torch.clamp(previous_luma, 0.0, 1.0) + feedback[:, 0:1] = _motion_detector( + previous["motion_lr"], previous["render_size"] + ) + feedback[:, 1:2] = previous_luma_derivative + feedback[:, 2:3] = previous_luma + feedback[:, 3:4] = _depth_edge(previous["depth"]) + luma_derivative = torch.clamp(torch.abs(luma - previous_luma), 0.0, 1.0) + + autoencoder_input = torch.cat( + [ + _reflect_pad_to_multiple(history_tm), + _reflect_pad_to_multiple(colour_tm), + _reflect_pad_to_multiple(motion_signal), + _reflect_pad_to_multiple(feedback), + _reflect_pad_to_multiple(luma_derivative), + ], + dim=1, + ) + return autoencoder_input.to(torch.float16) + + +def _autoencoder_sample(dataset: NSSDataset, index: int) -> torch.Tensor: + current = dataset[index][0] + previous = dataset[index - 1][0] if index > 0 else None + return _make_autoencoder_input(current, previous) + + +def _metadata( + src: Path, tensor: torch.Tensor, shard: int | None = None +) -> dict[str, str]: + metadata = { + "format": "nss_v1_autoencoder_calibration", + "source": str(src), + "samples": str(tensor.shape[0]), + "shape": json.dumps(list(tensor.shape)), + "spatial_multiple": str(NSS_V1_SPATIAL_MULTIPLE), + "preprocess": "cpu_approximation_of_nss_v1_slang_pre_process", + "channels": json.dumps( + [ + "history.r", + "history.g", + "history.b", + "colour.r", + "colour.g", + "colour.b", + "motion_detector", + "feedback.r", + "feedback.g", + "feedback.b", + "feedback.a", + "luma_derivative", + ] + ), + } + if shard is not None: + metadata["shard"] = str(shard) + return metadata + + +def _write_tensor( + src: Path, dst: Path, tensor: torch.Tensor, shard: int | None = None +) -> None: + dst.parent.mkdir(parents=True, exist_ok=True) + save_file( + {"input": tensor.contiguous()}, dst, metadata=_metadata(src, tensor, shard) + ) + + +def _remove_stale_shards(dst: Path) -> None: + if not dst.exists(): + return + if not dst.is_dir(): + raise NotADirectoryError(f"Expected shard output directory, got {dst}") + for path in dst.glob("*.safetensors"): + path.unlink() + + +def _generate_verification_dataset( + src: Path, + dst: Path, + num_samples: int, + shard_size: int, +) -> None: + dataset = _model_gym_dataset(src) + sample_limit = min(num_samples, len(dataset)) + _remove_stale_shards(dst) + dst.mkdir(parents=True, exist_ok=True) + + shard_idx = 0 + tensors: list[torch.Tensor] = [] + sources: list[Path] = [] + for index in range(sample_limit): + tensors.append(_autoencoder_sample(dataset, index)) + sources.append(dataset.frame_indexes[index][0]) + if len(tensors) < shard_size and index + 1 < sample_limit: + continue + shard = torch.cat(tensors, dim=0) + _write_tensor( + sources[0], + dst / f"{shard_idx:04d}.safetensors", + shard, + shard_idx, + ) + shard_idx += 1 + tensors.clear() + sources.clear() + + +def _raw_source_path() -> Path: + return nss_test_data_root() / "source" + + +def _raw_dataset_root(snapshot_path: Path) -> Path: + if (snapshot_path / "nss" / "train").is_dir() or ( + snapshot_path / "nss" / "test" + ).is_dir(): + return snapshot_path / "nss" + return snapshot_path + + +def _env_int(name: str, default: int) -> int: + value = os.environ.get(name) + if value is None: + return default + parsed = int(value) + if parsed <= 0: + raise ValueError(f"{name} must be positive.") + return parsed + + +def _has_safetensors(path: Path) -> bool: + return path.is_file() or (path.is_dir() and any(path.glob("*.safetensors"))) + + +def _download_raw_sources( + allow_patterns: list[str], *, force_download: bool = False +) -> Path: + snapshot = Path( + snapshot_download( + repo_id=_DATASET_REPO_ID, + repo_type="dataset", + revision="5039ce015d7c877980fad44f87893fc5ac0927e2", + allow_patterns=allow_patterns, + local_dir=_raw_source_path(), + force_download=force_download, + ) + ) + return _raw_dataset_root(snapshot) + + +def _delete_raw_split(raw_root: Path, split: str, keep_env: str) -> None: + if not os.environ.get(keep_env): + shutil.rmtree(raw_root / split, ignore_errors=True) + + +def _delete_raw_sources() -> None: + if os.environ.get("NSS_KEEP_RAW_TRAIN_DATA") or os.environ.get( + "NSS_KEEP_RAW_EVALUATION_DATA" + ): + return + shutil.rmtree(_raw_source_path(), ignore_errors=True) + + +def _has_test_calibration_samples( + path: Path, num_samples: int, spatial_size: tuple[int, int] +) -> bool: + if not path.is_dir() or len(list(path.glob("*.safetensors"))) != num_samples: + return False + return all( + nss_input_shape(file_path)[1:] == (12, *spatial_size) + for file_path in path.glob("*.safetensors") + ) + + +def ensure_generated_test_calibration_dataset( + num_samples: int = 3663, + spatial_size: tuple[int, int] = (128, 128), + force_download: bool = False, +) -> Path: + """Generate evenly distributed, test-ready NSS calibration samples.""" + + if num_samples <= 0: + raise ValueError("num_samples must be positive.") + + calibration_path = nss_test_calibration_path(num_samples, spatial_size) + if _has_test_calibration_samples(calibration_path, num_samples, spatial_size): + return calibration_path + + raw_root = _download_raw_sources( + _CALIBRATION_SOURCE_ALLOW_PATTERNS, force_download=force_download + ) + dataset = _model_gym_dataset(raw_root / "train") + total_samples = len(dataset) + if num_samples > total_samples: + raise ValueError( + f"Requested {num_samples} calibration samples, but only found " + f"{total_samples}." + ) + + _remove_stale_shards(calibration_path) + calibration_path.mkdir(parents=True, exist_ok=True) + sample_indices = ( + [0] + if num_samples == 1 + else [ + index * (total_samples - 1) // (num_samples - 1) + for index in range(num_samples) + ] + ) + for output_index, sample_index in enumerate(sample_indices): + tensor = _autoencoder_sample(dataset, sample_index) + if tensor.shape[-2:] != spatial_size: + tensor = F.interpolate( + tensor.to(torch.float32), + size=spatial_size, + mode="bilinear", + align_corners=False, + ).to(torch.float16) + _write_tensor( + dataset.frame_indexes[sample_index][0], + calibration_path / f"{output_index:04d}.safetensors", + tensor, + output_index, + ) + + _delete_raw_split(raw_root, "train", "NSS_KEEP_RAW_TRAIN_DATA") + return calibration_path + + +def ensure_generated_verification_dataset( + force_download: bool = False, +) -> Path: + """Generate the held-out NSS verification input without calibration data.""" + + verification_path = nss_test_verification_path() + if _has_safetensors(verification_path): + return verification_path + + raw_root = _download_raw_sources( + _EVALUATION_SOURCE_ALLOW_PATTERNS, force_download=force_download + ) + _generate_verification_dataset( + raw_root / "test", + verification_path, + _env_int("NSS_GENERATED_EVALUATION_SAMPLES", 1), + _env_int("NSS_GENERATED_EVALUATION_SHARD_SIZE", 10), + ) + _delete_raw_split(raw_root, "test", "NSS_KEEP_RAW_EVALUATION_DATA") + return verification_path + + +def ensure_generated_test_datasets( + calibration_samples: int = 3663, + spatial_size: tuple[int, int] = (128, 128), +) -> tuple[Path, Path]: + """Generate the NSS artifacts consumed directly by ``test_nss.py``.""" + + calibration_path = ensure_generated_test_calibration_dataset( + calibration_samples, spatial_size + ) + verification_path = ensure_generated_verification_dataset() + _delete_raw_sources() + return calibration_path, verification_path + + +def generate_test_datasets_from_scratch( + calibration_samples: int = 3663, + spatial_size: tuple[int, int] = (128, 128), +) -> tuple[Path, Path]: + """Download and regenerate the NSS artifacts consumed by ``test_nss.py``.""" + + calibration_path = nss_test_calibration_path(calibration_samples, spatial_size) + verification_path = nss_test_verification_path() + for path in (calibration_path, verification_path, _raw_source_path()): + if path.is_dir(): + shutil.rmtree(path) + elif path.exists(): + path.unlink() + + calibration_path = ensure_generated_test_calibration_dataset( + calibration_samples, + spatial_size, + force_download=True, + ) + verification_path = ensure_generated_verification_dataset(force_download=True) + _delete_raw_sources() + return calibration_path, verification_path diff --git a/backends/arm/scripts/install_models_for_test.sh b/backends/arm/scripts/install_models_for_test.sh index 1e91cd9c08f..c7239ee2760 100644 --- a/backends/arm/scripts/install_models_for_test.sh +++ b/backends/arm/scripts/install_models_for_test.sh @@ -8,7 +8,8 @@ set -e pip install -r backends/arm/requirements-arm-models-test.txt # Install model gym repository -MODEL_GYM_REF="${MODEL_GYM_REF:-v0.3.0}" +MODEL_GYM_REF="${MODEL_GYM_REF:-main}" +rm -rf neural-graphics-model-gym git clone --depth 1 --branch "$MODEL_GYM_REF" https://github.com/arm/neural-graphics-model-gym.git cd neural-graphics-model-gym # Remove model-converter installation from model-gym repository (to prevent overwriting executorch version) @@ -20,3 +21,11 @@ fi pip install . --no-deps cd .. rm -rf neural-graphics-model-gym + +# Prepare the fixed NSS artifacts before pytest. The calibration data is +# generated from raw 128x128 training frames; evaluation retains the +# deployment-resolution input. +python3 -c ' +from executorch.backends.arm.scripts.generate_neural_graphics_test_data import generate_test_datasets_from_scratch +generate_test_datasets_from_scratch() +' diff --git a/backends/arm/scripts/neural_graphics_test_data.py b/backends/arm/scripts/neural_graphics_test_data.py new file mode 100644 index 00000000000..e404394dca2 --- /dev/null +++ b/backends/arm/scripts/neural_graphics_test_data.py @@ -0,0 +1,138 @@ +# 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. +"""Load generated NSS autoencoder calibration and verification data.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Iterator + +import torch +from safetensors import safe_open + + +NSS_DATASET_REVISION = "main" +_TEST_CALIBRATION_DIR = "calibration/nss_v1_autoencoder_cpu_calibration" +_VERIFICATION_DIR = "evaluation/nss_v1_autoencoder_cpu_evaluation" + + +def nss_test_data_root() -> Path: + if "NSS_GENERATED_DATASET_ROOT" in os.environ: + return Path(os.environ["NSS_GENERATED_DATASET_ROOT"]) + else: + return ( + Path(__file__).resolve().parents[1] + / "test" + / "models" + / "nss_data" + / NSS_DATASET_REVISION + ) + + +def nss_test_calibration_path( + num_samples: int = 3663, spatial_size: tuple[int, int] = (128, 128) +) -> Path: + """Return the preprocessed calibration dataset path used by NSS tests.""" + + height, width = spatial_size + return nss_test_data_root() / ( + f"{_TEST_CALIBRATION_DIR}_{num_samples}_{height}x{width}" + ) + + +def nss_test_verification_path() -> Path: + """Return the preprocessed verification dataset path used by NSS tests.""" + + return nss_test_data_root() / _VERIFICATION_DIR + + +def nss_input_shape(path: Path) -> tuple[int, int, int, int]: + """Return the shape of the ``input`` tensor in a safetensors file.""" + + with safe_open(path, framework="pt", device="cpu") as handle: + keys = set(handle.keys()) + if "input" not in keys: + raise KeyError(f"{path} does not contain an `input` tensor. Found {keys}.") + + shape = tuple(handle.get_slice("input").get_shape()) + + if len(shape) != 4: + raise ValueError(f"Expected NCHW `input`, got shape {shape} in {path}.") + if shape[1] != 12: + raise ValueError(f"Expected 12 NSS input channels, got shape {shape}.") + return shape # type: ignore[return-value] + + +def _load_input_slice(path: Path, start: int, stop: int) -> torch.Tensor: + if not path.is_file(): + raise FileNotFoundError(path) + + shape = nss_input_shape(path) + if start < 0 or stop <= start or stop > shape[0]: + raise ValueError(f"Invalid slice [{start}:{stop}] for shape {shape}.") + + with safe_open(path, framework="pt", device="cpu") as handle: + tensor = ( + handle.get_slice("input")[start:stop].to(dtype=torch.float32).contiguous() + ) + + return tensor.to(memory_format=torch.channels_last) + + +def _safetensor_files(path: Path) -> list[Path]: + if path.is_file(): + return [path] + if path.is_dir(): + files = sorted(path.glob("*.safetensors")) + if files: + return files + raise FileNotFoundError(f"No safetensors found at {path}") + + +def _load_sample(path: Path, start: int = 0) -> tuple[torch.Tensor]: + if start < 0: + raise ValueError("start must be non-negative.") + + skipped = 0 + for file_path in _safetensor_files(path): + file_samples = nss_input_shape(file_path)[0] + if start < skipped + file_samples: + local_index = start - skipped + return (_load_input_slice(file_path, local_index, local_index + 1),) + skipped += file_samples + + raise ValueError(f"Sample {start} is outside the {skipped} samples in {path}.") + + +def iter_calibration_samples( + path: Path, + *, + num_samples: int = 8, +) -> Iterator[tuple[torch.Tensor]]: + """Stream NSS calibration samples without retaining them in memory.""" + + if num_samples <= 0: + raise ValueError("num_samples must be positive.") + + files = _safetensor_files(path) + if num_samples > len(files): + raise ValueError( + f"Requested {num_samples} samples from {path}, but only found {len(files)}." + ) + + for file_path in files[:num_samples]: + yield (_load_input_slice(file_path, 0, 1),) + + +def load_verification_inputs( + path: Path | None = None, + *, + start: int = 0, +) -> tuple[torch.Tensor]: + """Load one held-out verification sample in ``example_inputs`` format.""" + + path = nss_test_verification_path() if path is None else path + return _load_sample(path, start) diff --git a/backends/arm/scripts/pre-push b/backends/arm/scripts/pre-push index be12e14fccf..0cdbb85e17f 100755 --- a/backends/arm/scripts/pre-push +++ b/backends/arm/scripts/pre-push @@ -74,6 +74,8 @@ run_docgen_check() { } run_vgf_op_support_checks() { + local coverage_output + echo -e "${INFO} Generating VGF operator support documentation" if ! python "$VGF_OP_SUPPORT_SCRIPT"; then @@ -90,7 +92,8 @@ run_vgf_op_support_checks() { echo -e "${INFO} Checking VGF operator support coverage" - if ! python "$VGF_OP_SUPPORT_SCRIPT" --check; then + if ! coverage_output=$(python "$VGF_OP_SUPPORT_SCRIPT" --check 2>&1); then + echo "$coverage_output" >&2 echo -e "${ERROR} VGF operator support coverage check failed" >&2 FAILED=1 else diff --git a/backends/arm/scripts/run_fvp.sh b/backends/arm/scripts/run_fvp.sh index 7289f37484a..e0e2ba2aa7b 100755 --- a/backends/arm/scripts/run_fvp.sh +++ b/backends/arm/scripts/run_fvp.sh @@ -25,6 +25,7 @@ timeout="600" etrecord_file="" trace_file="" semihosting_cwd="" +semihosting_cmd_line="" ethosu_fast=0 help() { @@ -38,6 +39,7 @@ help() { echo " --etrecord= If ETDump is used you can supply a ETRecord file matching the PTE" echo " --trace_file= File to write PMU trace output to" echo " --semihosting-cwd= Enable target semihosting with this host working directory" + echo " --semihosting-cmd-line= Command line passed to a semihosting runner" echo " --fast Use fast Ethos-U model simulation for Ethos-U targets" exit 0 } @@ -53,6 +55,7 @@ for arg in "$@"; do --etrecord=*) etrecord_file="${arg#*=}";; --trace_file=*) trace_file="${arg#*=}";; --semihosting-cwd=*) semihosting_cwd="${arg#*=}";; + --semihosting-cmd-line=*) semihosting_cmd_line="${arg#*=}";; --fast) ethosu_fast=1;; *) ;; @@ -139,6 +142,16 @@ if [[ -n "${semihosting_cwd}" ]]; then -C "mps4_board.subsystem.cpu0.semihosting-cwd=${semihosting_cwd}" ) fi +if [[ -n "${semihosting_cmd_line}" ]]; then + [[ -n "${semihosting_cwd}" ]] \ + || { echo "--semihosting-cmd-line requires --semihosting-cwd"; exit 1; } + semihosting_args_u55+=( + -C "cpu0.semihosting-cmd_line=${semihosting_cmd_line}" + ) + semihosting_args_u85+=( + -C "mps4_board.subsystem.cpu0.semihosting-cmd_line=${semihosting_cmd_line}" + ) +fi if [[ ${target} == cortex-m* ]]; then [[ -z "${bundle_file}" ]] \ diff --git a/backends/arm/test/misc/test_docgen_op_support.py b/backends/arm/test/misc/test_docgen_op_support.py index 4956b15e06b..3622cfc5228 100644 --- a/backends/arm/test/misc/test_docgen_op_support.py +++ b/backends/arm/test/misc/test_docgen_op_support.py @@ -364,6 +364,33 @@ def test_matching_evidence_accepts_stage_equivalent_alias() -> None: assert records[0].asserted_op == alias +def test_collect_backend_custom_partition_ops_discovers_fp_and_int_profiles() -> None: + from executorch.backends.arm.tosa import TosaSpecification + + tosa_spec = TosaSpecification.create_from_string(docgen.BACKEND_TOSA_SPEC) + custom_ops = docgen._collect_backend_custom_partition_ops(tosa_spec) + canonical_by_profile = { + profile: { + docgen._canonical_pytorch_op_from_target(target) for target in targets + } + for profile, targets in custom_ops.items() + } + + expected = "torch.ops.aten.grid_sampler_2d.default" + assert expected in canonical_by_profile["FP"] + assert expected in canonical_by_profile["INT"] + + +def test_collect_backend_supported_ops_includes_vgf_custom_partition_ops() -> None: + repo_root = Path(__file__).resolve().parents[4] + + expected = docgen._collect_backend_supported_ops(repo_root) + row = expected["torch.ops.aten.grid_sampler_2d.default"] + + assert row.support_profiles == {"FP", "INT"} + assert "VgfPartitioner.register_custom_partition_op" in row.evidence + + def test_run_check_reports_missing_profile( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] ) -> None: @@ -411,6 +438,7 @@ def test_run_check_reports_missing_profile( def test_run_check_strict_ast_fails_on_unresolved_attribution( monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], ) -> None: unresolved = [ docgen.UnresolvedPipelineEvidence( @@ -431,7 +459,10 @@ def test_run_check_strict_ast_fails_on_unresolved_attribution( monkeypatch.setattr(docgen, "_collect_backend_supported_ops", lambda _root: {}) assert docgen.run_check(Path("/repo"), strict_ast=False) == 0 + assert "Unresolved VgfPipeline attribution" not in capsys.readouterr().out + assert docgen.run_check(Path("/repo"), strict_ast=True) == 1 + assert "Unresolved VgfPipeline attribution" in capsys.readouterr().out def test_main_writes_requested_markdown_and_html( diff --git a/backends/arm/test/misc/test_mixed_type_lowering.py b/backends/arm/test/misc/test_mixed_type_lowering.py index 6a2a1e4cbd5..14663b6fc51 100644 --- a/backends/arm/test/misc/test_mixed_type_lowering.py +++ b/backends/arm/test/misc/test_mixed_type_lowering.py @@ -33,14 +33,17 @@ def repeat_op_dict(op_dict, times): } q_tosa_ops = { "CAST": {"INT8": 1}, - "MUL": {"FP32": 1}, # scale multiplication - "ADD": {"FP32": 2}, # zero-point addition, rounding - "SUB": {"FP32": 1}, # for rounding + "MUL": {"FP32": 2}, # scale multiplication + round()'s internal multiply + "ADD": {"FP32": 1}, # zero-point addition + "SUB": {"FP32": 2}, # for rounding + "CEIL": {"FP32": 1}, # for rounding "CLAMP": {"FP32": 1}, # clamp - "GREATER_EQUAL": {"BOOL": 1}, # for rounding "SELECT": {"FP32": 1}, # for rounding - "CEIL": {"FP32": 1}, # for rounding - "FLOOR": {"FP32": 1}, # for rounding + "FLOOR": {"FP32": 2}, # for rounding + "EQUAL": {"BOOL": 2}, # for rounding + "GREATER": {"BOOL": 1}, # for rounding + "LOGICAL_AND": {"BOOL": 1}, # for rounding + "LOGICAL_OR": {"BOOL": 1}, # for rounding } diff --git a/backends/arm/test/misc/test_mobilesam.py b/backends/arm/test/misc/test_mobilesam.py new file mode 100644 index 00000000000..323549c017b --- /dev/null +++ b/backends/arm/test/misc/test_mobilesam.py @@ -0,0 +1,48 @@ +# 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 pathlib import Path + +import numpy as np +import pytest +import torch + +from examples.arm.mobilesam_prompt_segmentation_example_ethos_u.model_export.export_mobilesam import ( + iou, + mask, +) +from examples.arm.mobilesam_prompt_segmentation_example_ethos_u.runtime.visualize_fvp_output import ( + load_mask, + OUTPUT_SIZE, +) + + +def test_mask_and_iou() -> None: + first = mask(torch.tensor([[[[-1.0, 1.0], [1.0, -1.0]]]])) + second = np.array([[0, 1], [0, 1]], dtype=np.uint8) + + assert first.tolist() == [[0, 1], [1, 0]] + assert iou(first, second) == pytest.approx(1 / 3) + + +def test_load_mask(tmp_path: Path) -> None: + output = np.ones((OUTPUT_SIZE, OUTPUT_SIZE), dtype=np.float32) + output[0, 0] = -1 + path = tmp_path / "output.bin" + output.tofile(path) + + loaded = load_mask(path) + + assert loaded.shape == (OUTPUT_SIZE, OUTPUT_SIZE) + assert loaded[0, 0] == 0 + assert loaded[1, 1] == 1 + + +def test_load_mask_rejects_wrong_output_size(tmp_path: Path) -> None: + path = tmp_path / "output.bin" + np.ones(3, dtype=np.float32).tofile(path) + + with pytest.raises(ValueError, match="Expected"): + load_mask(path) diff --git a/backends/arm/test/misc/test_pass_pipeline_config.py b/backends/arm/test/misc/test_pass_pipeline_config.py index 5b820102d4f..2ac52322125 100644 --- a/backends/arm/test/misc/test_pass_pipeline_config.py +++ b/backends/arm/test/misc/test_pass_pipeline_config.py @@ -118,14 +118,25 @@ def test_sdpa_safe_softmax_guard_config_controls_guard_removal_pass(): assert RemoveSafeSoftmaxGuardPass in manager._skip_pass_types - stable_compile_spec = TosaCompileSpec( + auto_compile_spec = TosaCompileSpec( + TosaSpecification.create_from_string("TOSA-1.00+INT") + ) + auto_config = ArmPassPipelineConfig( + sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO + ) + auto_compile_spec.set_pass_pipeline_config(auto_config) + auto_manager = ArmPassManager(auto_compile_spec) + + assert RemoveSafeSoftmaxGuardPass in auto_manager._skip_pass_types + + remove_compile_spec = TosaCompileSpec( TosaSpecification.create_from_string("TOSA-1.00+INT") ) remove_config = ArmPassPipelineConfig( sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.REMOVE ) - stable_compile_spec.set_pass_pipeline_config(remove_config) - remove_manager = ArmPassManager(stable_compile_spec) + remove_compile_spec.set_pass_pipeline_config(remove_config) + remove_manager = ArmPassManager(remove_compile_spec) assert RemoveSafeSoftmaxGuardPass not in remove_manager._skip_pass_types @@ -260,11 +271,11 @@ def test_leaky_relu_decompose_config_reaches_backend_pipeline(): def test_sdpa_safe_softmax_guard_config_serializes(): config = ArmPassPipelineConfig( - sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.REMOVE + sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO ) roundtripped = ArmPassPipelineConfig.from_dict(config.to_dict()) - assert roundtripped.sdpa_safe_softmax_guard is SDPASafeSoftmaxGuardPolicy.REMOVE + assert roundtripped.sdpa_safe_softmax_guard is SDPASafeSoftmaxGuardPolicy.AUTO def test_sdpa_safe_softmax_guard_preserves_positional_config_arguments(): diff --git a/backends/arm/test/misc/test_process_node.py b/backends/arm/test/misc/test_process_node.py index 02d2a5e012b..994c1628836 100644 --- a/backends/arm/test/misc/test_process_node.py +++ b/backends/arm/test/misc/test_process_node.py @@ -10,6 +10,7 @@ import torch import tosa_serializer as ts from executorch.backends.arm.process_node import _add_const, process_placeholder +from executorch.backends.arm.tosa.constant_pool import TosaSerializerWithConstantPool from executorch.backends.arm.tosa.mapping import TosaArg, TosaSpecialDtype from executorch.backends.arm.tosa.specification import TosaSpecification from executorch.exir import to_edge @@ -39,7 +40,7 @@ def __init__(self) -> None: self.name = None self.serialized_bytes = None - def addConst(self, shape, dtype, values, name): + def addUnpooledConst(self, shape, dtype, values, name): self.shape = shape self.dtype = dtype self.values = np.asarray(values) @@ -111,7 +112,7 @@ def test_add_const_fp4_in_packed_storage() -> None: TosaArg, SimpleNamespace(dtype=ts.DType.FP4E2M1, shape=(1, 1, 8)), ) - tosa_graph = ts.TosaSerializer() + tosa_graph = TosaSerializerWithConstantPool() _add_const(tosa_graph, packed_values, tosa_arg, name="fp4_weight") @@ -140,7 +141,7 @@ def _test_add_const_fp6_in_packed_storage(dtype: int) -> None: TosaArg, SimpleNamespace(dtype=dtype, shape=(1, 1, 32)), ) - tosa_graph = ts.TosaSerializer() + tosa_graph = TosaSerializerWithConstantPool() _add_const(tosa_graph, values, tosa_arg, name="fp6_weight") diff --git a/backends/arm/test/misc/test_tosa_constant_pool.py b/backends/arm/test/misc/test_tosa_constant_pool.py new file mode 100644 index 00000000000..4a5f4502844 --- /dev/null +++ b/backends/arm/test/misc/test_tosa_constant_pool.py @@ -0,0 +1,144 @@ +# 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 gc +import weakref + +import numpy as np +import pytest +import tosa_serializer as ts +from executorch.backends.arm.tosa.constant_pool import TosaSerializerWithConstantPool + + +def _serializer(path_prefix=""): + return TosaSerializerWithConstantPool( + path_prefix, + targetMajor=1, + targetMinor=0, + targetPatch=0, + targetDraft=False, + ) + + +@pytest.mark.parametrize("dtype", [ts.DType.INT8, ts.DType.SHAPE]) +def test_identical_constants_are_reused(dtype): + serializer = _serializer() + + first = serializer.addConst([1], dtype, [0], name="first") + second = serializer.addConst([1], dtype, [0], name="second") + + assert isinstance(serializer, ts.TosaSerializer) + assert second is first + assert first.name == "first" + block = serializer.currRegion.currBasicBlock + assert len(block.operators) == 1 + constants = block.shapes if dtype == ts.DType.SHAPE else block.tensors + assert list(constants.keys()) == ["first"] + + +@pytest.mark.parametrize("dtype", [ts.DType.INT8, ts.DType.SHAPE]) +def test_constant_pool_does_not_keep_serializer_alive(dtype): + serializer = _serializer() + serializer.addConst([1], dtype, [0], name="first") + serializer.addConst([1], dtype, [0], name="duplicate") + serializer.serialize() + serializer_ref = weakref.ref(serializer) + + del serializer + gc.collect() + + assert serializer_ref() is None + + +def test_unpooled_constants_are_not_reused(): + serializer = _serializer() + + first = serializer.addUnpooledConst([1], ts.DType.INT8, [0], name="first") + second = serializer.addUnpooledConst([1], ts.DType.INT8, [0], name="second") + + assert second is not first + assert len(serializer.currRegion.currBasicBlock.operators) == 2 + assert list(serializer.currRegion.currBasicBlock.tensors.keys()) == [ + "first", + "second", + ] + + +def test_unnamed_constant_uses_serializer_generated_name(): + serializer = _serializer() + + constant = serializer.addConst([1], ts.DType.INT8, [0]) + + assert constant.name + + +@pytest.mark.parametrize( + "first,second", + [ + (([1], ts.DType.INT8, [0]), ([1], ts.DType.INT16, [0])), + (([1], ts.DType.INT8, [0]), ([2], ts.DType.INT8, [0, 0])), + (([1], ts.DType.INT8, [0]), ([1], ts.DType.INT8, [1])), + ], +) +def test_constants_with_different_keys_remain_separate(first, second): + serializer = _serializer() + + first_const = serializer.addConst(*first, name="first") + second_const = serializer.addConst(*second, name="second") + + assert second_const is not first_const + assert len(serializer.currRegion.currBasicBlock.operators) == 2 + + +def test_float_constants_use_exact_serialized_values(): + serializer = _serializer() + + positive_zero = serializer.addConst( + [1], ts.DType.FP32, np.array([0.0]), name="positive_zero" + ) + negative_zero = serializer.addConst( + [1], ts.DType.FP32, np.array([-0.0]), name="negative_zero" + ) + repeated_negative_zero = serializer.addConst( + [1], + ts.DType.FP32, + np.array([-0.0]), + name="repeated_negative_zero", + ) + + assert negative_zero is not positive_zero + assert repeated_negative_zero is negative_zero + assert len(serializer.currRegion.currBasicBlock.operators) == 2 + + +def test_constants_are_scoped_to_basic_blocks(): + serializer = _serializer() + first = serializer.addConst([1], ts.DType.INT8, [0], name="first") + + serializer.startRegion("main") + serializer.currRegion.addBasicBlock("main") + second = serializer.addConst([1], ts.DType.INT8, [0], name="second") + + assert second is not first + assert second.name == "second" + + +def test_start_region_preserves_path_prefix(): + serializer = _serializer("artifacts") + + serializer.startRegion("other") + + assert serializer.currRegion.pathPrefix == "artifacts" + + +def test_constant_pool_serialization_is_deterministic(): + def serialize(): + serializer = _serializer() + serializer.addConst([1], ts.DType.INT8, [0], name="first") + serializer.addConst([1], ts.DType.INT8, [0], name="duplicate") + serializer.addConst([1], ts.DType.INT8, [1], name="second") + return bytes(serializer.serialize()) + + assert serialize() == serialize() diff --git a/backends/arm/test/misc/test_tosa_operator_support.py b/backends/arm/test/misc/test_tosa_operator_support.py index a3dce64fefc..87d2d8a4c0d 100644 --- a/backends/arm/test/misc/test_tosa_operator_support.py +++ b/backends/arm/test/misc/test_tosa_operator_support.py @@ -3,10 +3,16 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import pytest import torch +from executorch.backends.arm.operator_support.index_tensor_support import ( + IndexTensorSupported, +) from executorch.backends.arm.operator_support.tosa_supported_operators import ( + CheckFPComparisonInputs, CheckKnownUnsupportedTOSASemantics, ) +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.backend.utils import WhyNoPartitionReporter from executorch.exir.dialects._ops import ops as exir_ops from torch._subclasses.fake_tensor import FakeTensorMode @@ -27,6 +33,49 @@ def _checker() -> CheckKnownUnsupportedTOSASemantics: return CheckKnownUnsupportedTOSASemantics(WhyNoPartitionReporter()) +def _fp_comparison_checker() -> CheckFPComparisonInputs: + return CheckFPComparisonInputs(WhyNoPartitionReporter()) + + +@pytest.mark.parametrize( + "target", + ( + exir_ops.edge.aten.eq.Tensor, + exir_ops.edge.aten.ne.Tensor, + exir_ops.edge.aten.ge.Tensor, + exir_ops.edge.aten.gt.Tensor, + exir_ops.edge.aten.le.Tensor, + exir_ops.edge.aten.lt.Tensor, + ), +) +@pytest.mark.parametrize( + "dtype", + (torch.bool, torch.uint8, torch.int32, torch.int64), +) +def test_fp_comparison_rejects_unsupported_inputs(target, dtype) -> None: + graph = torch.fx.Graph() + x = _placeholder(graph, "x", (3, 4), dtype) + y = _placeholder(graph, "y", (3, 4), dtype) + node = graph.call_function(target, (x, y)) + node.meta["val"] = _fake_tensor((3, 4), torch.bool) + + assert not _fp_comparison_checker().is_node_supported({}, node) + + +@pytest.mark.parametrize( + "dtype", + (torch.float16, torch.float32, torch.bfloat16, torch.int8, torch.int16), +) +def test_fp_comparison_accepts_supported_inputs(dtype) -> None: + graph = torch.fx.Graph() + x = _placeholder(graph, "x", (3, 4), dtype) + y = _placeholder(graph, "y", (3, 4), dtype) + node = graph.call_function(exir_ops.edge.aten.eq.Tensor, (x, y)) + node.meta["val"] = _fake_tensor((3, 4), torch.bool) + + assert _fp_comparison_checker().is_node_supported({}, node) + + def test_rejects_argmax_without_int32_cast_user() -> None: graph = torch.fx.Graph() x = _placeholder(graph, "x", (3, 4)) @@ -87,3 +136,19 @@ def test_rejects_argmax_with_mixed_int32_cast_and_raw_user() -> None: raw_user.meta["val"] = _fake_tensor((3,), torch.int64) assert not _checker().is_node_supported({}, node) + + +@pytest.mark.parametrize("dtype", (torch.bool, torch.uint8)) +def test_rejects_index_tensor_mask(dtype: torch.dtype) -> None: + graph = torch.fx.Graph() + x = _placeholder(graph, "x", (5, 2, 3)) + index = _placeholder(graph, "index", (5,), dtype) + node = graph.call_function(exir_ops.edge.aten.index.Tensor, (x, [index])) + node.meta["val"] = _fake_tensor((2, 2, 3)) + + checker = IndexTensorSupported( + TosaSpecification.create_from_string("TOSA-1.0+INT+u55"), + WhyNoPartitionReporter(), + ) + + assert not checker.is_node_supported({}, node) diff --git a/backends/arm/test/misc/test_transpose_counts.py b/backends/arm/test/misc/test_transpose_counts.py index bd73ddfe0cb..643a14b3301 100644 --- a/backends/arm/test/misc/test_transpose_counts.py +++ b/backends/arm/test/misc/test_transpose_counts.py @@ -543,7 +543,7 @@ def forward(self, x: torch.Tensor): "groupnorm_channels_last": TransposeCountCase( GroupNormModule(), (torch.randn(1, 4, 4, 4).to(memory_format=torch.channels_last),), - 2, + 1, ), "cumsum_rank4_dim3_channels_last": TransposeCountCase( CumsumModule(), diff --git a/backends/arm/test/misc/test_tosa_data_layout_visitors.py b/backends/arm/test/misc/tosa_dialect/test_tosa_data_layout_visitors.py similarity index 100% rename from backends/arm/test/misc/test_tosa_data_layout_visitors.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_data_layout_visitors.py diff --git a/backends/arm/test/misc/test_tosa_dialect_activation.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_activation.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_activation.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_activation.py diff --git a/backends/arm/test/misc/test_tosa_dialect_argmax.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_argmax.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_argmax.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_argmax.py diff --git a/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py index 920454f5a9b..c4a811f3a2c 100644 --- a/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py +++ b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py @@ -81,33 +81,6 @@ def _to_fake(mode: FakeTensorMode, *values): (2, 3), torch.int8, ), - pytest.param( - "EQUAL", - "TOSA-1.1+INT", - torch.randint(1, 16, (2, 1, 3), dtype=torch.int32), - torch.randint(1, 8, (1, 4, 3), dtype=torch.int32), - {}, - (2, 4, 3), - torch.bool, - ), - pytest.param( - "GREATER", - "TOSA-1.1+FP", - torch.randn((2, 1, 3), dtype=torch.float32), - torch.randn((1, 4, 3), dtype=torch.float32), - {}, - (2, 4, 3), - torch.bool, - ), - pytest.param( - "GREATER_EQUAL", - "TOSA-1.1+INT", - torch.randint(1, 16, (2, 1, 3), dtype=torch.int32), - torch.randint(1, 8, (1, 4, 3), dtype=torch.int32), - {}, - (2, 4, 3), - torch.bool, - ), pytest.param( "INTDIV", "TOSA-1.1+INT", @@ -400,31 +373,6 @@ def test_intdiv_supports_int32_on_fp_profile() -> None: assert tuple(output.shape) == tuple(input1.shape) -def test_equal_rejects_int8() -> None: - input1 = torch.randint(-8, 8, (2, 3), dtype=torch.int8) - input2 = torch.randint(-8, 8, (2, 3), dtype=torch.int8) - - with TosaLoweringContext( - TosaSpecification.create_from_string("TOSA-1.1+INT") - ), FakeTensorMode() as mode: - with pytest.raises(TosaValueError, match="Unsupported dtype"): - exir_ops.backend.tosa.EQUAL.default(*_to_fake(mode, input1, input2)) - - -@pytest.mark.parametrize("op_name", ["EQUAL", "GREATER", "GREATER_EQUAL"]) -def test_compare_ops_reject_int32_on_fp_profile(op_name: str) -> None: - input1 = torch.randint(1, 16, (2, 3), dtype=torch.int32) - input2 = torch.randint(1, 8, (2, 3), dtype=torch.int32) - - with TosaLoweringContext( - TosaSpecification.create_from_string("TOSA-1.1+FP") - ), FakeTensorMode() as mode: - with pytest.raises(TosaValueError, match="doesn't support int32"): - getattr(exir_ops.backend.tosa, op_name).default( - *_to_fake(mode, input1, input2) - ) - - @pytest.mark.parametrize("op_name", ["MAXIMUM", "MINIMUM"]) def test_extrema_ops_reject_int32_on_fp_profile(op_name: str) -> None: input1 = torch.randint(1, 16, (2, 3), dtype=torch.int32) diff --git a/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_comparison.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_comparison.py new file mode 100644 index 00000000000..6fbf25ed380 --- /dev/null +++ b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_comparison.py @@ -0,0 +1,97 @@ +# 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 executorch.backends.arm.tosa.dialect # noqa: F401 +import pytest +import torch +from executorch.backends.arm.tosa.dialect.lib import TosaValueError +from executorch.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) +from executorch.exir.dialects._ops import ops as exir_ops +from torch._subclasses.fake_tensor import FakeTensorMode + + +def _to_fake(mode: FakeTensorMode, *values): + return [ + mode.from_tensor(value) if isinstance(value, torch.Tensor) else value + for value in values + ] + + +@pytest.mark.parametrize( + ( + "op_name", + "spec", + "input1", + "input2", + "expected_shape", + ), + [ + pytest.param( + "EQUAL", + "TOSA-1.1+INT", + torch.randint(1, 16, (2, 1, 3), dtype=torch.int32), + torch.randint(1, 8, (1, 4, 3), dtype=torch.int32), + (2, 4, 3), + ), + pytest.param( + "GREATER", + "TOSA-1.1+FP", + torch.randn((2, 1, 3), dtype=torch.float32), + torch.randn((1, 4, 3), dtype=torch.float32), + (2, 4, 3), + ), + pytest.param( + "GREATER_EQUAL", + "TOSA-1.1+INT", + torch.randint(1, 16, (2, 1, 3), dtype=torch.int32), + torch.randint(1, 8, (1, 4, 3), dtype=torch.int32), + (2, 4, 3), + ), + ], +) +def test_tosa_comparison_ops( + op_name: str, + spec: str, + input1: torch.Tensor, + input2: torch.Tensor, + expected_shape: tuple[int, ...], +) -> None: + with TosaLoweringContext( + TosaSpecification.create_from_string(spec) + ), FakeTensorMode() as mode: + output = getattr(exir_ops.backend.tosa, op_name).default( + *_to_fake(mode, input1, input2) + ) + + assert output.dtype == torch.bool + assert tuple(output.shape) == expected_shape + + +def test_equal_rejects_int8() -> None: + input1 = torch.randint(-8, 8, (2, 3), dtype=torch.int8) + input2 = torch.randint(-8, 8, (2, 3), dtype=torch.int8) + + with TosaLoweringContext( + TosaSpecification.create_from_string("TOSA-1.1+INT") + ), FakeTensorMode() as mode: + with pytest.raises(TosaValueError, match="Unsupported dtype"): + exir_ops.backend.tosa.EQUAL.default(*_to_fake(mode, input1, input2)) + + +@pytest.mark.parametrize("op_name", ["EQUAL", "GREATER", "GREATER_EQUAL"]) +def test_compare_ops_reject_int32_on_fp_profile(op_name: str) -> None: + input1 = torch.randint(1, 16, (2, 3), dtype=torch.int32) + input2 = torch.randint(1, 8, (2, 3), dtype=torch.int32) + + with TosaLoweringContext( + TosaSpecification.create_from_string("TOSA-1.1+FP") + ), FakeTensorMode() as mode: + with pytest.raises(TosaValueError, match="doesn't support int32"): + getattr(exir_ops.backend.tosa, op_name).default( + *_to_fake(mode, input1, input2) + ) diff --git a/backends/arm/test/misc/test_tosa_dialect_fft.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_fft.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_fft.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_fft.py diff --git a/backends/arm/test/misc/test_tosa_dialect_max_pool2d_adaptive.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_max_pool2d_adaptive.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_max_pool2d_adaptive.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_max_pool2d_adaptive.py diff --git a/backends/arm/test/misc/test_tosa_dialect_scatter.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_scatter.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_scatter.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_scatter.py diff --git a/backends/arm/test/misc/test_tosa_dialect_unary_ops.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_unary_ops.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_unary_ops.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_unary_ops.py diff --git a/backends/arm/test/misc/test_tosa_shape_node_visitors.py b/backends/arm/test/misc/tosa_dialect/test_tosa_shape_node_visitors.py similarity index 96% rename from backends/arm/test/misc/test_tosa_shape_node_visitors.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_shape_node_visitors.py index d88424599e8..3813d7a0b0f 100644 --- a/backends/arm/test/misc/test_tosa_shape_node_visitors.py +++ b/backends/arm/test/misc/tosa_dialect/test_tosa_shape_node_visitors.py @@ -17,6 +17,7 @@ NodeVisitor, ) from executorch.backends.arm.test.runner_utils import TosaReferenceModelDispatch +from executorch.backends.arm.tosa.constant_pool import TosaSerializerWithConstantPool from executorch.backends.arm.tosa.mapping import TosaArg from executorch.backends.arm.tosa.specification import TosaSpecification from torch.fx import Node @@ -75,7 +76,7 @@ def _shape_spec() -> TosaSpecification: def _serializer() -> ts.TosaSerializer: - return ts.TosaSerializer( + return TosaSerializerWithConstantPool( "", targetMajor=1, targetMinor=1, @@ -289,6 +290,22 @@ def test_const_shape_node_visitor_serializes_const_operator() -> None: assert _serialized_op_codes(tosa_graph) == [ts.Op.CONST_SHAPE] +def test_const_shape_node_visitor_preserves_output_name() -> None: + visitor = get_node_visitors(_shape_spec())["tosa.CONST_SHAPE.default"] + tosa_graph = _serializer() + tosa_graph.addConst([2], ts.DType.SHAPE, [2, 3], name="helper") + + _define_node( + visitor, + SimpleNamespace(name="node", meta={"val": [2, 3]}, kwargs={}), + tosa_graph, + [SimpleNamespace(special=[2, 3])], + SimpleNamespace(name="output", shape=(2,)), + ) + + assert list(tosa_graph.currRegion.currBasicBlock.shapes) == ["helper", "output"] + + def test_dim_shape_node_visitor_serializes_operator() -> None: visitor = get_node_visitors(_shape_spec())["tosa.DIM.default"] tosa_graph = _serializer() diff --git a/backends/arm/test/misc/test_tosa_shape_support.py b/backends/arm/test/misc/tosa_dialect/test_tosa_shape_support.py similarity index 100% rename from backends/arm/test/misc/test_tosa_shape_support.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_shape_support.py diff --git a/backends/arm/test/misc/test_tosa_spec.py b/backends/arm/test/misc/tosa_dialect/test_tosa_spec.py similarity index 100% rename from backends/arm/test/misc/test_tosa_spec.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_spec.py diff --git a/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_model.py b/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_model.py index 34f87c1ca8b..792cc689c39 100644 --- a/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_model.py +++ b/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_model.py @@ -197,7 +197,7 @@ class DeepSeekR1DistillQwenModelTestCase: "base_model": DeepSeekR1DistillQwenModelTestCase( model_cls=BaseModelWrapper, config_factory=_make_deepseek_r1_distill_qwen_1_5b_model_config, - atol=0.1, + atol=0.17, rtol=0.1, tosa_spec="TOSA-1.0+FP+bf16", ), diff --git a/backends/arm/test/models/Qwen3_VL/test_qwen3_vl_model.py b/backends/arm/test/models/Qwen3_VL/test_qwen3_vl_model.py index 2b83f75030a..a826c4fbdb9 100644 --- a/backends/arm/test/models/Qwen3_VL/test_qwen3_vl_model.py +++ b/backends/arm/test/models/Qwen3_VL/test_qwen3_vl_model.py @@ -314,6 +314,8 @@ def _test_qwen3_vl_full_models_vgf_no_quant_bf16( def _test_qwen3_vl_text_model_tosa_mxfp8_bf16( config_factory=_make_qwen3_vl_e2e_test_config, + frobenius_threshold: float = 0.1, + cosine_threshold: float = 0.98, ): # The Qwen 3 VL FP8 model only quantizes the TextModel model, inputs = TextModelWrapper.prepare_model_and_inputs(config_factory) @@ -326,8 +328,8 @@ def _test_qwen3_vl_text_model_tosa_mxfp8_bf16( aten_op=aten_op_mxfp_linear, exir_op=[], filter_fn=_is_linear, - frobenius_threshold=0.1, - cosine_threshold=0.98, + frobenius_threshold=frobenius_threshold, + cosine_threshold=cosine_threshold, mxfp_config=mxfp_config, tosa_version="1.1", tosa_extensions=["bf16", "mxfp"], @@ -404,4 +406,8 @@ def test_qwen3_vl_2b_instruct_full_models_vgf_no_quant_bf16( @pytest.mark.slow @pytest.mark.xlarge def test_qwen3_vl_2b_instruct_text_model_tosa_mxfp8_bf16(): - _test_qwen3_vl_text_model_tosa_mxfp8_bf16(_make_qwen3_vl_2b_instruct_layer_config) + _test_qwen3_vl_text_model_tosa_mxfp8_bf16( + _make_qwen3_vl_2b_instruct_layer_config, + frobenius_threshold=0.3, + cosine_threshold=0.95, + ) diff --git a/backends/arm/test/models/stable_diffusion_3_5_large/test_CLIPTextModelWithProjection_sd35_large.py b/backends/arm/test/models/stable_diffusion_3_5_large/test_CLIPTextModelWithProjection_sd35_large.py new file mode 100644 index 00000000000..34d61805e4d --- /dev/null +++ b/backends/arm/test/models/stable_diffusion_3_5_large/test_CLIPTextModelWithProjection_sd35_large.py @@ -0,0 +1,218 @@ +# 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 pytest +import torch +from executorch.backends.arm.test import common +from executorch.backends.arm.test.models.stable_diffusion_3_5_large.test_configs_sd35_large import ( + get_int64_to_int32_passes, + get_tiny_sd35_large_text_encoder_2_config, + get_tiny_sd35_large_text_encoder_config, +) +from executorch.backends.arm.test.tester.test_pipeline import ( + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) +from executorch.examples.models.stable_diffusion_3_5_large.model import ( + SD3CLIPTextEncoderWrapper, +) +from transformers import CLIPTextModelWithProjection + +input_t = Tuple[torch.Tensor] + + +class TestCLIPTextModelWithProjection: + """Test helper for SD3.5 Large CLIPTextModelWithProjection configs.""" + + ops_after_partitioner_FP = { + "executorch_exir_dialects_edge__ops_aten_argmax_default": 1, + "executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor": 1, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2, + "torch.ops.higher_order.executorch_call_delegate": 2, + } + + ops_after_partitioner_vgf_no_quantize = { + "executorch_exir_dialects_edge__ops_aten_argmax_default": 1, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2, + "torch.ops.higher_order.executorch_call_delegate": 2, + } + ops_after_partitioner_vgf_quantize = { + "executorch_exir_dialects_edge__ops_aten_argmax_default": 1, + "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2, + "torch.ops.higher_order.executorch_call_delegate": 1, + } + + def create_dummy_inputs( + self, + config, + batch_size: int = 1, + seq_length: int = 2, + dtype: torch.dtype = torch.long, + ) -> tuple[torch.Tensor]: + """Create dummy inputs for the CLIPTextModelWithProjection tests.""" + # SD3.5 Large uses (batch_size, seq_length) = (1, 77) for both CLIP-L + # and CLIP-bigG. Keep this unit-test default smaller for TOSA runtime. + return ( + torch.randint( + low=0, + high=config.vocab_size, + size=(batch_size, seq_length), + dtype=dtype, + ), + ) + + def create_model( + self, + config, + ) -> SD3CLIPTextEncoderWrapper: + """Instantiate wrapped CLIPTextModelWithProjection for tests.""" + return SD3CLIPTextEncoderWrapper( + CLIPTextModelWithProjection(config).to(dtype=config.dtype) # type: ignore[call-arg] + ).eval() + + @staticmethod + def ops_after_partitioner_INT(config) -> dict[str, int]: + if config.num_hidden_layers == 2: + return { + "executorch_exir_dialects_edge__ops_aten_add_Tensor": 2, + "executorch_exir_dialects_edge__ops_aten_argmax_default": 1, + "executorch_exir_dialects_edge__ops_aten_where_self": 2, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 12, + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 18, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 12, + "torch.ops.higher_order.executorch_call_delegate": 7, + } + + raise ValueError( + f"Unexpected CLIP config: hidden_act={config.hidden_act}, " + f"num_hidden_layers={config.num_hidden_layers}" + ) + + +@pytest.mark.parametrize( + ("config_factory", "atol"), + ( + (get_tiny_sd35_large_text_encoder_config, 1e-2), # FP atol + (get_tiny_sd35_large_text_encoder_2_config, 1.5e-2), # FP atol + ), + ids=("text_encoder", "text_encoder_2"), +) +def test_clip_text_model_with_projection_tosa_FP(config_factory, atol): + """Run the CLIPTextModelWithProjection TOSA FP test for a given config.""" + test_helper = TestCLIPTextModelWithProjection() + config = config_factory() + + with torch.no_grad(): + pipeline = TosaPipelineFP[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=atol, + transform_passes=get_int64_to_int32_passes(), + ) + pipeline.change_args( + "check_count.exir", TestCLIPTextModelWithProjection.ops_after_partitioner_FP + ) + pipeline.run() + + +@pytest.mark.parametrize( + ("config_factory", "atol"), + ( + (get_tiny_sd35_large_text_encoder_config, 5.5e-2), # INT atol + (get_tiny_sd35_large_text_encoder_2_config, 6e-2), # INT atol + ), + ids=("text_encoder", "text_encoder_2"), +) +def test_clip_text_model_with_projection_tosa_INT(config_factory, atol): + """Run the CLIPTextModelWithProjection TOSA INT test for a given config.""" + test_helper = TestCLIPTextModelWithProjection() + config = config_factory() + + with torch.no_grad(): + pipeline = TosaPipelineINT[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=atol, + frobenius_threshold=None, + cosine_threshold=None, + ) + pipeline.change_args( + "check_count.exir", + TestCLIPTextModelWithProjection.ops_after_partitioner_INT(config), + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +@pytest.mark.parametrize( + ("config_factory",), + ( + (get_tiny_sd35_large_text_encoder_config,), + (get_tiny_sd35_large_text_encoder_2_config,), + ), + ids=("text_encoder", "text_encoder_2"), +) +def test_clip_text_model_with_projection_vgf_no_quant(config_factory): + """Run the CLIPTextModelWithProjection VGF no-quant test.""" + test_helper = TestCLIPTextModelWithProjection() + config = config_factory() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=5e-3, + transform_passes=get_int64_to_int32_passes(), + quantize=False, + ) + pipeline.change_args( + "check_count.exir", + TestCLIPTextModelWithProjection.ops_after_partitioner_vgf_no_quantize, + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +@pytest.mark.parametrize( + ("config_factory", "atol"), + ( + (get_tiny_sd35_large_text_encoder_config, 5.5e-2), + (get_tiny_sd35_large_text_encoder_2_config, 6e-2), + ), + ids=("text_encoder", "text_encoder_2"), +) +def test_clip_text_model_with_projection_vgf_quant(config_factory, atol): + """Run the CLIPTextModelWithProjection VGF quant test.""" + test_helper = TestCLIPTextModelWithProjection() + config = config_factory() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=atol, + quantize=True, + ) + pipeline.change_args( + "check_count.exir", + TestCLIPTextModelWithProjection.ops_after_partitioner_vgf_quantize, + ) + pipeline.run() diff --git a/backends/arm/test/models/stable_diffusion_3_5_large/test_SD3Transformer2DModel_sd35_large.py b/backends/arm/test/models/stable_diffusion_3_5_large/test_SD3Transformer2DModel_sd35_large.py new file mode 100644 index 00000000000..842c40ffb31 --- /dev/null +++ b/backends/arm/test/models/stable_diffusion_3_5_large/test_SD3Transformer2DModel_sd35_large.py @@ -0,0 +1,172 @@ +# 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 pytest +import torch +from executorch.backends.arm.test import common +from executorch.backends.arm.test.models.stable_diffusion_3_5_large.test_configs_sd35_large import ( + get_tiny_sd35_large_transformer_config, +) +from executorch.backends.arm.test.tester.test_pipeline import ( + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) +from executorch.examples.models.stable_diffusion_3_5_large.model import ( + SD3TransformerWrapper, +) + +input_t4 = Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + + +class TestSD3Transformer2DModel: + """Test helper for SD3.5 Large SD3Transformer2DModel config.""" + + ops_after_partitioner_FP = { + "executorch_exir_dialects_edge__ops_aten_unsqueeze_copy_default": 1, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1, + "torch.ops.higher_order.executorch_call_delegate": 1, + } + + ops_after_partitioner_INT = { + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1, + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 3, + "torch.ops.higher_order.executorch_call_delegate": 1, + } + + ops_after_partitioner_vgf_quantize = { + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1, + "torch.ops.higher_order.executorch_call_delegate": 1, + } + ops_after_partitioner_vgf_no_quantize = ops_after_partitioner_FP + + def create_config(self): + """Create a tiny SD3.5 Large-like MMDiT config for tests.""" + return get_tiny_sd35_large_transformer_config() + + def create_dummy_inputs( + self, + batch_size: int = 2, + latent_channels: int = 4, + latent_size: int = 32, + seq_length: int = 77, + joint_attention_dim: int = 16, + pooled_projection_dim: int = 32, + max_timestep: int = 1000, + dtype: torch.dtype = torch.float32, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Create dummy inputs for the SD3Transformer2DModel tests.""" + # SD3.5 Large uses latent channels=16, latent size=128, T5 seq length=256, + # joint_attention_dim=4096, and pooled_projection_dim=2048. Keep this + # unit-test default smaller for TOSA runtime and VGF memory limits. + return ( + torch.randn( + batch_size, + latent_channels, + latent_size, + latent_size, + dtype=dtype, + ), + torch.randint(low=0, high=max_timestep, size=(batch_size,)), + torch.randn( + batch_size, + seq_length, + joint_attention_dim, + dtype=dtype, + ), + torch.randn(batch_size, pooled_projection_dim, dtype=dtype), + ) + + def create_model(self) -> SD3TransformerWrapper: + """Instantiate wrapped SD3Transformer2DModel for tests.""" + SD3Transformer2DModel = pytest.importorskip( + "diffusers.models.transformers" + ).SD3Transformer2DModel + return SD3TransformerWrapper( + SD3Transformer2DModel(**self.create_config()) + ).eval() + + +def test_sd3_transformer_tosa_FP(): + """Run the SD3Transformer2DModel TOSA FP test.""" + test_helper = TestSD3Transformer2DModel() + + with torch.no_grad(): + pipeline = TosaPipelineFP[input_t4]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + ) + pipeline.change_args( + "check_count.exir", TestSD3Transformer2DModel.ops_after_partitioner_FP + ) + pipeline.run() + + +def test_sd3_transformer_tosa_INT(): + """Run the SD3Transformer2DModel TOSA INT test.""" + test_helper = TestSD3Transformer2DModel() + + with torch.no_grad(): + pipeline = TosaPipelineINT[input_t4]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + frobenius_threshold=None, + cosine_threshold=None, + ) + pipeline.change_args( + "check_count.exir", TestSD3Transformer2DModel.ops_after_partitioner_INT + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_sd3_transformer_vgf_no_quant(): + """Run the SD3Transformer2DModel VGF no-quant test.""" + test_helper = TestSD3Transformer2DModel() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t4]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + quantize=False, + ) + pipeline.change_args( + "check_count.exir", + TestSD3Transformer2DModel.ops_after_partitioner_vgf_no_quantize, + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_sd3_transformer_vgf_quant(): + """Run the SD3Transformer2DModel VGF quant test.""" + test_helper = TestSD3Transformer2DModel() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t4]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + quantize=True, + ) + pipeline.change_args( + "check_count.exir", + TestSD3Transformer2DModel.ops_after_partitioner_vgf_quantize, + ) + pipeline.run() diff --git a/backends/arm/test/models/stable_diffusion_3_5_large/test_T5EncoderModel_sd35_large.py b/backends/arm/test/models/stable_diffusion_3_5_large/test_T5EncoderModel_sd35_large.py new file mode 100644 index 00000000000..c6e9d268248 --- /dev/null +++ b/backends/arm/test/models/stable_diffusion_3_5_large/test_T5EncoderModel_sd35_large.py @@ -0,0 +1,173 @@ +# 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.models.stable_diffusion_3_5_large.test_configs_sd35_large import ( + get_int64_to_int32_passes, + get_tiny_sd35_large_t5_config, +) +from executorch.backends.arm.test.tester.test_pipeline import ( + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) +from executorch.examples.models.stable_diffusion_3_5_large.model import ( + SD3T5TextEncoderWrapper, +) +from transformers import T5EncoderModel + +input_t = Tuple[torch.Tensor] + + +class TestT5EncoderModel: + """Test helper for SD3.5 Large T5EncoderModel config.""" + + ops_after_partitioner_FP = { + "executorch_exir_dialects_edge__ops_aten_clamp_Tensor": 4, + "executorch_exir_dialects_edge__ops_aten_where_self": 1, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2, + "torch.ops.higher_order.executorch_call_delegate": 6, + } + + ops_after_partitioner_INT = { + "executorch_exir_dialects_edge__ops_aten_isinf_default": 4, + "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 5, + "executorch_exir_dialects_edge__ops_aten_where_self": 5, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 21, + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 30, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 24, + "aten.scalar_tensor.default": 9, + "torch.ops.higher_order.executorch_call_delegate": 24, + } + + ops_after_partitioner_vgf_quantize = { + "executorch_exir_dialects_edge__ops_aten_clamp_Tensor": 4, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1, + "torch.ops.higher_order.executorch_call_delegate": 5, + } + + ops_after_partitioner_vgf_no_quantize = ops_after_partitioner_vgf_quantize + + def create_dummy_inputs( + self, + config, + batch_size: int = 1, + seq_length: int = 2, + dtype: torch.dtype = torch.long, + ) -> tuple[torch.Tensor]: + """Create dummy inputs for the T5EncoderModel tests.""" + # SD3.5 Large uses (batch_size, seq_length) = (1, 256) for T5. + # Keep this unit-test default smaller for TOSA runtime. + return ( + torch.randint( + low=0, + high=config.vocab_size, + size=(batch_size, seq_length), + dtype=dtype, + ), + ) + + def create_config(self): + """Create a tiny SD3.5 Large-like T5 config for tests.""" + return get_tiny_sd35_large_t5_config() + + def create_model(self, config) -> SD3T5TextEncoderWrapper: + """Instantiate wrapped T5EncoderModel for tests.""" + return SD3T5TextEncoderWrapper( + T5EncoderModel(config).to(dtype=config.dtype) # type: ignore[call-arg] + ).eval() + + +def test_t5_encoder_tosa_FP(): + """Run the T5EncoderModel TOSA FP test.""" + test_helper = TestT5EncoderModel() + config = test_helper.create_config() + + with torch.no_grad(): + pipeline = TosaPipelineFP[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=2e-2, + transform_passes=get_int64_to_int32_passes(), + ) + pipeline.change_args( + "check_count.exir", TestT5EncoderModel.ops_after_partitioner_FP + ) + pipeline.run() + + +def test_t5_encoder_tosa_INT(): + """Run the T5EncoderModel TOSA INT test.""" + test_helper = TestT5EncoderModel() + config = test_helper.create_config() + + with torch.no_grad(): + pipeline = TosaPipelineINT[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=3e-2, + frobenius_threshold=None, + cosine_threshold=None, + ) + pipeline.change_args( + "check_count.exir", TestT5EncoderModel.ops_after_partitioner_INT + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_t5_encoder_vgf_no_quant(): + """Run the T5EncoderModel VGF no-quant test.""" + test_helper = TestT5EncoderModel() + config = test_helper.create_config() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=8e-3, + transform_passes=get_int64_to_int32_passes(), + quantize=False, + ) + pipeline.change_args( + "check_count.exir", + TestT5EncoderModel.ops_after_partitioner_vgf_no_quantize, + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_t5_encoder_vgf_quant(): + """Run the T5EncoderModel VGF quant test.""" + test_helper = TestT5EncoderModel() + config = test_helper.create_config() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=1.2e-2, + quantize=True, + ) + pipeline.change_args( + "check_count.exir", + TestT5EncoderModel.ops_after_partitioner_vgf_quantize, + ) + pipeline.run() diff --git a/backends/arm/test/models/stable_diffusion_3_5_large/test_configs_sd35_large.py b/backends/arm/test/models/stable_diffusion_3_5_large/test_configs_sd35_large.py new file mode 100644 index 00000000000..d4de15e5d74 --- /dev/null +++ b/backends/arm/test/models/stable_diffusion_3_5_large/test_configs_sd35_large.py @@ -0,0 +1,311 @@ +# 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 hashlib +import json +import os +import warnings +from pathlib import Path +from typing import Any + +from executorch.backends.arm._passes import ( + ArmPass, + ConvertInt64ConstOpsToInt32Pass, + ConvertInt64OutputOpsToInt32Pass, + InsertInt32CastsAfterInt64PlaceholdersPass, +) +from transformers import CLIPTextConfig, T5Config + + +_EXECUTORCH_SD35_UPSTREAM_SYNC_ENV_VAR = "EXECUTORCH_SD35_UPSTREAM_SYNC" +_sd35_large_upstream_checked: set[str] = set() + +_SD35_LARGE_UPSTREAM_FINGERPRINTS: dict[str, tuple[tuple[str, ...], str]] = { + "text_encoder": ( + ( + "hidden_act", + "hidden_size", + "intermediate_size", + "max_position_embeddings", + "num_attention_heads", + "num_hidden_layers", + "projection_dim", + "vocab_size", + ), + "0cb164627ea428c39166d8ecf70462f7ddba34a57071d383a1116603e114bf50", + ), + "text_encoder_2": ( + ( + "hidden_act", + "hidden_size", + "intermediate_size", + "max_position_embeddings", + "num_attention_heads", + "num_hidden_layers", + "projection_dim", + "vocab_size", + ), + "4c082c890625573879240ef700b1dbe18e2a263ad8e76d222cba759553176e1e", + ), + "text_encoder_3": ( + ( + "d_ff", + "d_kv", + "d_model", + "dense_act_fn", + "feed_forward_proj", + "num_heads", + "num_layers", + "relative_attention_num_buckets", + "vocab_size", + ), + "da069a817a2fe4eb347fc5aefd7690bf97f8661d33dbe689058977809651023d", + ), + "transformer": ( + ( + "sample_size", + "patch_size", + "in_channels", + "num_layers", + "attention_head_dim", + "num_attention_heads", + "caption_projection_dim", + "joint_attention_dim", + "pooled_projection_dim", + "out_channels", + "pos_embed_max_size", + "qk_norm", + ), + "bc87c7eefb80f7bc6e80479f5bd2af929fd14c2331a19ddb4e27a989546ebfb0", + ), + "vae": ( + ( + "sample_size", + "in_channels", + "out_channels", + "down_block_types", + "up_block_types", + "block_out_channels", + "layers_per_block", + "latent_channels", + "norm_num_groups", + "act_fn", + "mid_block_add_attention", + "force_upcast", + "use_quant_conv", + "use_post_quant_conv", + "scaling_factor", + "shift_factor", + ), + "8019be48e6681895b9b7c54c0f2f48c3ddc9975d7b6c7ef2061b7743f8d55b71", + ), +} + + +def _load_upstream_sd35_large_config(subfolder: str) -> dict[str, Any]: + from executorch.examples.models.stable_diffusion_3_5_large.model import MODEL_ID + from huggingface_hub import hf_hub_download + + config_path = hf_hub_download( # nosec B615 + repo_id=MODEL_ID, + filename="config.json", + subfolder=subfolder, + token=os.environ.get("HF_TOKEN"), + etag_timeout=1.0, + ) + return json.loads(Path(config_path).read_text()) + + +def _fingerprint_sd35_large_config( + config: dict[str, Any], fields: tuple[str, ...] +) -> str: + selected_config = {field: config[field] for field in fields} + serialized_config = json.dumps( + selected_config, + sort_keys=True, + separators=(",", ":"), + ) + return hashlib.sha256(serialized_config.encode()).hexdigest() + + +def _warn_if_sd35_large_config_differs_from_upstream( + subfolder: str, warning_name: str +) -> None: + if os.environ.get(_EXECUTORCH_SD35_UPSTREAM_SYNC_ENV_VAR, "0") != "1": + return + if subfolder in _sd35_large_upstream_checked: + return + + _sd35_large_upstream_checked.add(subfolder) + try: + upstream_config = _load_upstream_sd35_large_config(subfolder) + fields, expected_fingerprint = _SD35_LARGE_UPSTREAM_FINGERPRINTS[subfolder] + upstream_fingerprint = _fingerprint_sd35_large_config(upstream_config, fields) + except Exception as exc: + warnings.warn( + f"Unable to validate {warning_name} against upstream metadata: {exc}", + RuntimeWarning, + stacklevel=2, + ) + return + + if upstream_fingerprint != expected_fingerprint: + warnings.warn( + f"Upstream {warning_name} architecture changed; review the tiny test config", + RuntimeWarning, + stacklevel=2, + ) + + +def get_int64_to_int32_passes() -> list[ArmPass]: + return [ + ConvertInt64ConstOpsToInt32Pass(), + ConvertInt64OutputOpsToInt32Pass(), + InsertInt32CastsAfterInt64PlaceholdersPass(), + ] + + +def get_tiny_sd35_large_text_encoder_config() -> CLIPTextConfig: + """Create a tiny SD3.5 Large-like CLIP-L text encoder config for tests.""" + _warn_if_sd35_large_config_differs_from_upstream( + "text_encoder", "SD3.5 Large CLIP text encoder" + ) + return CLIPTextConfig( # type: ignore[call-arg] + architectures=["CLIPTextModelWithProjection"], + attention_dropout=0.0, + bos_token_id=0, + dropout=0.0, + eos_token_id=2, + hidden_act="quick_gelu", + hidden_size=32, + initializer_factor=1.0, + initializer_range=0.02, + intermediate_size=128, + layer_norm_eps=1e-5, + max_position_embeddings=16, + num_attention_heads=4, + num_hidden_layers=2, + pad_token_id=1, + projection_dim=32, + dtype="float16", + vocab_size=256, + ) + + +def get_tiny_sd35_large_text_encoder_2_config() -> CLIPTextConfig: + """Create a tiny SD3.5 Large-like CLIP-bigG text encoder config for + tests. + """ + _warn_if_sd35_large_config_differs_from_upstream( + "text_encoder_2", "SD3.5 Large CLIP text encoder 2" + ) + return CLIPTextConfig( # type: ignore[call-arg] + architectures=["CLIPTextModelWithProjection"], + attention_dropout=0.0, + bos_token_id=0, + dropout=0.0, + eos_token_id=2, + hidden_act="gelu", + hidden_size=48, + initializer_factor=1.0, + initializer_range=0.02, + intermediate_size=192, + layer_norm_eps=1e-5, + max_position_embeddings=16, + num_attention_heads=6, + num_hidden_layers=2, + pad_token_id=1, + projection_dim=48, + dtype="float16", + vocab_size=256, + ) + + +def get_tiny_sd35_large_t5_config() -> T5Config: + """Create a tiny SD3.5 Large-like T5 config for tests.""" + _warn_if_sd35_large_config_differs_from_upstream( + "text_encoder_3", "SD3.5 Large T5 text encoder" + ) + return T5Config( # type: ignore[call-arg] + architectures=["T5EncoderModel"], + classifier_dropout=0.0, + d_ff=64, + d_kv=8, + d_model=32, + decoder_start_token_id=0, + dense_act_fn="gelu_new", + dropout_rate=0.1, + eos_token_id=1, + feed_forward_proj="gated-gelu", + initializer_factor=1.0, + is_encoder_decoder=True, + is_gated_act=True, + layer_norm_epsilon=1e-6, + num_decoder_layers=2, + num_heads=4, + num_layers=2, + output_past=True, + pad_token_id=0, + relative_attention_max_distance=128, + relative_attention_num_buckets=8, + tie_word_embeddings=False, + dtype="float16", + vocab_size=256, + use_cache=True, + ) + + +def get_tiny_sd35_large_transformer_config() -> dict[str, Any]: + """Create a tiny SD3.5 Large-like MMDiT config for tests.""" + _warn_if_sd35_large_config_differs_from_upstream( + "transformer", "SD3.5 Large transformer" + ) + return { + "sample_size": 32, + "patch_size": 2, + "in_channels": 4, + "num_layers": 2, + "attention_head_dim": 8, + "num_attention_heads": 2, + "caption_projection_dim": 16, + "joint_attention_dim": 16, + "pooled_projection_dim": 32, + "out_channels": 4, + "pos_embed_max_size": 32, + "qk_norm": "rms_norm", + } + + +def get_tiny_sd35_large_vae_config() -> dict[str, Any]: + """Create a tiny SD3.5 Large-like VAE config for tests.""" + _warn_if_sd35_large_config_differs_from_upstream("vae", "SD3.5 Large VAE") + return { + "sample_size": 32, + "in_channels": 3, + "out_channels": 3, + "down_block_types": ( + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + ), + "up_block_types": ( + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + ), + "block_out_channels": (4, 8, 8, 8), + "layers_per_block": 1, + "latent_channels": 16, + "norm_num_groups": 1, + "act_fn": "silu", + "mid_block_add_attention": True, + "force_upcast": False, + "use_quant_conv": False, + "use_post_quant_conv": False, + "scaling_factor": 1.5305, + "shift_factor": 0.0609, + } diff --git a/backends/arm/test/models/stable_diffusion_3_5_large/test_model_sd35_large.py b/backends/arm/test/models/stable_diffusion_3_5_large/test_model_sd35_large.py new file mode 100644 index 00000000000..336e1f8e1b6 --- /dev/null +++ b/backends/arm/test/models/stable_diffusion_3_5_large/test_model_sd35_large.py @@ -0,0 +1,369 @@ +# 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 types import SimpleNamespace + +import pytest +import torch +from executorch.backends.arm.test.models.stable_diffusion_3_5_large.test_configs_sd35_large import ( + get_tiny_sd35_large_t5_config, + get_tiny_sd35_large_text_encoder_2_config, + get_tiny_sd35_large_text_encoder_config, + get_tiny_sd35_large_transformer_config, + get_tiny_sd35_large_vae_config, +) +from executorch.examples.models.stable_diffusion_3_5_large import ( + model as sd35_large_model, +) +from transformers import CLIPTextModelWithProjection, T5EncoderModel + + +@pytest.mark.parametrize( + ("clip_skip", "hidden_state_index"), + ( + pytest.param(None, -2, id="default_clip_skip"), + pytest.param(1, -3, id="clip_skip_1"), + ), +) +def test_clip_text_encoder_wrapper_returns_selected_hidden_state_and_pooled_projection( + clip_skip, hidden_state_index +): + """Verify CLIP wrapper outputs.""" + config = get_tiny_sd35_large_text_encoder_config() + config.num_hidden_layers = 3 # Set to 3 layers to test clip_skip=1 behavior + text_encoder = CLIPTextModelWithProjection(config).to(dtype=config.dtype) + text_encoder.eval() + wrapper = sd35_large_model.SD3CLIPTextEncoderWrapper( + text_encoder, clip_skip=clip_skip + ) + input_ids = torch.randint(0, config.vocab_size, (2, 7)) + + with torch.no_grad(): + hidden_states, pooled_projection = wrapper(input_ids) + expected = text_encoder(input_ids, output_hidden_states=True, return_dict=True) + + torch.testing.assert_close( + hidden_states, expected.hidden_states[hidden_state_index] + ) + torch.testing.assert_close(pooled_projection, expected[0]) + + +def test_t5_text_encoder_wrapper_returns_last_hidden_state(): + """Verify T5 text encoder wrapper returns last hidden state.""" + config = get_tiny_sd35_large_t5_config() + text_encoder = T5EncoderModel(config) + text_encoder.eval() + wrapper = sd35_large_model.SD3T5TextEncoderWrapper(text_encoder) + input_ids = torch.randint(0, config.vocab_size, (2, 7)) + + with torch.no_grad(): + hidden_states = wrapper(input_ids) + expected = text_encoder(input_ids, return_dict=True) + + torch.testing.assert_close(hidden_states, expected.last_hidden_state) + + +def test_transformer_wrapper_returns_sample_tensor(): + """Verify transformer wrapper returns the sample tensor.""" + SD3Transformer2DModel = pytest.importorskip( + "diffusers.models.transformers" + ).SD3Transformer2DModel + transformer = SD3Transformer2DModel(**get_tiny_sd35_large_transformer_config()) + transformer.eval() + wrapper = sd35_large_model.SD3TransformerWrapper(transformer) + batch_size = 2 + latents = torch.randn(batch_size, 4, 32, 32) + timestep = torch.randint(0, 1000, (batch_size,)) + encoder_hidden_states = torch.randn(batch_size, 154, 16) + pooled_projections = torch.randn(batch_size, 32) + + with torch.no_grad(): + sample = wrapper( + latents, + timestep, + encoder_hidden_states, + pooled_projections, + ) + expected = transformer( + hidden_states=latents, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + pooled_projections=pooled_projections, + return_dict=True, + ) + + torch.testing.assert_close(sample, expected.sample) + + +def test_vae_decoder_wrapper_scales_shifts_decodes_and_clamps(): + """Verify VAE decoder wrapper scales, shifts, decodes, and clamps.""" + AutoencoderKL = pytest.importorskip("diffusers.models.autoencoders").AutoencoderKL + vae_config = get_tiny_sd35_large_vae_config() + vae = AutoencoderKL(**vae_config) + vae.eval() + wrapper = sd35_large_model.SD3VAEDecoderWrapper(vae) + latents = torch.randn(1, vae_config["latent_channels"], 4, 4) + + with torch.no_grad(): + image = wrapper(latents) + expected_latents = latents / vae.config.scaling_factor + vae.config.shift_factor + expected = vae.decode(expected_latents, return_dict=True).sample + # Normalize decoder output from [-1, 1] to image range [0, 1]. + expected = (expected / 2 + 0.5).clamp(0, 1) + + torch.testing.assert_close(image, expected) + assert torch.all(image >= 0) + assert torch.all(image <= 1) + + +@pytest.mark.parametrize( + "getter_name", + ( + "get_text_encoder_wrapper", + "get_text_encoder_2_wrapper", + "get_text_encoder_3_wrapper", + "get_transformer_wrapper", + "get_vae_decoder_wrapper", + ), +) +def test_model_loader_getters_require_loaded_components(getter_name): + """Verify model loader getters require loaded components.""" + loader = sd35_large_model.StableDiffusion3ModelLoader(dtype=torch.float32) + + with pytest.raises(ValueError, match="Models not loaded"): + getattr(loader, getter_name)() + + +def test_model_loader_text_encoder_getters_wrap_loaded_components(): + """Verify model loader text encoder getters wrap loaded components.""" + loader = sd35_large_model.StableDiffusion3ModelLoader(dtype=torch.float32) + loader.text_encoder = CLIPTextModelWithProjection( + get_tiny_sd35_large_text_encoder_config() + ) + loader.text_encoder_2 = CLIPTextModelWithProjection( + get_tiny_sd35_large_text_encoder_2_config() + ) + loader.text_encoder_3 = T5EncoderModel(get_tiny_sd35_large_t5_config()) + + assert loader.get_text_encoder_wrapper().text_encoder is loader.text_encoder + assert loader.get_text_encoder_2_wrapper().text_encoder is loader.text_encoder_2 + assert loader.get_text_encoder_3_wrapper().text_encoder is loader.text_encoder_3 + + +def test_model_loader_transformer_getter_wraps_loaded_component(): + """Verify model loader transformer getter wraps loaded component.""" + SD3Transformer2DModel = pytest.importorskip( + "diffusers.models.transformers" + ).SD3Transformer2DModel + loader = sd35_large_model.StableDiffusion3ModelLoader(dtype=torch.float32) + loader.transformer = SD3Transformer2DModel( + **get_tiny_sd35_large_transformer_config() + ) + + assert loader.get_transformer_wrapper().transformer is loader.transformer + + +def test_model_loader_vae_getter_wraps_loaded_component(): + """Verify model loader VAE getter wraps loaded component.""" + AutoencoderKL = pytest.importorskip("diffusers.models.autoencoders").AutoencoderKL + loader = sd35_large_model.StableDiffusion3ModelLoader(dtype=torch.float32) + loader.vae = AutoencoderKL(**get_tiny_sd35_large_vae_config()) + + assert loader.get_vae_decoder_wrapper().vae is loader.vae + + +def _patch_model_loaders(monkeypatch): + """Patch model loaders and return call records.""" + + class FakeModel: + def __init__(self): + self.eval_called = False + + def to(self, dtype): + return self + + def eval(self): + self.eval_called = True + return self + + calls = SimpleNamespace( + tokenizer=[], + text_encoder=[], + t5=[], + transformer=[], + vae=[], + ) + + class FakeTokenizer: + @staticmethod + def from_pretrained(model_id, **kwargs): + calls.tokenizer.append((model_id, kwargs)) + return SimpleNamespace(model_max_length=77) + + class FakeTextEncoder: + @staticmethod + def from_pretrained(model_id, **kwargs): + calls.text_encoder.append((model_id, kwargs)) + return FakeModel() + + class FakeT5: + @staticmethod + def from_pretrained(model_id, **kwargs): + calls.t5.append((model_id, kwargs)) + return FakeModel() + + class FakeTransformer: + @staticmethod + def from_pretrained(model_id, **kwargs): + calls.transformer.append((model_id, kwargs)) + return FakeModel() + + class FakeVAE: + @staticmethod + def from_pretrained(model_id, **kwargs): + calls.vae.append((model_id, kwargs)) + return FakeModel() + + monkeypatch.setattr(sd35_large_model, "CLIPTokenizer", FakeTokenizer) + monkeypatch.setattr( + sd35_large_model, "CLIPTextModelWithProjection", FakeTextEncoder + ) + monkeypatch.setattr(sd35_large_model, "T5EncoderModel", FakeT5) + monkeypatch.setattr(sd35_large_model, "SD3Transformer2DModel", FakeTransformer) + monkeypatch.setattr(sd35_large_model, "AutoencoderKL", FakeVAE) + + return calls + + +def test_load_models_uses_component_subfolders(monkeypatch): + """Verify model loading uses the expected component subfolders.""" + calls = _patch_model_loaders(monkeypatch) + + loader = sd35_large_model.StableDiffusion3ModelLoader( + model_id="test/sd3", + dtype=torch.float32, + ) + + assert loader.load_models() + assert calls.tokenizer == [ + ("test/sd3", {"subfolder": "tokenizer"}), + ("test/sd3", {"subfolder": "tokenizer_2"}), + ] + assert calls.text_encoder == [ + ("test/sd3", {"subfolder": "text_encoder", "torch_dtype": torch.float32}), + ("test/sd3", {"subfolder": "text_encoder_2", "torch_dtype": torch.float32}), + ] + assert calls.t5 == [ + ("test/sd3", {"subfolder": "text_encoder_3", "torch_dtype": torch.float32}) + ] + assert calls.transformer == [ + ("test/sd3", {"subfolder": "transformer", "torch_dtype": torch.float32}) + ] + assert calls.vae == [ + ("test/sd3", {"subfolder": "vae", "torch_dtype": torch.float32}) + ] + assert loader.text_encoder.eval_called + assert loader.text_encoder_2.eval_called + assert loader.text_encoder_3.eval_called + assert loader.transformer.eval_called + assert loader.vae.eval_called + + +def test_load_models_loads_only_requested_component(monkeypatch): + """Verify model loading can load only requested components.""" + calls = _patch_model_loaders(monkeypatch) + + loader = sd35_large_model.StableDiffusion3ModelLoader( + model_id="test/sd3", + dtype=torch.float32, + ) + + assert loader.load_models( + [sd35_large_model.StableDiffusionComponent.TEXT_ENCODER_3] + ) + assert calls.tokenizer == [] + assert calls.text_encoder == [] + assert calls.t5 == [ + ("test/sd3", {"subfolder": "text_encoder_3", "torch_dtype": torch.float32}) + ] + assert calls.transformer == [] + assert calls.vae == [] + assert loader.text_encoder is None + assert loader.text_encoder_2 is None + assert loader.text_encoder_3 is not None + assert loader.transformer is None + assert loader.vae is None + + +@pytest.mark.parametrize( + ("latent_size", "expected_latent_size"), + ( + pytest.param(None, 32, id="default_latent_size"), + pytest.param(16, 16, id="override_latent_size"), + ), +) +def test_get_dummy_inputs_builds_expected_component_inputs( + latent_size, expected_latent_size +): + """Verify dummy inputs have expected component shapes.""" + loader = sd35_large_model.StableDiffusion3ModelLoader(dtype=torch.float32) + loader.tokenizer = SimpleNamespace(model_max_length=77) + loader.text_encoder = object() + loader.text_encoder_2 = object() + loader.text_encoder_3 = object() + loader.transformer = SimpleNamespace( + config=SimpleNamespace( + in_channels=4, + sample_size=32, + joint_attention_dim=16, + pooled_projection_dim=32, + ) + ) + loader.vae = SimpleNamespace(config=SimpleNamespace(latent_channels=16)) + + dummy_inputs = loader.get_dummy_inputs( + max_sequence_length=256, + latent_size=latent_size, + ) + + assert set(dummy_inputs) == { + sd35_large_model.StableDiffusionComponent.TEXT_ENCODER, + sd35_large_model.StableDiffusionComponent.TEXT_ENCODER_2, + sd35_large_model.StableDiffusionComponent.TEXT_ENCODER_3, + sd35_large_model.StableDiffusionComponent.TRANSFORMER, + sd35_large_model.StableDiffusionComponent.VAE_DECODER, + } + assert dummy_inputs[sd35_large_model.StableDiffusionComponent.TEXT_ENCODER][ + 0 + ].shape == (1, 77) + assert ( + dummy_inputs[sd35_large_model.StableDiffusionComponent.TEXT_ENCODER][0].dtype + == torch.long + ) + assert dummy_inputs[sd35_large_model.StableDiffusionComponent.TEXT_ENCODER_2][ + 0 + ].shape == (1, 77) + assert dummy_inputs[sd35_large_model.StableDiffusionComponent.TEXT_ENCODER_3][ + 0 + ].shape == (1, 256) + + transformer_inputs = dummy_inputs[ + sd35_large_model.StableDiffusionComponent.TRANSFORMER + ] + assert transformer_inputs[0].shape == ( + 1, + 4, + expected_latent_size, + expected_latent_size, + ) + assert transformer_inputs[0].dtype == torch.float32 + assert transformer_inputs[1].shape == (1,) + assert transformer_inputs[1].dtype == torch.float32 + assert transformer_inputs[2].shape == (1, 333, 16) + assert transformer_inputs[3].shape == (1, 32) + + assert dummy_inputs[sd35_large_model.StableDiffusionComponent.VAE_DECODER][ + 0 + ].shape == (1, 16, expected_latent_size, expected_latent_size) diff --git a/backends/arm/test/models/stable_diffusion_3_5_large/test_vae_AutoencoderKL_sd35_large.py b/backends/arm/test/models/stable_diffusion_3_5_large/test_vae_AutoencoderKL_sd35_large.py new file mode 100644 index 00000000000..30c6df3e693 --- /dev/null +++ b/backends/arm/test/models/stable_diffusion_3_5_large/test_vae_AutoencoderKL_sd35_large.py @@ -0,0 +1,152 @@ +# 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 pytest +import torch +from executorch.backends.arm.test import common +from executorch.backends.arm.test.models.stable_diffusion_3_5_large.test_configs_sd35_large import ( + get_tiny_sd35_large_vae_config, +) +from executorch.backends.arm.test.tester.test_pipeline import ( + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) +from executorch.examples.models.stable_diffusion_3_5_large.model import ( + SD3VAEDecoderWrapper, +) + +input_t = Tuple[torch.Tensor] + + +class TestAutoencoderKL: + """Test helper for SD3.5 Large AutoencoderKL config.""" + + ops_after_partitioner_FP = { + "torch.ops.higher_order.executorch_call_delegate": 1, + } + + ops_after_partitioner_INT = { + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 1, + "torch.ops.higher_order.executorch_call_delegate": 1, + } + + ops_after_partitioner_vgf_quantize = ops_after_partitioner_FP + ops_after_partitioner_vgf_no_quantize = ops_after_partitioner_FP + + def create_config(self): + """Create a tiny SD3.5 Large-like AutoencoderKL config for tests.""" + return get_tiny_sd35_large_vae_config() + + def create_dummy_inputs( + self, + batch_size: int = 1, + latent_channels: int = 16, + latent_size: int = 4, + dtype: torch.dtype = torch.float32, + ) -> tuple[torch.Tensor]: + """Create dummy inputs for the SD3 VAE decoder tests.""" + # SD3.5 Large uses VAE decoder latent channels=16 and latent size=128. + # Keep this unit-test default spatial size smaller for TOSA runtime. + return ( + torch.randn( + batch_size, + latent_channels, + latent_size, + latent_size, + dtype=dtype, + ), + ) + + def create_model(self) -> SD3VAEDecoderWrapper: + """Instantiate wrapped AutoencoderKL decoder for tests.""" + diffusers_autoencoders = pytest.importorskip("diffusers.models.autoencoders") + AutoencoderKL = diffusers_autoencoders.AutoencoderKL + return SD3VAEDecoderWrapper(AutoencoderKL(**self.create_config())).eval() + + +def test_vae_tosa_FP(): + """Run the AutoencoderKL TOSA FP test.""" + test_helper = TestAutoencoderKL() + + with torch.no_grad(): + pipeline = TosaPipelineFP[input_t]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + ) + pipeline.change_args( + "check_count.exir", TestAutoencoderKL.ops_after_partitioner_FP + ) + pipeline.run() + + +def test_vae_tosa_INT(): + """Run the AutoencoderKL TOSA INT test.""" + test_helper = TestAutoencoderKL() + + with torch.no_grad(): + pipeline = TosaPipelineINT[input_t]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=9e-2, + frobenius_threshold=None, + cosine_threshold=None, + ) + pipeline.change_args( + "check_count.exir", TestAutoencoderKL.ops_after_partitioner_INT + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_vae_vgf_no_quant(): + """Run the AutoencoderKL VGF no-quant test.""" + test_helper = TestAutoencoderKL() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + quantize=False, + ) + pipeline.change_args( + "check_count.exir", + TestAutoencoderKL.ops_after_partitioner_vgf_no_quantize, + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_vae_vgf_quant(): + """Run the AutoencoderKL VGF quant test.""" + test_helper = TestAutoencoderKL() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + quantize=True, + qtol=2, + ) + pipeline.change_args( + "check_count.exir", + TestAutoencoderKL.ops_after_partitioner_vgf_quantize, + ) + pipeline.run() diff --git a/backends/arm/test/models/test_deit_tiny_arm.py b/backends/arm/test/models/test_deit_tiny_arm.py index bfbef6cb835..35462e6e6da 100644 --- a/backends/arm/test/models/test_deit_tiny_arm.py +++ b/backends/arm/test/models/test_deit_tiny_arm.py @@ -119,6 +119,28 @@ def test_deit_tiny_tosa_FP_remove_sdpa_safe_softmax_guard(deit_tiny): pipeline.run() +def test_deit_tiny_tosa_FP_auto_remove_sdpa_safe_softmax_guard(deit_tiny): + pipeline = TosaPipelineFP[input_t]( + deit_tiny, + model_inputs, + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + ) + pipeline.tester.compile_spec.set_pass_pipeline_config( + ArmPassPipelineConfig(sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO) + ) + pipeline.count_tosa_ops( + { + "EQUAL": 0, + "LOGICAL_NOT": 0, + "REDUCE_ANY": 0, + "SELECT": 0, + } + ) + pipeline.run() + + def test_deit_tiny_tosa_INT(deit_tiny): pipeline = TosaPipelineINT[input_t]( deit_tiny, diff --git a/backends/arm/test/models/test_nss.py b/backends/arm/test/models/test_nss.py index 61523964021..f3f1e0d54eb 100644 --- a/backends/arm/test/models/test_nss.py +++ b/backends/arm/test/models/test_nss.py @@ -3,10 +3,17 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import os +from pathlib import Path from typing import Tuple import pytest import torch +from executorch.backends.arm.scripts.neural_graphics_test_data import ( + iter_calibration_samples, + load_verification_inputs, + nss_test_calibration_path, +) from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( @@ -19,12 +26,30 @@ from huggingface_hub import hf_hub_download -from ng_model_gym.usecases.nss.model.model_blocks import ( # type: ignore[import-not-found,import-untyped] +from ng_model_gym.usecases.nss.model.model_blocks_v1 import ( # type: ignore[import-not-found,import-untyped] AutoEncoderV1, ) +from torch.export import Dim input_t = Tuple[torch.Tensor] # Input x +_RELEASE_REFS = ( + os.environ.get("GITHUB_REF", ""), + os.environ.get("GITHUB_REF_NAME", ""), + os.environ.get("GITHUB_BASE_REF", ""), +) +_IS_FROZEN_RELEASE = any( + ref.removeprefix("refs/heads/").startswith("release/") for ref in _RELEASE_REFS +) +pytestmark = pytest.mark.skipif( + _IS_FROZEN_RELEASE, + reason="NSS tests depend on resources fetched from main.", +) + +_NSS_HEIGHT = 8 * Dim("_nss_height", min=16, max=68) +_NSS_WIDTH = 8 * Dim("_nss_width", min=16, max=120) +_NSS_QUANTIZATION_DYNAMIC_SHAPES = ({2: _NSS_HEIGHT, 3: _NSS_WIDTH},) + class NSS(torch.nn.Module): def __init__(self, *args, **kwargs): @@ -35,46 +60,87 @@ def __init__(self, *args, **kwargs): def nss() -> AutoEncoderV1: """Get an instance of NSS with weights loaded.""" - weights = hf_hub_download( + weights = hf_hub_download( # nosec B615 repo_id="Arm/neural-super-sampling", - filename="nss_v0.1.0_fp32.pt", - revision="2e9b606acd9fa25071825a12f0764f1c3bef9480", + filename="nss_v1_0_1_high_fp32.pt", + revision="main", ) - nss_model = NSS() - nss_model.load_state_dict( - torch.load(weights, map_location=torch.device("cpu"), weights_only=True), - strict=False, + checkpoint = torch.load( + weights, map_location=torch.device("cpu"), weights_only=True ) + state_dict = { + f"auto_encoder.{key.removeprefix('autoencoder.')}": value + for key, value in checkpoint["model_state_dict"].items() + } + + nss_model = NSS() + nss_model.load_state_dict(state_dict, strict=True) return nss_model.auto_encoder def example_inputs(): - return (torch.randn((1, 12, 544, 960)),) + return load_verification_inputs() + + +def random_inputs(): + return (torch.rand((1, 12, 544, 960)),) + + +input_test_data = { + "real_data": True, + "random_data": False, +} -def test_nss_tosa_FP(): +def _nss_calibration_path() -> Path: + path = nss_test_calibration_path() + if not path.exists(): + raise RuntimeError( + "NSS calibration data is prepared by " + "backends/arm/scripts/install_models_for_test.sh." + ) + return path + + +def _set_nss_calibration_samples(pipeline): + quantize_stage = pipeline._stages[pipeline.find_pos("quantize")].args[0] + quantize_stage.dynamic_shapes = _NSS_QUANTIZATION_DYNAMIC_SHAPES + quantize_stage.calibration_samples = iter_calibration_samples( + _nss_calibration_path(), num_samples=3663 + ) + return pipeline + + +@common.parametrize("use_real_data", input_test_data) +def test_nss_tosa_FP(use_real_data): pipeline = TosaPipelineFP[input_t]( nss().eval(), - example_inputs(), + example_inputs() if use_real_data else random_inputs(), aten_op=[], exir_op=[], use_to_edge_transform_and_lower=True, ) - pipeline.add_stage_after("export", pipeline.tester.dump_operator_distribution) + if use_real_data: + pipeline.add_stage_after("export", pipeline.tester.dump_operator_distribution) pipeline.run() -def test_nss_tosa_INT(): +@common.parametrize("use_real_data", input_test_data) +def test_nss_tosa_INT(use_real_data): + pipeline_kwargs = ( + {"frobenius_threshold": 0.32, "qtol": 12} if use_real_data else {"qtol": 7} + ) pipeline = TosaPipelineINT[input_t]( nss().eval(), - example_inputs(), + example_inputs() if use_real_data else random_inputs(), aten_op=[], exir_op=[], use_to_edge_transform_and_lower=True, - frobenius_threshold=None, - cosine_threshold=None, + **pipeline_kwargs, ) + if use_real_data: + _set_nss_calibration_samples(pipeline) pipeline.run() @@ -89,6 +155,7 @@ def test_nss_u55_INT(): run_on_fvp=True, use_to_edge_transform_and_lower=True, ) + _set_nss_calibration_samples(pipeline) pipeline.run() @@ -105,17 +172,16 @@ def test_nss_u85_INT(): run_on_fvp=True, use_to_edge_transform_and_lower=True, ) + _set_nss_calibration_samples(pipeline) pipeline.run() -@pytest.mark.xfail( - reason="[MLETORCH-1430]: Double types are not supported in buffers in MSL" -) @common.SkipIfNoModelConverter -def test_nss_vgf_FP(): +@common.parametrize("use_real_data", input_test_data) +def test_nss_vgf_FP(use_real_data): pipeline = VgfPipeline[input_t]( nss().eval(), - example_inputs(), + example_inputs() if use_real_data else random_inputs(), aten_op=[], exir_op=[], use_to_edge_transform_and_lower=True, @@ -128,10 +194,11 @@ def test_nss_vgf_FP(): @common.SkipIfNoModelConverter -def test_nss_vgf_INT(): +@common.parametrize("use_real_data", input_test_data) +def test_nss_vgf_INT(use_real_data): pipeline = VgfPipeline[input_t]( nss().eval(), - example_inputs(), + example_inputs() if use_real_data else random_inputs(), aten_op=[], exir_op=[], symmetric_io_quantization=True, @@ -140,9 +207,8 @@ def test_nss_vgf_INT(): quantize=True, # Override tosa version to test INT-only path tosa_version="TOSA-1.0+INT", + qtol=12 if use_real_data else 7, ) + if use_real_data: + _set_nss_calibration_samples(pipeline) pipeline.run() - - -ModelUnderTest = nss().eval() -ModelInputs = example_inputs() diff --git a/backends/arm/test/modules/test_static_cache.py b/backends/arm/test/modules/test_static_cache.py index 86649f1e589..20ddac33f79 100644 --- a/backends/arm/test/modules/test_static_cache.py +++ b/backends/arm/test/modules/test_static_cache.py @@ -12,6 +12,7 @@ InsertInt32CastsAfterInt64PlaceholdersPass, ) from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.arm_tester import ToExecutorch from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, @@ -19,6 +20,11 @@ TosaPipelineINT, VgfPipeline, ) +from executorch.examples.models.llama.source_transformation.custom_kv_cache import ( + StaticQuantizedKVCache, +) +from executorch.exir import ExecutorchBackendConfig +from executorch.exir.passes.init_mutable_pass import InitializedMutableBufferPass from torch.export.graph_signature import InputKind, OutputKind from transformers import LlamaConfig @@ -41,39 +47,23 @@ InputKind.USER_INPUT: 3, } -EXPECTED_STATIC_QUANTIZED_INPUT_COUNTS = { - InputKind.BUFFER: 4, - InputKind.USER_INPUT: 3, -} - EXPECTED_OUTPUT_COUNTS = { OutputKind.BUFFER_MUTATION: 2, OutputKind.USER_OUTPUT: 2, } -DYNAMIC_KVQ_OPS = [ - "torch.ops.quantized_decomposed.choose_qparams_per_token_asymmetric.default", - "torch.ops.quantized_decomposed.quantize_per_token.default", - "torch.ops.quantized_decomposed.dequantize_per_token.default", - "torch.ops.llama.update_cache.default", - "torch.ops.llama.update_cache_with_indices.default", -] - - -def _reject_dynamic_kvq_ops(pipeline): - pipeline.add_stage_after( - "export", pipeline.tester.check_not, DYNAMIC_KVQ_OPS, suffix="dynamic_kvq_ops" +def _initialize_cache_buffers(pipeline, pattern: list[str]) -> None: + pipeline.change_args( + "to_executorch", + ToExecutorch( + ExecutorchBackendConfig(passes=[InitializedMutableBufferPass(pattern)]) + ), ) @torch.no_grad() class StaticQuantizedCacheModule(torch.nn.Module): - key_cache: torch.Tensor - value_cache: torch.Tensor - key_scale: torch.Tensor - value_scale: torch.Tensor - def __init__( self, config: LlamaConfig, @@ -90,40 +80,23 @@ def __init__( self.hidden_size = hidden_size self.num_attention_heads = num_attention_heads self.head_dim = self.hidden_size // self.num_attention_heads - cache_shape = (1, self.num_attention_heads, max_cache_len, self.head_dim) - scale_shape = (1, 1, 1, self.head_dim) - - self.register_buffer("key_cache", torch.zeros(cache_shape, dtype=torch.int8)) - self.register_buffer("value_cache", torch.zeros(cache_shape, dtype=torch.int8)) - self.register_buffer( - "key_scale", torch.full(scale_shape, scale, dtype=torch.float32) - ) - self.register_buffer( - "value_scale", torch.full(scale_shape, scale, dtype=torch.float32) + self.cache = StaticQuantizedKVCache( + max_batch_size=1, + max_context_length=max_cache_len, + n_heads=self.num_attention_heads, + head_dim=self.head_dim, + scale=scale, + use_custom_update_cache_op=False, + use_per_channel=False, ) - # PT2E activation quantization does not create persistent int8 mutable buffers. - def _quantize(self, value: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: - return torch.clamp(torch.round(value / scale), -128, 127).to(torch.int8) - def forward( self, key_states: torch.Tensor, value_states: torch.Tensor, cache_position: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: - key_q = self._quantize(key_states, self.key_scale) - value_q = self._quantize(value_states, self.value_scale) - - self.key_cache[:, :, cache_position] = key_q - self.value_cache[:, :, cache_position] = value_q - - key = self.key_cache.to(torch.float32) * self.key_scale - value = self.value_cache.to(torch.float32) * self.value_scale - key[:, :, cache_position] = key_states - value[:, :, cache_position] = value_states - - return key.clone(), value.clone() + return self.cache.update(cache_position, key_states, value_states) def get_inputs(self) -> input_t: key_states = torch.randn( @@ -236,17 +209,18 @@ def test_static_cache_tosa_FP(test_data): exir_op=[], transform_passes=[InsertInt32CastsAfterInt64PlaceholdersPass()], ) + _initialize_cache_buffers(pipeline, ["cache_layer_"]) pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS) pipeline.run() -@pytest.mark.xfail(reason="BUFFER_MUTATION count mismatch: MLETORCH-1971") @common.parametrize("test_data", test_configs) def test_static_cache_tosa_INT(test_data): - module = StaticCacheModule(test_data).eval() + module = StaticQuantizedCacheModule(test_data).eval() pipeline = TosaPipelineINT[input_t]( - module, module.get_inputs(), aten_op=[], exir_op=[], fold_quantize=False + module, module.get_inputs(), aten_op=[], exir_op=[] ) + _initialize_cache_buffers(pipeline, ["k_cache", "v_cache"]) pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS) pipeline.run() @@ -255,41 +229,26 @@ def test_static_cache_tosa_INT(test_data): @pytest.mark.xfail(reason="Scatter operator is not supported on U55.") @common.parametrize("test_data", test_configs) def test_static_cache_u55_INT(test_data): - module = StaticCacheModule(test_data).eval() + module = StaticQuantizedCacheModule(test_data).eval() pipeline = EthosU55PipelineINT[input_t]( module, module.get_inputs(), aten_ops=[], ) + _initialize_cache_buffers(pipeline, ["k_cache", "v_cache"]) pipeline.run() -@common.parametrize( - "test_data", - test_configs, - xfails={ - "multihead_attention": ( - "BUFFER_MUTATION count mismatch: MLETORCH-1971" - "Incorrect numerical behavior: MLBEDSW-11589" - ), - "grouped_query_attention": ( - "BUFFER_MUTATION count mismatch: MLETORCH-1971" - "Incorrect numerical behavior: MLBEDSW-11589" - ), - "multi_query_attention": ( - "BUFFER_MUTATION count mismatch: MLETORCH-1971" - "Incorrect numerical behavior: MLBEDSW-11589" - ), - }, -) +@common.XfailIfNoCorstone320 +@common.parametrize("test_data", test_configs) def test_static_cache_u85_INT(test_data): - module = StaticCacheModule(test_data).eval() + module = StaticQuantizedCacheModule(test_data).eval() pipeline = EthosU85PipelineINT[input_t]( module, module.get_inputs(), aten_ops=[], - fold_quantize=False, ) + _initialize_cache_buffers(pipeline, ["k_cache", "v_cache"]) # U85: keep _to_dim_order_copy portable for int64->int32 cast of cache_position (not delegatable). pipeline.tester.use_portable_ops = True pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS) @@ -308,72 +267,14 @@ def test_static_cache_vgf_no_quant(test_data): transform_passes=[InsertInt32CastsAfterInt64PlaceholdersPass()], quantize=False, ) + _initialize_cache_buffers(pipeline, ["cache_layer_"]) pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS) pipeline.run() @common.SkipIfNoModelConverter -@pytest.mark.xfail(reason="BUFFER_MUTATION count mismatch: MLETORCH-1971") @common.parametrize("test_data", test_configs) def test_static_cache_vgf_quant(test_data): - module = StaticCacheModule(test_data).eval() - pipeline = VgfPipeline[input_t]( - module, - module.get_inputs(), - aten_op=[], - exir_op=[], - quantize=True, - fold_quantize=False, - tosa_spec="TOSA-1.0+INT", - ) - pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS) - pipeline.run() - - -@common.parametrize("test_data", test_configs) -def test_static_quantized_cache_tosa_INT(test_data): - module = StaticQuantizedCacheModule(test_data).eval() - pipeline = TosaPipelineINT[input_t]( - module, module.get_inputs(), aten_op=[], exir_op=[], fold_quantize=False - ) - _reject_dynamic_kvq_ops(pipeline) - pipeline.change_args( - "check_count.exir", - {"torch.ops.higher_order.executorch_call_delegate": 2}, - ) - pipeline.count_program_io_kinds( - EXPECTED_STATIC_QUANTIZED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS - ) - pipeline.run() - - -@common.parametrize( - "test_data", - test_configs, - xfails={ - config: "Incorrect numerical behavior: MLBEDSW-11589" for config in test_configs - }, -) -def test_static_quantized_cache_u85_INT(test_data): - module = StaticQuantizedCacheModule(test_data).eval() - pipeline = EthosU85PipelineINT[input_t]( - module, module.get_inputs(), aten_ops=[], fold_quantize=False - ) - _reject_dynamic_kvq_ops(pipeline) - pipeline.change_args( - "check_count.exir", - {"torch.ops.higher_order.executorch_call_delegate": 2}, - ) - pipeline.tester.use_portable_ops = True - pipeline.count_program_io_kinds( - EXPECTED_STATIC_QUANTIZED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS - ) - pipeline.run() - - -@common.SkipIfNoModelConverter -@common.parametrize("test_data", test_configs) -def test_static_quantized_cache_vgf_quant(test_data): module = StaticQuantizedCacheModule(test_data).eval() pipeline = VgfPipeline[input_t]( module, @@ -381,12 +282,8 @@ def test_static_quantized_cache_vgf_quant(test_data): aten_op=[], exir_op=[], quantize=True, - fold_quantize=False, tosa_spec="TOSA-1.0+INT", - n_expected_delegates=2, - ) - _reject_dynamic_kvq_ops(pipeline) - pipeline.count_program_io_kinds( - EXPECTED_STATIC_QUANTIZED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS ) + _initialize_cache_buffers(pipeline, ["k_cache", "v_cache"]) + pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS) pipeline.run() diff --git a/backends/arm/test/ops/test_arange.py b/backends/arm/test/ops/test_arange.py index 90ab437b9e7..31165ba48a0 100644 --- a/backends/arm/test/ops/test_arange.py +++ b/backends/arm/test/ops/test_arange.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,6 +17,11 @@ TosaPipelineINT, VgfPipeline, ) +from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec +from executorch.backends.arm.tosa.partitioner import TOSAPartitioner +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.dialects._ops import ops as exir_ops +from torch._subclasses.fake_tensor import FakeTensorMode input_t = tuple[torch.Tensor] test_data_t = tuple[Callable[[], input_t], tuple[float, float, float, torch.dtype]] @@ -173,6 +178,75 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: } +class _LinspaceToFloatAdd(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + values = torch.linspace(0.0, 1.0, x.shape[0], dtype=torch.float64) + return values.to(torch.float32) + x + + +@pytest.mark.parametrize( + ("dtype", "expected"), + ((torch.float32, True), (torch.float64, False)), +) +def test_linspace_preservation_depends_on_dtype( + dtype: torch.dtype, expected: bool +) -> None: + exported_program = torch.export.export( + LinspaceAdd(0.0, 1.0, 10, dtype), + (torch.randn(10),), + ) + partitioner = TOSAPartitioner(TosaCompileSpec("TOSA-1.0+FP")) + _, filter_fn = partitioner.ops_to_not_decompose(exported_program) + linspace = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.linspace.default + ) + + assert filter_fn is not None + assert filter_fn(linspace) is expected + + +def test_ops_to_not_decompose_are_not_preserved_for_fp64() -> None: + exported_program = torch.export.export( + LinspaceAdd(0.0, 1.0, 10, torch.float32), + (torch.randn(10),), + ) + partitioner = TOSAPartitioner(TosaCompileSpec("TOSA-1.0+FP")) + ops_to_not_decompose, filter_fn = partitioner.ops_to_not_decompose(exported_program) + graph = torch.fx.Graph() + fake_tensor = FakeTensorMode().from_tensor(torch.empty(1, dtype=torch.float64)) + + assert filter_fn is not None + for op in ops_to_not_decompose: + node = graph.call_function(op) + node.meta["val"] = fake_tensor + assert not filter_fn(node), f"FP64 {op} should be decomposed" + + +def test_linspace_fp64_decomposes_for_portable_fallback() -> None: + inputs = (torch.randn(10),) + exported_program = torch.export.export(_LinspaceToFloatAdd(), inputs) + partitioner = TOSAPartitioner(TosaCompileSpec("TOSA-1.0+FP")) + + edge_manager = to_edge_transform_and_lower( + exported_program, + partitioner=[partitioner], + ) + targets = { + node.target + for node in edge_manager.exported_program().graph.nodes + if node.op == "call_function" + } + + assert exir_ops.edge.aten.linspace.default not in targets + assert exir_ops.edge.aten.arange.start_step in targets + + program = edge_manager.to_executorch().executorch_program + operators = {(op.name, op.overload) for op in program.execution_plan[0].operators} + assert not any("linspace" in name for name, _ in operators) + + @common.parametrize("test_data", LinspaceAdd.test_data) def test_linspace_tosa_FP(test_data: test_data_t): input_data, init_data = test_data diff --git a/backends/arm/test/ops/test_batch_norm.py b/backends/arm/test/ops/test_batch_norm.py index 4fb458ee918..b3df52c8f64 100644 --- a/backends/arm/test/ops/test_batch_norm.py +++ b/backends/arm/test/ops/test_batch_norm.py @@ -22,6 +22,7 @@ Input = Tuple[torch.Tensor] ATEN_BATCH_NORM = "torch.ops.aten.batch_norm.default" ATEN_CONV2D = "torch.ops.aten.conv2d.default" +ATEN_CONV_TRANSPOSE2D = "torch.ops.aten.conv_transpose2d.input" @dataclass(frozen=True) @@ -110,6 +111,30 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.batch_norm(self.conv2d(x)) +class BatchNorm2dConvTranspose(torch.nn.Module): + aten_ops = [ATEN_CONV_TRANSPOSE2D, ATEN_BATCH_NORM] + + def __init__(self, groups: int) -> None: + super().__init__() + self.conv_transpose2d = torch.nn.ConvTranspose2d( + in_channels=4, + out_channels=6, + kernel_size=3, + padding=1, + groups=groups, + ) + self.batch_norm = _make_batch_norm( + 6, + affine=True, + weight=torch.rand(6), + bias=torch.rand(6), + track_running_stats=True, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.batch_norm(self.conv_transpose2d(x)) + + class BatchNorm2dNoStats(torch.nn.Module): def __init__( self, @@ -205,6 +230,27 @@ def test_native_batch_norm_legit_no_training_tosa_FP_conv_fuses_before_decompose pipeline.run() +@common.parametrize("groups", {"groups=1": 1, "groups=2": 2}) +def test_conv_transpose_batch_norm_fuses_before_decompose_tosa_FP( + groups: int, +) -> None: + model = BatchNorm2dConvTranspose(groups) + pipeline = TosaPipelineFP[Input]( + model, + (torch.rand(1, 4, 5, 6),), + aten_op=model.aten_ops, + ) + pipeline.count_tosa_ops( + { + "TRANSPOSE_CONV2D": groups, + "CONCAT": int(groups > 1), + "RSQRT": 0, + "SUB": 0, + } + ) + pipeline.run() + + @common.parametrize("case", batch_norm_cases) def test_native_batch_norm_legit_no_training_tosa_INT_conv(case: BatchNormCase) -> None: test_data, model_params = case.make_input_and_parameters() @@ -254,6 +300,17 @@ def test_native_batch_norm_legit_no_training_vgf_no_quant_conv( ).run() +@common.SkipIfNoModelConverter +def test_grouped_conv_transpose_batch_norm_vgf_no_quant() -> None: + model = BatchNorm2dConvTranspose(groups=2) + VgfPipeline[Input]( + model, + (torch.rand(1, 4, 5, 6),), + aten_op=model.aten_ops, + quantize=False, + ).run() + + @common.parametrize("case", batch_norm_cases) @common.SkipIfNoModelConverter def test_native_batch_norm_legit_no_training_vgf_quant_conv( diff --git a/backends/arm/test/ops/test_conv2d.py b/backends/arm/test/ops/test_conv2d.py index 977ffdb9a7c..5b608135406 100644 --- a/backends/arm/test/ops/test_conv2d.py +++ b/backends/arm/test/ops/test_conv2d.py @@ -566,7 +566,8 @@ def _get_dtype_count(model: torch.nn.Module): # Set nbr_conv to be the amount of groups set if necessary. nbr_convs: int = model.nbr_convs if model.groups is None else model.groups # noqa return { - "CONST": {"INT4": nbr_convs * 2}, # One for the weight, one for the zp. + # Each convolution has a distinct weight and shares the symmetric zero point. + "CONST": {"INT4": nbr_convs + 1}, "CONV2D": {"INT32": nbr_convs}, "RESCALE": {"INT8": nbr_convs}, } diff --git a/backends/arm/test/ops/test_conv3d.py b/backends/arm/test/ops/test_conv3d.py index 09cd44525c5..a43121274ce 100644 --- a/backends/arm/test/ops/test_conv3d.py +++ b/backends/arm/test/ops/test_conv3d.py @@ -585,7 +585,8 @@ def forward(self, x): def _get_dtype_count(model: torch.nn.Module): nbr_convs: int = model.nbr_convs # noqa return { - "CONST": {"INT4": nbr_convs * 2}, + # Each convolution has a distinct weight and shares the symmetric zero point. + "CONST": {"INT4": nbr_convs + 1}, "CONV3D": {"INT32": nbr_convs}, "RESCALE": {"INT8": nbr_convs}, } diff --git a/backends/arm/test/ops/test_depthwise_conv.py b/backends/arm/test/ops/test_depthwise_conv.py index a81a656017f..3dd8598dddf 100644 --- a/backends/arm/test/ops/test_depthwise_conv.py +++ b/backends/arm/test/ops/test_depthwise_conv.py @@ -264,7 +264,8 @@ def _get_dtype_count(model: torch.nn.Module): nbr_convs: int = model.nbr_convs # noqa return { - "CONST": {"INT4": nbr_convs * 2}, + # Each convolution has a distinct weight and shares the symmetric zero point. + "CONST": {"INT4": nbr_convs + 1}, "DEPTHWISE_CONV2D": {"INT32": nbr_convs}, "RESCALE": {"INT8": nbr_convs}, } diff --git a/backends/arm/test/ops/test_grid_sampler.py b/backends/arm/test/ops/test_grid_sampler.py index c5a1f3560bd..4c8ab0de8ce 100644 --- a/backends/arm/test/ops/test_grid_sampler.py +++ b/backends/arm/test/ops/test_grid_sampler.py @@ -60,3 +60,18 @@ def test_grid_sampler_vgf_no_quant(test_data): run_on_vulkan_runtime=False, ) pipeline.run() + + +@common.parametrize("test_data", test_data_suite, xfails=xfails, strict=False) +@common.SkipIfNoModelConverter +def test_grid_sampler_vgf_quant(test_data): + test_data = test_data() + pipeline = VgfPipeline[input_t]( + GridSampler2d(), + test_data, + aten_op, + exir_op, + quantize=True, + run_on_vulkan_runtime=False, + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_index_select.py b/backends/arm/test/ops/test_index_select.py index 5410bc09a4e..0729b9dbd08 100644 --- a/backends/arm/test/ops/test_index_select.py +++ b/backends/arm/test/ops/test_index_select.py @@ -6,10 +6,19 @@ from typing import Tuple +import pytest import torch +from executorch.backends.arm._passes.decompose_index_select_to_gather_pass import ( + DecomposeIndexSelectToGatherPass, +) +from executorch.backends.arm.operator_support.ethos_u55_support import ( + EthosU55IndexSelectCheck, +) from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.arm_tester import ArmTester from executorch.backends.arm.test.tester.test_pipeline import ( + EthosU55PipelineINT, EthosU85PipelineINT, OpNotSupportedPipeline, TosaPipelineFP, @@ -17,6 +26,9 @@ VgfPipeline, ) +from executorch.exir.backend.utils import WhyNoPartitionReporter +from executorch.exir.dialects._ops import ops as exir_ops + class IndexSelect(torch.nn.Module): aten_op = "torch.ops.aten.index_select.default" @@ -26,6 +38,17 @@ def forward(self, input_: torch.Tensor, dim: int, index_: torch.Tensor): return torch.index_select(input_, dim=dim, index=index_) +class ConstantIndexSelect(torch.nn.Module): + def __init__(self, dim: int, indices: list[int], dtype: torch.dtype = torch.int32): + super().__init__() + self.dim = dim + self.register_buffer("indices", torch.tensor(indices, dtype=dtype)) + + def forward(self, input_: torch.Tensor): + return torch.index_select(input_, dim=self.dim, index=self.indices) + + +input_t1 = Tuple[torch.Tensor] input_params = Tuple[torch.Tensor, int, torch.Tensor] # ---- FP profile: only float inputs ---- @@ -229,3 +252,141 @@ def test_index_select_vgf_quant(test_data: input_params): quantize=True, ) pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_index_select_u55_INT_constant_contiguous(): + pipeline = EthosU55PipelineINT[input_t1]( + ConstantIndexSelect(2, [1, 2, 3]), + (torch.rand(1, 2, 5, 3),), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_index_select_u55_INT_constant_contiguous_negative_dim(): + pipeline = EthosU55PipelineINT[input_t1]( + ConstantIndexSelect(-1, [1, 2]), + (torch.rand(1, 2, 4, 5),), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.parametrize( + "indices", + { + "noncontiguous": [1, 3], + "descending": [3, 1], + "duplicate": [1, 1], + }, +) +@common.XfailIfNoCorstone300 +def test_index_select_u55_INT_constant_slices_concat(indices): + pipeline = EthosU55PipelineINT[input_t1]( + ConstantIndexSelect(2, indices), + (torch.rand(1, 2, 5, 3),), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +def test_index_select_u55_INT_constant_contiguous_symbolic_dim_not_delegated(): + selected_dim = torch.export.Dim("selected_dim", min=4, max=8) + tester = ArmTester( + ConstantIndexSelect(2, [1, 2, 3]), + (torch.rand(1, 2, 5, 3),), + common.get_u55_compile_spec(), + dynamic_shapes={"input_": {2: selected_dim}}, + ) + tester.quantize().export().to_edge().partition() + + targets = { + node.target + for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes + } + assert exir_ops.edge.aten.index_select.default in targets + assert torch.ops.higher_order.executorch_call_delegate not in targets + + +@common.parametrize( + "indices", + { + "negative_index": [-1], + "upper_bound_index": [5], + }, +) +def test_index_select_u55_constant_out_of_bounds_raises(indices): + tester = ArmTester( + ConstantIndexSelect(2, indices), + (torch.rand(1, 2, 5, 3),), + common.get_u55_compile_spec(), + ) + tester.export().to_edge() + exported_program = tester.stages[tester.cur].artifact.exported_program() + + with pytest.raises(RuntimeError, match="index_select index out of range"): + DecomposeIndexSelectToGatherPass(exported_program).call( + exported_program.graph_module + ) + + +def test_index_select_u55_scalar_not_supported(): + tester = ArmTester( + ConstantIndexSelect(0, [0]), + (torch.tensor(1.0),), + common.get_u55_compile_spec(), + ) + tester.export().to_edge() + exported_program = tester.stages[tester.cur].artifact.exported_program() + index_select_node = next( + node + for node in exported_program.graph.nodes + if node.target == exir_ops.edge.aten.index_select.default + ) + + assert not EthosU55IndexSelectCheck( + exported_program, WhyNoPartitionReporter() + ).is_node_supported({}, index_select_node) + + +def test_index_select_u55_INT_constant_int64_delegated(): + tester = ArmTester( + ConstantIndexSelect(2, [1, 2, 3], torch.int64), + (torch.rand(1, 2, 5, 3),), + common.get_u55_compile_spec(), + ) + tester.quantize().export().to_edge().partition() + + targets = { + node.target + for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes + } + assert torch.ops.higher_order.executorch_call_delegate in targets + + +@common.XfailIfNoCorstone300 +def test_index_select_u55_INT_constant_contiguous_a16w8(): + pipeline = EthosU55PipelineINT[input_t1]( + ConstantIndexSelect(2, [1, 2, 3]), + (torch.rand(1, 2, 5, 3),), + aten_ops=[], + exir_ops=[], + a16w8_quantization=True, + ) + pipeline.run() + + +def test_index_select_u55_INT_constant_empty_not_delegated(): + pipeline = OpNotSupportedPipeline[input_t1]( + ConstantIndexSelect(2, []), + (torch.rand(1, 2, 5, 3),), + {IndexSelect.exir_op: 1}, + quantize=True, + u55_subset=True, + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_index_tensor.py b/backends/arm/test/ops/test_index_tensor.py index de6a1ac5f6b..f6b17bf5c32 100644 --- a/backends/arm/test/ops/test_index_tensor.py +++ b/backends/arm/test/ops/test_index_tensor.py @@ -8,12 +8,15 @@ import torch from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.arm_tester import ArmTester from executorch.backends.arm.test.tester.test_pipeline import ( + EthosU55PipelineINT, OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, VgfPipeline, ) +from executorch.exir.dialects._ops import ops as exir_ops class IndexTensorTestCommon: @@ -41,6 +44,58 @@ def forward(self, x: torch.Tensor): return x[self.index] +class IndexTensorLeadingInt64Buffers(torch.nn.Module): + """NCHW indexing with leading full slices and int64 index buffers.""" + + def __init__(self): + super().__init__() + self.register_buffer("rows", torch.tensor([[0], [2]], dtype=torch.int64)) + self.register_buffer("columns", torch.tensor([[1, 3]], dtype=torch.int64)) + + def forward(self, x: torch.Tensor): + return x[:, :, self.rows, self.columns] + + +class ConstantIndexTensor(torch.nn.Module): + def __init__(self, indices: list[int]): + super().__init__() + self.register_buffer("index", torch.tensor(indices, dtype=torch.int32)) + + def forward(self, x: torch.Tensor): + return x[self.index] + + +class ConstantIndexTensorDim(torch.nn.Module): + def __init__(self, dim: int, indices: list[int]): + super().__init__() + self.dim = dim + self.register_buffer("index", torch.tensor(indices, dtype=torch.int32)) + + def forward(self, x: torch.Tensor): + indices = [slice(None)] * x.dim() + indices[self.dim] = self.index + return x[tuple(indices)] + + +class ConstantTensorIndex(torch.nn.Module): + def __init__(self, index: torch.Tensor): + super().__init__() + self.register_buffer("index", index) + + def forward(self, x: torch.Tensor): + return x[self.index] + + +class ConstantMultiIndexTensor(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("index_0", torch.tensor([0, 1], dtype=torch.int32)) + self.register_buffer("index_1", torch.tensor([1, 0], dtype=torch.int32)) + + def forward(self, x: torch.Tensor): + return x[self.index_0, self.index_1] + + def test_index_tensor_tosa_FP_int64_buffer_index(): # This mirrors torchvision Swin relative_position_bias_table[index]. The # int64 get_attr must be cast before index.Tensor is decomposed to GATHER: @@ -57,8 +112,73 @@ def test_index_tensor_tosa_FP_int64_buffer_index(): pipeline.run() +def test_index_tensor_tosa_FP_leading_full_slice_int64_buffer_indices(): + pipeline = TosaPipelineFP[Tuple[torch.Tensor]]( + IndexTensorLeadingInt64Buffers(), + (torch.rand(1, 2, 4, 5),), + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + atol=IndexTensorTestCommon.atol, + rtol=IndexTensorTestCommon.rtol, + ) + pipeline.count_tosa_ops({"GATHER": 1, "TRANSPOSE": 0}) + pipeline.run() + + input_params_slice = Tuple[torch.Tensor, int, int, str, Tuple[torch.Tensor]] input_params = Tuple[torch.Tensor, Tuple[torch.Tensor]] +input_t2 = Tuple[torch.Tensor, torch.Tensor] +input_t3 = Tuple[torch.Tensor, torch.Tensor, torch.Tensor] + + +class IndexTensorLeadingFullSlice(torch.nn.Module): + def forward(self, x: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + return x[:, index, :] + + +class IndexTensorLeadingFullSlicesNCHW(torch.nn.Module): + def forward( + self, x: torch.Tensor, rows: torch.Tensor, columns: torch.Tensor + ) -> torch.Tensor: + return x[:, :, rows, columns] + + +class IndexTensorLeadingFullSliceWithTrailingDim(torch.nn.Module): + def forward( + self, x: torch.Tensor, rows: torch.Tensor, columns: torch.Tensor + ) -> torch.Tensor: + return x[:, rows, columns, :] + + +leading_full_slice_test_data = { + "nchw_broadcast_indices": lambda: ( + IndexTensorLeadingFullSlicesNCHW(), + ( + torch.arange(2 * 3 * 4 * 5, dtype=torch.float32).reshape(2, 3, 4, 5), + torch.tensor([[0], [2]], dtype=torch.int32), + torch.tensor([[1, 3, 4]], dtype=torch.int32), + ), + ), + "leading_and_trailing_dims": lambda: ( + IndexTensorLeadingFullSliceWithTrailingDim(), + ( + torch.arange(2 * 4 * 5 * 3, dtype=torch.float32).reshape(2, 4, 5, 3), + torch.tensor([[0], [2]], dtype=torch.int32), + torch.tensor([[1, 3, 4]], dtype=torch.int32), + ), + ), +} + +zero_sized_test_data = { + "zero_sized_leading_dimension": ( + torch.empty(0, 3, 4), + torch.tensor([0, 2], dtype=torch.int32), + ), + "empty_index_tensor": ( + torch.rand(2, 3, 4), + torch.empty(0, dtype=torch.int32), + ), +} class IndexTensor_Ellipsis(torch.nn.Module): @@ -115,7 +235,6 @@ def forward( IndexTensor_Ellipsis.test_data_ellipsis, xfails={ # More info in index_tensor_support.py - "test_4d_ellipsis_before": "Ellipsis before index unsupported", "test_4d_ellipsis_middle": "Ellipsis before index unsupported", }, ) @@ -139,7 +258,6 @@ def test_index_tensor_tosa_FP_ellipsis(test_data: input_params): IndexTensor_Ellipsis.test_data_ellipsis, xfails={ # More info in index_tensor_support.py - "test_4d_ellipsis_before": "Ellipsis before index unsupported", "test_4d_ellipsis_middle": "Ellipsis before index unsupported", }, ) @@ -227,8 +345,6 @@ def forward( IndexTensor_Slice.test_data, xfails={ # More info in index_tensor_support.py - "test_4d_slice_before_1d_idx": "Slice before index unsupported", - "test_3d_slice_before_2d_idx": "Slice before index unsupported", "test_4d_slice_middle": "Slice before index unsupported", }, ) @@ -252,8 +368,6 @@ def test_index_tensor_tosa_FP_slice(test_data: input_params_slice): IndexTensor_Slice.test_data, xfails={ # More info in index_tensor_support.py - "test_4d_slice_before_1d_idx": "Slice before index unsupported", - "test_3d_slice_before_2d_idx": "Slice before index unsupported", "test_4d_slice_middle": "Slice before index unsupported", }, ) @@ -424,8 +538,7 @@ class IndexTensor(torch.nn.Module): ), } - # xfail - None (unsqueeze) unsupported - test_data_none: dict[input_params] = { + test_data_leading_none: dict[input_params] = { "test_3d_3_idx_with_none_before": ( torch.rand(12, 3, 7), ( @@ -441,6 +554,9 @@ class IndexTensor(torch.nn.Module): torch.randint(3, size=(12,), dtype=torch.int32), ), ), + } + + test_data_none: dict[input_params] = test_data_leading_none | { "test_3d_3_idx_with_none_around": ( torch.rand(12, 3, 7), ( @@ -530,25 +646,22 @@ def test_index_tensor_tosa_INT(test_data: input_params): IndexTensor.test_data_none, xfails={ # More info in index_tensor_support.py - "test_3d_3_idx_with_none_before": "None (Unsqueeze) unsupported", - "test_3d_3_idx_with_2_none_before": "None (Unsqueeze) unsupported", - "test_3d_3_idx_with_none_around": "None (Unsqueeze) unsupported", "test_3d_3_idx_with_none_middle": "None (Unsqueeze) unsupported", }, ) def test_index_tensor_tosa_FP_none(test_data: input_params): test_input = test_data with torch.no_grad(): - ( - TosaPipelineFP[input_params]( - IndexTensor(), - test_input, - IndexTensorTestCommon.aten_op, - IndexTensorTestCommon.exir_op, - atol=IndexTensorTestCommon.atol, - rtol=IndexTensorTestCommon.rtol, - ).run() + pipeline = TosaPipelineFP[input_params]( + IndexTensor(), + test_input, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + atol=IndexTensorTestCommon.atol, + rtol=IndexTensorTestCommon.rtol, ) + pipeline.count_tosa_ops({"GATHER": 1, "TRANSPOSE": 0}) + pipeline.run() @common.parametrize( @@ -556,23 +669,100 @@ def test_index_tensor_tosa_FP_none(test_data: input_params): IndexTensor.test_data_none, xfails={ # More info in index_tensor_support.py - "test_3d_3_idx_with_none_before": "None (Unsqueeze) unsupported", - "test_3d_3_idx_with_2_none_before": "None (Unsqueeze) unsupported", - "test_3d_3_idx_with_none_around": "None (Unsqueeze) unsupported", "test_3d_3_idx_with_none_middle": "None (Unsqueeze) unsupported", }, ) def test_index_tensor_tosa_INT_none(test_data: input_params): test_input = test_data with torch.no_grad(): - ( - TosaPipelineINT[input_params]( - IndexTensor(), - test_input, - IndexTensorTestCommon.aten_op, - IndexTensorTestCommon.exir_op, - ).run() + pipeline = TosaPipelineINT[input_params]( + IndexTensor(), + test_input, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, ) + pipeline.count_tosa_ops({"GATHER": 1, "TRANSPOSE": 0}) + pipeline.run() + + +@common.parametrize("test_data", leading_full_slice_test_data) +def test_index_tensor_tosa_FP_leading_full_slices(test_data): + model, test_inputs = test_data() + pipeline = TosaPipelineFP[input_t3]( + model, + test_inputs, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + atol=IndexTensorTestCommon.atol, + rtol=IndexTensorTestCommon.rtol, + ) + pipeline.count_tosa_ops({"GATHER": 1, "TRANSPOSE": 0}) + pipeline.run() + + +@common.parametrize("test_data", leading_full_slice_test_data) +def test_index_tensor_tosa_INT_leading_full_slices(test_data): + model, test_inputs = test_data() + pipeline = TosaPipelineINT[input_t3]( + model, + test_inputs, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + ) + pipeline.count_tosa_ops({"GATHER": 1, "TRANSPOSE": 0}) + pipeline.run() + + +@common.parametrize("test_data", zero_sized_test_data) +def test_index_tensor_zero_sized_not_delegated_tosa_FP(test_data: input_t2): + OpNotSupportedPipeline[input_t2]( + IndexTensorLeadingFullSlice(), + test_data, + {IndexTensorTestCommon.exir_op: 1}, + ).run() + + +@common.parametrize("test_data", IndexTensor.test_data_leading_none) +@common.SkipIfNoModelConverter +def test_index_tensor_vgf_leading_full_slices(test_data: input_params): + pipeline = VgfPipeline[input_params]( + IndexTensor(), + test_data, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + atol=IndexTensorTestCommon.atol, + rtol=IndexTensorTestCommon.rtol, + quantize=False, + ) + pipeline.run() + + +@common.parametrize("test_data", leading_full_slice_test_data) +@common.SkipIfNoModelConverter +def test_index_tensor_leading_full_slice_indexing_vgf_no_quant(test_data): + model, test_inputs = test_data() + VgfPipeline[input_t3]( + model, + test_inputs, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + atol=IndexTensorTestCommon.atol, + rtol=IndexTensorTestCommon.rtol, + quantize=False, + ).run() + + +@common.parametrize("test_data", leading_full_slice_test_data) +@common.SkipIfNoModelConverter +def test_index_tensor_leading_full_slice_indexing_vgf_quant(test_data): + model, test_inputs = test_data() + VgfPipeline[input_t3]( + model, + test_inputs, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + quantize=True, + ).run() @common.parametrize("test_data", IndexTensor.test_data_int | IndexTensor.test_data_fp) @@ -622,3 +812,116 @@ def test_index_tensor_vgf_quant(test_data: input_params): quantize=True, ) pipeline.run() + + +@common.parametrize( + "indices", + { + "contiguous": [1, 2, 3], + "noncontiguous": [1, 3], + "descending": [3, 1], + "duplicate": [1, 1], + "negative": [-1, -3], + }, +) +@common.XfailIfNoCorstone300 +def test_index_tensor_u55_INT_constant(indices): + pipeline = EthosU55PipelineINT[Tuple[torch.Tensor]]( + ConstantIndexTensor(indices), + (torch.rand(5, 2, 3),), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.parametrize( + "test_data", + { + "dim1_contiguous": (1, [1, 2, 3]), + "dim1_noncontiguous": (1, [1, 3]), + "dim2_descending": (2, [3, 1]), + "dim2_duplicate": (2, [1, 1]), + "dim2_negative": (2, [-1, -3]), + }, +) +@common.XfailIfNoCorstone300 +def test_index_tensor_u55_INT_constant_later_dim(test_data): + dim, indices = test_data + pipeline = EthosU55PipelineINT[Tuple[torch.Tensor]]( + ConstantIndexTensorDim(dim, indices), + (torch.rand(5, 5, 5),), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_index_tensor_u55_INT_constant_a16w8(): + pipeline = EthosU55PipelineINT[Tuple[torch.Tensor]]( + ConstantIndexTensor([1, 2, 3]), + (torch.rand(5, 2, 3),), + aten_ops=[], + exir_ops=[], + a16w8_quantization=True, + ) + pipeline.run() + + +def test_index_tensor_u55_INT_constant_empty_not_delegated(): + pipeline = OpNotSupportedPipeline[Tuple[torch.Tensor]]( + ConstantIndexTensor([]), + (torch.rand(5, 2, 3),), + {IndexTensorTestCommon.exir_op: 1}, + quantize=True, + u55_subset=True, + ) + pipeline.run() + + +def test_index_tensor_u55_INT_constant_multi_not_delegated(): + pipeline = OpNotSupportedPipeline[Tuple[torch.Tensor]]( + ConstantMultiIndexTensor(), + (torch.rand(5, 2, 3),), + {IndexTensorTestCommon.exir_op: 1}, + quantize=True, + u55_subset=True, + ) + pipeline.run() + + +@common.parametrize( + "index", + { + "multidimensional": torch.tensor([[0, 1]], dtype=torch.int32), + "boolean": torch.tensor([False, True, False, True, False]), + }, +) +def test_index_tensor_u55_INT_constant_shape_or_dtype_not_delegated(index): + pipeline = OpNotSupportedPipeline[Tuple[torch.Tensor]]( + ConstantTensorIndex(index), + (torch.rand(5, 2, 3),), + {IndexTensorTestCommon.exir_op: 1}, + quantize=True, + u55_subset=True, + ) + pipeline.run() + + +def test_index_tensor_u55_INT_constant_symbolic_dim_not_delegated(): + indexed_dim = torch.export.Dim("indexed_dim", min=4, max=8) + tester = ArmTester( + ConstantIndexTensor([1, 2, 3]), + (torch.rand(5, 2, 3),), + common.get_u55_compile_spec(), + dynamic_shapes={"x": {0: indexed_dim}}, + ) + tester.quantize().export().to_edge().partition() + + targets = { + node.target + for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes + } + assert exir_ops.edge.aten.index.Tensor in targets + assert torch.ops.higher_order.executorch_call_delegate not in targets diff --git a/backends/arm/test/ops/test_isinf.py b/backends/arm/test/ops/test_isinf.py new file mode 100644 index 00000000000..4ed3ef6bfb7 --- /dev/null +++ b/backends/arm/test/ops/test_isinf.py @@ -0,0 +1,79 @@ +# 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 collections.abc import Callable + +import torch + +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.test_pipeline import ( + OpNotSupportedPipeline, + TosaPipelineFP, + VgfPipeline, +) + +aten_op = "torch.ops.aten.isinf.default" +exir_op = "executorch_exir_dialects_edge__ops_aten_isinf_default" + + +class IsInf(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.isinf(x) + + +test_data_suite = { + "finite": lambda: torch.tensor([-1.0, 0.0, 3.14]), + "inf": lambda: torch.tensor([-float("inf"), 0.0, float("inf"), float("nan")]), + "integer": lambda: torch.tensor([-5, 0, 9], dtype=torch.int32), + "rank4": lambda: torch.tensor( + [[[[float("inf"), 0.0]]], [[[-float("inf"), float("nan")]]]] + ), +} + + +@common.parametrize( + "test_data", + {name: data for name, data in test_data_suite.items() if name != "integer"}, +) +def test_isinf_tosa_FP(test_data: Callable[[], torch.Tensor]) -> None: + TosaPipelineFP( + IsInf(), + (test_data(),), + aten_op, + exir_op, + ).run() + + +def test_isinf_tosa_FP_falls_back_for_integer() -> None: + OpNotSupportedPipeline( + IsInf(), + (test_data_suite["integer"](),), + {exir_op: 1}, + quantize=False, + ).run() + + +def test_isinf_tosa_INT_falls_back() -> None: + test_data = (test_data_suite["inf"](),) + pipeline = OpNotSupportedPipeline( + IsInf(), + test_data, + {exir_op: 1}, + quantize=True, + ) + quantize_stage = pipeline._stages[pipeline.find_pos("quantize")].args[0] + quantize_stage.calibration_samples = [(torch.ones_like(test_data[0]),)] + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_isinf_vgf_no_quant() -> None: + VgfPipeline( + IsInf(), + (test_data_suite["inf"](),), + aten_op, + exir_op, + quantize=False, + ).run() diff --git a/backends/arm/test/ops/test_isnan.py b/backends/arm/test/ops/test_isnan.py new file mode 100644 index 00000000000..0c6476c01ab --- /dev/null +++ b/backends/arm/test/ops/test_isnan.py @@ -0,0 +1,77 @@ +# 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 collections.abc import Callable + +import torch + +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.test_pipeline import ( + OpNotSupportedPipeline, + TosaPipelineFP, + VgfPipeline, +) + +aten_op = "torch.ops.aten.isnan.default" +exir_op = "executorch_exir_dialects_edge__ops_aten_isnan_default" + + +class IsNan(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.isnan(x) + + +test_data_suite = { + "finite": lambda: torch.tensor([-1.0, 0.0, 3.14]), + "nan": lambda: torch.tensor([float("nan"), 0.0, float("inf")]), + "integer": lambda: torch.tensor([-5, 0, 9], dtype=torch.int32), + "rank4": lambda: torch.tensor([[[[float("nan"), 0.0]]], [[[float("inf"), -3.0]]]]), +} + + +@common.parametrize( + "test_data", + {name: data for name, data in test_data_suite.items() if name != "integer"}, +) +def test_isnan_tosa_FP(test_data: Callable[[], torch.Tensor]) -> None: + TosaPipelineFP( + IsNan(), + (test_data(),), + aten_op, + exir_op, + ).run() + + +def test_isnan_tosa_FP_falls_back_for_integer() -> None: + OpNotSupportedPipeline( + IsNan(), + (test_data_suite["integer"](),), + {exir_op: 1}, + quantize=False, + ).run() + + +def test_isnan_tosa_INT_falls_back() -> None: + test_data = (test_data_suite["nan"](),) + pipeline = OpNotSupportedPipeline( + IsNan(), + test_data, + {exir_op: 1}, + quantize=True, + ) + quantize_stage = pipeline._stages[pipeline.find_pos("quantize")].args[0] + quantize_stage.calibration_samples = [(torch.ones_like(test_data[0]),)] + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_isnan_vgf_no_quant() -> None: + VgfPipeline( + IsNan(), + (test_data_suite["nan"](),), + aten_op, + exir_op, + quantize=False, + ).run() diff --git a/backends/arm/test/ops/test_max_pool.py b/backends/arm/test/ops/test_max_pool.py index 22dfe09b070..1dbe5536c31 100644 --- a/backends/arm/test/ops/test_max_pool.py +++ b/backends/arm/test/ops/test_max_pool.py @@ -82,6 +82,12 @@ [3, 2, 1], ), } + +test_data_suite_rank3 = { + "bev_class_slice": lambda: (torch.rand(1, 8, 8), [1, 1, 0]), + "spatial_pool": lambda: (torch.rand(3, 9, 11), [3, 2, 1]), +} + test_data_suite_fp8 = { "rand_fp8e4m3": lambda: ( torch.rand(1, 8, 20, 20).to(torch.float8_e4m3fn), @@ -168,7 +174,11 @@ def forward(self, x): @common.parametrize( - "test_data", test_data_suite | test_data_suite_fp16 | test_data_suite_bf16 + "test_data", + test_data_suite + | test_data_suite_rank3 + | test_data_suite_fp16 + | test_data_suite_bf16, ) def test_max_pool2d_tosa_FP(test_data: torch.Tensor): test_data, model_params = test_data() @@ -197,7 +207,7 @@ def test_max_pool2d_tosa_FP_fp8(test_data: torch.Tensor): pipeline.run() -@common.parametrize("test_data", test_data_suite) +@common.parametrize("test_data", test_data_suite | test_data_suite_rank3) def test_max_pool2d_tosa_INT(test_data: torch.Tensor): test_data, model_params = test_data() pipeline = TosaPipelineINT[input_t1]( @@ -374,22 +384,28 @@ def test_max_pool2d_tosa_INT_dilation(test_data): # VGF tests @common.parametrize( - "test_data", test_data_suite | test_data_suite_bf16 | test_data_suite_fp16 + "test_data", + test_data_suite + | test_data_suite_rank3 + | test_data_suite_bf16 + | test_data_suite_fp16, ) @common.SkipIfNoModelConverter def test_max_pool2d_vgf_no_quant(test_data: torch.Tensor): test_data, model_params = test_data() + run_on_vulkan_runtime = test_data.dim() == 4 pipeline = VgfPipeline[input_t1]( MaxPool2d(*model_params), (test_data,), aten_op, exir_op, quantize=False, + run_on_vulkan_runtime=run_on_vulkan_runtime, ) pipeline.run() -@common.parametrize("test_data", test_data_suite) +@common.parametrize("test_data", test_data_suite | test_data_suite_rank3) @common.SkipIfNoModelConverter def test_max_pool2d_vgf_quant(test_data: torch.Tensor): test_data, model_params = test_data() diff --git a/backends/arm/test/ops/test_pixel_shuffling.py b/backends/arm/test/ops/test_pixel_shuffling.py index 4980e24bab3..03f8e9a17fe 100644 --- a/backends/arm/test/ops/test_pixel_shuffling.py +++ b/backends/arm/test/ops/test_pixel_shuffling.py @@ -34,6 +34,10 @@ "rand_4d_channels_last": "Known U55 partitioning limitation for large 4D pixel shuffle layouts.", } +mixed_precision_xfails = { + "rand_4d_channels_last": "Permute propagation stops at f(x, g(x)) shapes such as the round decomposition.", +} + class PixelUnShuffle(nn.Module): @@ -110,7 +114,9 @@ def test_pixel_unshuffle_tosa_FP(test_data: input_t1): pipeline.run() -@common.parametrize("test_data", PixelUnShuffle.test_data_generators) +@common.parametrize( + "test_data", PixelUnShuffle.test_data_generators, xfails=mixed_precision_xfails +) def test_pixel_unshuffle_no_target_tosa_mixed_precision(test_data: input_t1): inputs, expected_transposes = test_data() pipeline = TosaPipelineINT[input_t1]( @@ -140,7 +146,9 @@ def test_pixel_shuffle_tosa_FP(test_data: input_t1): pipeline.run() -@common.parametrize("test_data", PixelShuffle.test_data_generators) +@common.parametrize( + "test_data", PixelShuffle.test_data_generators, xfails=mixed_precision_xfails +) def test_pixel_shuffle_no_target_tosa_mixed_precision(test_data: input_t1): inputs, expected_transposes = test_data() pipeline = TosaPipelineINT[input_t1]( diff --git a/backends/arm/test/ops/test_reflection_pad1d.py b/backends/arm/test/ops/test_reflection_pad1d.py new file mode 100644 index 00000000000..45a8f0ee9bd --- /dev/null +++ b/backends/arm/test/ops/test_reflection_pad1d.py @@ -0,0 +1,75 @@ +# 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.arm_tester import ArmTester +from executorch.backends.arm.test.tester.test_pipeline import EthosU55PipelineINT + + +input_t1 = Tuple[torch.Tensor] + +test_data_suite_u55 = { + "rank2_symmetric": lambda: (torch.rand(2, 5), (1, 1)), + "rank3_symmetric": lambda: (torch.rand(1, 2, 5), (1, 1)), + "asymmetric": lambda: (torch.rand(1, 2, 5), (1, 3)), + "maximum_legal": lambda: (torch.rand(1, 2, 5), (4, 4)), + "batched": lambda: (torch.rand(2, 2, 5), (1, 1)), +} + + +class ReflectionPad1d(torch.nn.Module): + def __init__(self, padding): + super().__init__() + self.padding = padding + + def forward(self, x): + return torch.nn.functional.pad(x, self.padding, mode="reflect") + + +@common.parametrize("test_data", test_data_suite_u55) +@common.XfailIfNoCorstone300 +def test_reflection_pad1d_u55_INT(test_data): + data, padding = test_data() + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad1d(padding), + (data,), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_reflection_pad1d_u55_INT_a16w8(): + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad1d((1, 1)), + (torch.rand(1, 2, 5),), + aten_ops=[], + exir_ops=[], + a16w8_quantization=True, + ) + pipeline.run() + + +def test_reflection_pad1d_u55_INT_symbolic_width_not_delegated(): + width = torch.export.Dim("width", min=3, max=8) + tester = ArmTester( + ReflectionPad1d((1, 1)), + (torch.rand(1, 2, 5),), + common.get_u55_compile_spec(), + dynamic_shapes={"x": {2: width}}, + ) + tester.quantize().export().to_edge().partition() + + targets = { + node.target + for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes + } + assert torch.ops.aten.scalar_tensor.default in targets + assert torch.ops.higher_order.executorch_call_delegate not in targets diff --git a/backends/arm/test/ops/test_reflection_pad2d.py b/backends/arm/test/ops/test_reflection_pad2d.py new file mode 100644 index 00000000000..a9b91bc64a8 --- /dev/null +++ b/backends/arm/test/ops/test_reflection_pad2d.py @@ -0,0 +1,87 @@ +# 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.arm_tester import ArmTester +from executorch.backends.arm.test.tester.test_pipeline import EthosU55PipelineINT + + +input_t1 = Tuple[torch.Tensor] + +test_data_suite_u55 = { + "both_axes": lambda: (torch.rand(1, 2, 4, 3), (1, 1, 1, 1)), + "width_only": lambda: (torch.rand(1, 2, 4, 3), (1, 1, 0, 0)), + "height_only": lambda: (torch.rand(1, 2, 4, 3), (0, 0, 1, 1)), + "asymmetric": lambda: (torch.rand(1, 2, 5, 4), (1, 2, 3, 1)), + "maximum_legal": lambda: (torch.rand(1, 2, 4, 3), (2, 2, 3, 3)), + "batched": lambda: (torch.rand(2, 2, 4, 3), (1, 1, 1, 1)), +} + + +class ReflectionPad2d(torch.nn.Module): + def __init__(self, padding): + super().__init__() + self.padding = padding + + def forward(self, x): + return torch.nn.functional.pad(x, self.padding, mode="reflect") + + +@common.parametrize("test_data", test_data_suite_u55) +@common.XfailIfNoCorstone300 +def test_reflection_pad2d_u55_INT(test_data): + data, padding = test_data() + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad2d(padding), + (data,), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_reflection_pad2d_u55_INT_a16w8(): + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad2d((1, 1, 1, 1)), + (torch.rand(1, 2, 4, 3),), + aten_ops=[], + exir_ops=[], + a16w8_quantization=True, + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_reflection_pad2d_u55_INT_rank3(): + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad2d((1, 1, 1, 1)), + (torch.rand(2, 4, 3),), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +def test_reflection_pad2d_u55_INT_symbolic_width_not_delegated(): + width = torch.export.Dim("width", min=3, max=8) + tester = ArmTester( + ReflectionPad2d((1, 1, 1, 1)), + (torch.rand(1, 2, 4, 5),), + common.get_u55_compile_spec(), + dynamic_shapes={"x": {3: width}}, + ) + tester.quantize().export().to_edge().partition() + + targets = { + node.target + for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes + } + assert torch.ops.aten.scalar_tensor.default in targets + assert torch.ops.higher_order.executorch_call_delegate not in targets diff --git a/backends/arm/test/ops/test_reflection_pad3d.py b/backends/arm/test/ops/test_reflection_pad3d.py new file mode 100644 index 00000000000..4bd343eafe9 --- /dev/null +++ b/backends/arm/test/ops/test_reflection_pad3d.py @@ -0,0 +1,75 @@ +# 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.arm_tester import ArmTester +from executorch.backends.arm.test.tester.test_pipeline import EthosU55PipelineINT + + +input_t1 = Tuple[torch.Tensor] + +test_data_suite_u55 = { + "rank4_symmetric": lambda: (torch.rand(2, 4, 4, 4), (1, 1, 1, 1, 1, 1)), + "rank5_symmetric": lambda: (torch.rand(1, 2, 4, 4, 4), (1, 1, 1, 1, 1, 1)), + "asymmetric": lambda: (torch.rand(1, 2, 5, 5, 5), (1, 2, 2, 1, 3, 1)), + "maximum_legal": lambda: (torch.rand(1, 2, 4, 4, 4), (3, 3, 3, 3, 3, 3)), + "batched": lambda: (torch.rand(2, 2, 4, 4, 4), (1, 1, 1, 1, 1, 1)), +} + + +class ReflectionPad3d(torch.nn.Module): + def __init__(self, padding): + super().__init__() + self.padding = padding + + def forward(self, x): + return torch.nn.functional.pad(x, self.padding, mode="reflect") + + +@common.parametrize("test_data", test_data_suite_u55) +@common.XfailIfNoCorstone300 +def test_reflection_pad3d_u55_INT(test_data): + data, padding = test_data() + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad3d(padding), + (data,), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_reflection_pad3d_u55_INT_a16w8(): + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad3d((1, 1, 1, 1, 1, 1)), + (torch.rand(1, 2, 4, 4, 4),), + aten_ops=[], + exir_ops=[], + a16w8_quantization=True, + ) + pipeline.run() + + +def test_reflection_pad3d_u55_INT_symbolic_width_not_delegated(): + width = torch.export.Dim("width", min=3, max=8) + tester = ArmTester( + ReflectionPad3d((1, 1, 1, 1, 1, 1)), + (torch.rand(1, 2, 4, 4, 5),), + common.get_u55_compile_spec(), + dynamic_shapes={"x": {4: width}}, + ) + tester.quantize().export().to_edge().partition() + + targets = { + node.target + for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes + } + assert torch.ops.aten.scalar_tensor.default in targets + assert torch.ops.higher_order.executorch_call_delegate not in targets diff --git a/backends/arm/test/ops/test_roll.py b/backends/arm/test/ops/test_roll.py new file mode 100644 index 00000000000..9d3dfd43ed4 --- /dev/null +++ b/backends/arm/test/ops/test_roll.py @@ -0,0 +1,119 @@ +# 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, +) + + +input_t1 = Tuple[torch.Tensor] +aten_op = "torch.ops.aten.roll.default" +exir_op = "executorch_exir_dialects_edge__ops_aten_roll_default" + +test_data_fp = { + "bev_cyclic_shift_fp32": lambda: ( + torch.randn(1, 8, 8, 4, dtype=torch.float32), + (-2, -2), + (1, 2), + ), + "bev_cyclic_shift_fp16": lambda: ( + torch.randn(1, 8, 8, 4, dtype=torch.float16), + (-2, -2), + (1, 2), + ), +} + +test_data_bf16 = { + "bev_cyclic_shift_bf16": lambda: ( + torch.randn(1, 8, 8, 4, dtype=torch.bfloat16), + (-2, -2), + (1, 2), + ), +} + +test_data_quant = { + "bev_cyclic_shift": lambda: ( + torch.randn(1, 8, 8, 4, dtype=torch.float32), + (-2, -2), + (1, 2), + ), +} + + +class Roll(torch.nn.Module): + def __init__(self, shifts: tuple[int, ...], dims: tuple[int, ...]) -> None: + super().__init__() + self.shifts = shifts + self.dims = dims + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.roll(x, self.shifts, self.dims) + + +@common.parametrize("test_data", test_data_fp) +def test_roll_tosa_FP(test_data) -> None: + data, shifts, dims = test_data() + pipeline = TosaPipelineFP[input_t1](Roll(shifts, dims), (data,), aten_op, exir_op) + pipeline.count_tosa_ops({"SLICE": 4, "CONCAT": 2}) + pipeline.run() + + +@common.parametrize("test_data", test_data_bf16) +def test_roll_tosa_FP_bf16(test_data) -> None: + data, shifts, dims = test_data() + pipeline = TosaPipelineFP[input_t1]( + Roll(shifts, dims), + (data,), + aten_op, + exir_op, + tosa_extensions=["bf16"], + ) + pipeline.count_tosa_ops({"SLICE": 4, "CONCAT": 2}) + pipeline.run() + + +@common.parametrize("test_data", test_data_quant) +def test_roll_tosa_INT(test_data) -> None: + data, shifts, dims = test_data() + pipeline = TosaPipelineINT[input_t1](Roll(shifts, dims), (data,), aten_op, exir_op) + pipeline.count_tosa_ops({"SLICE": 4, "CONCAT": 2}) + pipeline.run() + + +@common.parametrize("test_data", test_data_fp | test_data_bf16) +@common.SkipIfNoModelConverter +def test_roll_vgf_no_quant(test_data) -> None: + data, shifts, dims = test_data() + pipeline = VgfPipeline[input_t1]( + Roll(shifts, dims), + (data,), + aten_op, + exir_op, + quantize=False, + run_on_vulkan_runtime=False, + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_quant) +@common.SkipIfNoModelConverter +def test_roll_vgf_quant(test_data) -> None: + data, shifts, dims = test_data() + pipeline = VgfPipeline[input_t1]( + Roll(shifts, dims), + (data,), + aten_op, + exir_op, + quantize=True, + run_on_vulkan_runtime=False, + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_round.py b/backends/arm/test/ops/test_round.py index 1f4470cebb3..f06f7e37737 100644 --- a/backends/arm/test/ops/test_round.py +++ b/backends/arm/test/ops/test_round.py @@ -21,6 +21,7 @@ aten_op = "torch.ops.aten.round.default" exir_op = "executorch_exir_dialects_edge__ops_aten_round_default" + test_data_suite = { # (test_name, test_data) "zeros": lambda: torch.zeros(1, 10, 10, 10), @@ -29,6 +30,14 @@ "randn_pos": lambda: torch.randn(10) + 10, "randn_neg": lambda: torch.randn(10) - 10, "ramp": lambda: torch.arange(-16, 16, 0.2), + "halfway_ties": lambda: torch.arange(-8, 8, 0.5), +} + +# One ulp either side of a tie. +test_data_suite_fp = { + "halfway_neighbors": lambda: torch.nextafter( + torch.tensor([0.5, 0.5]), torch.tensor([-float("inf"), float("inf")]) + ), } test_data_suite_bf16 = { @@ -41,7 +50,9 @@ def forward(self, x: torch.Tensor): return x.round() -@common.parametrize("test_data", test_data_suite | test_data_suite_bf16) +@common.parametrize( + "test_data", test_data_suite | test_data_suite_fp | test_data_suite_bf16 +) def test_round_tosa_FP(test_data: torch.Tensor): pipeline = TosaPipelineFP[input_t1]( Round(), @@ -88,7 +99,9 @@ def test_round_u85_INT(test_data: torch.Tensor): pipeline.run() -@common.parametrize("test_data", test_data_suite | test_data_suite_bf16) +@common.parametrize( + "test_data", test_data_suite | test_data_suite_fp | test_data_suite_bf16 +) @common.SkipIfNoModelConverter def test_round_vgf_no_quant(test_data: torch.Tensor): pipeline = VgfPipeline[input_t1]( diff --git a/backends/arm/test/ops/test_sum.py b/backends/arm/test/ops/test_sum.py index 8b4cfe46075..ee9ecaed05f 100644 --- a/backends/arm/test/ops/test_sum.py +++ b/backends/arm/test/ops/test_sum.py @@ -8,6 +8,7 @@ import pytest import torch + from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( @@ -81,6 +82,16 @@ def test_sum_dim_intlist_scalar_input_tosa_FP_not_delegated(): pipeline.run() +def test_sum_bool_tosa_INT() -> None: + pipeline = TosaPipelineINT( + Sum(), + (torch.ones(1, dtype=torch.bool), [], False), + aten_op, + exir_op=[], + ) + pipeline.run() + + @common.parametrize( "test_data", Sum.test_parameters | Sum.test_parameters_bf16 | Sum.test_parameters_fp16, diff --git a/backends/arm/test/passes/test_arm_pass_manager_errors.py b/backends/arm/test/passes/test_arm_pass_manager_errors.py new file mode 100644 index 00000000000..22818af3b7d --- /dev/null +++ b/backends/arm/test/passes/test_arm_pass_manager_errors.py @@ -0,0 +1,37 @@ +# 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 pytest + +import torch + +from executorch.backends.arm._passes.arm_pass_manager import ( + _ExportedProgramGraphPassAdapter, +) +from executorch.exir.pass_manager import ExportedProgramPassManager, PassType +from torch.fx import GraphModule +from torch.fx.passes.infra.pass_base import PassResult + + +def test_exported_program_adapter_preserves_pass_names_in_errors() -> None: + """Report wrapped pass names instead of the adapter class name.""" + + def successful_pass(graph_module: GraphModule) -> PassResult: + return PassResult(graph_module, False) + + def failing_pass(graph_module: GraphModule) -> PassResult: + raise RuntimeError("test failure") + + exported_program = torch.export.export(torch.nn.ReLU(), (torch.randn(2, 3),)) + passes: list[PassType] = [ + _ExportedProgramGraphPassAdapter(successful_pass), + _ExportedProgramGraphPassAdapter(failing_pass), + ] + + with pytest.raises(Exception) as error: + ExportedProgramPassManager(passes)(exported_program) + + assert "running the 'failing_pass' pass" in str(error.value) + assert "following passes: ['successful_pass']" in str(error.value) diff --git a/backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.py b/backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.py index 38a55c8ba10..1ed9961d4f0 100644 --- a/backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.py +++ b/backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.py @@ -298,6 +298,38 @@ def test_canonicalize_moves_permute_before_view() -> None: _validate_numerics(gm_before, result.graph_module, (x_data,)) +def test_canonicalize_sinks_singleton_view_below_permute() -> None: + builder = GraphBuilder() + x_data = torch.randn(2, 8, 64) + x = builder.placeholder("x", x_data) + v1 = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, + args=(x, [2, 8, 1, 64]), + ) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(v1, [0, 2, 3, 1]), + ) + builder.output([p1]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + pass_instance = CanonicalizeViewCopyPermutePass() + result = cast(PassResult, pass_instance.call(original)) + + assert result.modified + compute_nodes = [ + node for node in result.graph_module.graph.nodes if node.op == "call_function" + ] + assert [node.target for node in compute_nodes] == [ + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.view_copy.default, + ] + assert compute_nodes[0].args[1] == [0, 2, 1] + assert compute_nodes[1].args[1] == [2, 1, 64, 8] + _validate_numerics(gm_before, result.graph_module, (x_data,)) + + def test_canonicalize_follows_interleaved_chain_users() -> None: builder = GraphBuilder() x_data = torch.randn(4, 2, 4) diff --git a/backends/arm/test/passes/test_cast_int_comparison_inputs_pass.py b/backends/arm/test/passes/test_cast_int_comparison_inputs_pass.py new file mode 100644 index 00000000000..ca3f8cb5cd9 --- /dev/null +++ b/backends/arm/test/passes/test_cast_int_comparison_inputs_pass.py @@ -0,0 +1,152 @@ +# 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 operator +from collections.abc import Callable + +import pytest +import torch +from executorch.backends.arm._passes import CastIntComparisonInputsPass +from executorch.backends.arm.test.tester.test_pipeline import ( + PassPipeline, + TosaPipelineFP, +) +from executorch.backends.test.harness.stages import StageType +from executorch.exir.dialects._ops import ops as edge_ops + + +class Comparison(torch.nn.Module): + def __init__(self, op: Callable[[torch.Tensor, torch.Tensor], torch.Tensor]): + super().__init__() + self.op = op + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return self.op(x, y) + + +class ScalarComparison(torch.nn.Module): + def __init__(self, op: Callable[[torch.Tensor, int], torch.Tensor]): + super().__init__() + self.op = op + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.op(x, 0) + + +comparison_ops = { + "eq": operator.eq, + "ne": operator.ne, + "ge": operator.ge, + "gt": operator.gt, + "le": operator.le, + "lt": operator.lt, +} +aten_ops = { + "eq": "torch.ops.aten.eq.Tensor", + "ne": "torch.ops.aten.ne.Tensor", + "ge": "torch.ops.aten.ge.Tensor", + "gt": "torch.ops.aten.gt.Tensor", + "le": "torch.ops.aten.le.Tensor", + "lt": "torch.ops.aten.lt.Tensor", +} +aten_scalar_ops = { + name: target.replace("Tensor", "Scalar") for name, target in aten_ops.items() +} +exir_ops = { + name: f"executorch_exir_dialects_edge__ops_aten_{name}_Tensor" + for name in comparison_ops +} +exir_scalar_ops = { + name: target.replace("Tensor", "Scalar") for name, target in exir_ops.items() +} + + +def comparison_inputs(dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: + limits = torch.iinfo(dtype) + return ( + torch.tensor( + [limits.min, limits.min + 1, limits.max - 1, limits.max], dtype=dtype + ), + torch.tensor( + [limits.min + 1, limits.min, limits.max, limits.max - 1], dtype=dtype + ), + ) + + +@pytest.mark.parametrize("op", comparison_ops.values(), ids=comparison_ops.keys()) +@pytest.mark.parametrize( + ("dtypes", "expected_dtype"), + ( + ((torch.int8, torch.int8), torch.float16), + ((torch.int16, torch.int16), torch.float32), + ((torch.int8, torch.int16), torch.float32), + ), +) +def test_cast_int_comparison_inputs(op, dtypes, expected_dtype) -> None: + inputs = ( + comparison_inputs(dtypes[0])[0], + comparison_inputs(dtypes[1])[1], + ) + pipeline = PassPipeline( + Comparison(op), + inputs, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2 + }, + pass_list=[CastIntComparisonInputsPass], + ) + pipeline.run() + + graph_module = ( + pipeline.tester.get_artifact(StageType.RUN_PASSES) + .exported_program() + .graph_module + ) + cast_op = edge_ops.edge.dim_order_ops._to_dim_order_copy.default + cast_nodes = [node for node in graph_module.graph.nodes if node.target == cast_op] + assert len(cast_nodes) == 2 + assert all(node.kwargs["dtype"] == expected_dtype for node in cast_nodes) + + +def test_cast_int_comparison_inputs_keeps_int32() -> None: + inputs = ( + torch.tensor([2**24, 2**24 + 1], dtype=torch.int32), + torch.tensor([2**24 + 1, 2**24], dtype=torch.int32), + ) + pipeline = PassPipeline( + Comparison(operator.eq), + inputs, + ops_not_after_pass=[ + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default" + ], + pass_list=[CastIntComparisonInputsPass], + ) + pipeline.run() + + +@pytest.mark.parametrize("name", comparison_ops) +@pytest.mark.parametrize("dtype", (torch.int8, torch.int16)) +def test_int_comparison_tosa_fp(name, dtype) -> None: + inputs = comparison_inputs(dtype) + pipeline = TosaPipelineFP( + Comparison(comparison_ops[name]), + inputs, + aten_ops[name], + exir_ops[name], + ) + pipeline.run() + + +@pytest.mark.parametrize("name", comparison_ops) +@pytest.mark.parametrize("dtype", (torch.int8, torch.int16)) +def test_int_scalar_comparison_tosa_fp(name, dtype) -> None: + inputs = (comparison_inputs(dtype)[0],) + pipeline = TosaPipelineFP( + ScalarComparison(comparison_ops[name]), + inputs, + aten_scalar_ops[name], + exir_scalar_ops[name], + ) + pipeline.run() diff --git a/backends/arm/test/passes/test_convert_bool_sum_pass.py b/backends/arm/test/passes/test_convert_bool_sum_pass.py new file mode 100644 index 00000000000..1b5842a9ed7 --- /dev/null +++ b/backends/arm/test/passes/test_convert_bool_sum_pass.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 pytest + +import torch + +from executorch.backends.arm._passes import ConvertBoolSumPass +from executorch.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) +from torch._export.utils import _get_shape_env_from_gm +from torch._subclasses import FakeTensorMode +from torch.fx import Graph, GraphModule + + +class BoolSum(torch.nn.Module): + def forward(self, x: torch.Tensor, dim: int, keepdim: bool): + return x.sum(dim=dim, keepdim=keepdim) + + +def _run_pass(graph_module: GraphModule, tosa_spec: str) -> GraphModule: + with TosaLoweringContext( + TosaSpecification.create_from_string(tosa_spec), + _get_shape_env_from_gm(graph_module), + ): + return ConvertBoolSumPass().call(graph_module).graph_module + + +def _operator_targets(graph_module: GraphModule) -> list[torch._ops.OpOverload]: + return [ + node.target for node in graph_module.graph.nodes if node.op == "call_function" + ] + + +def test_sum_bool_skips_inexact_accumulator() -> None: + graph = Graph() + with FakeTensorMode(): + fake_input = torch.empty(torch.iinfo(torch.int32).max + 1, dtype=torch.bool) + fake_output = torch.empty((), dtype=torch.int64) + x = graph.placeholder("x") + x.meta["val"] = fake_input + output = graph.call_function(torch.ops.aten.sum.dim_IntList, (x, [0], False)) + output.meta["val"] = fake_output + graph.output(output) + + result = _run_pass(GraphModule(torch.nn.Module(), graph), "TOSA-1.0+INT") + + assert _operator_targets(result) == [torch.ops.aten.sum.dim_IntList] + + +@pytest.mark.parametrize( + "max_dynamic_dim,expect_transform", + [ + (1024, True), + (torch.iinfo(torch.int32).max + 1, False), + ], +) +def test_sum_bool_dynamic_shape(max_dynamic_dim: int, expect_transform: bool) -> None: + dynamic_dim = torch.export.Dim("dynamic_dim", min=1, max=max_dynamic_dim) + exported_program = torch.export.export( + BoolSum(), + (torch.ones(2, dtype=torch.bool), 0, False), + dynamic_shapes=({0: dynamic_dim}, None, None), + ) + + result = _run_pass(exported_program.graph_module, "TOSA-1.0+INT") + + if not expect_transform: + assert _operator_targets(result) == [torch.ops.aten.sum.dim_IntList] + return + + sum_node = next( + node + for node in result.graph.nodes + if node.target == torch.ops.aten.sum.dim_IntList + ) + (output,) = result(torch.ones(5, dtype=torch.bool), 0, False) + assert sum_node.kwargs["dtype"] == torch.int32 + assert output.dtype == torch.int64 + assert output == 5 + + +@pytest.mark.parametrize("tosa_spec", ["TOSA-1.0+FP", "TOSA-1.0+FP+INT"]) +def test_sum_bool_skips_float_profiles(tosa_spec: str) -> None: + exported_program = torch.export.export( + BoolSum(), + (torch.ones(2, dtype=torch.bool), 0, False), + ) + + result = _run_pass(exported_program.graph_module, tosa_spec) + + assert _operator_targets(result) == [torch.ops.aten.sum.dim_IntList] diff --git a/backends/arm/test/passes/test_decompose_index_tensor_to_gather_pass.py b/backends/arm/test/passes/test_decompose_index_tensor_to_gather_pass.py new file mode 100644 index 00000000000..2d99bca9d70 --- /dev/null +++ b/backends/arm/test/passes/test_decompose_index_tensor_to_gather_pass.py @@ -0,0 +1,53 @@ +# 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 pytest +import torch +from executorch.backends.arm._passes.decompose_index_tensor_to_gather_pass import ( + DecomposeIndexTensorToGatherPass, +) +from executorch.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) +from executorch.exir import to_edge +from torch.export import export + + +class ConstantIndexTensor(torch.nn.Module): + def __init__(self, dim: int, index: int): + super().__init__() + self.dim = dim + self.index: torch.Tensor + self.register_buffer("index", torch.tensor([index], dtype=torch.int32)) + + def forward(self, x: torch.Tensor): + indices: list[slice | torch.Tensor] = [slice(None)] * x.dim() + indices[self.dim] = self.index + return x[tuple(indices)] + + +@pytest.mark.parametrize( + "dim,index", + ( + (0, 5), + (1, -6), + ), +) +def test_constant_out_of_bounds_index_raises(dim: int, index: int): + exported_program = export( + ConstantIndexTensor(dim, index), + (torch.rand(5, 5, 5),), + ) + edge_program = to_edge(exported_program) + edge_exported_program = edge_program.exported_program() + decompose_pass = DecomposeIndexTensorToGatherPass(edge_exported_program) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT")): + with pytest.raises( + IndexError, + match=rf"index {index} is out of bounds for dimension {dim} with size 5", + ): + decompose_pass(edge_exported_program.graph_module) diff --git a/backends/arm/test/passes/test_decompose_roll_pass.py b/backends/arm/test/passes/test_decompose_roll_pass.py new file mode 100644 index 00000000000..510a6ed8c09 --- /dev/null +++ b/backends/arm/test/passes/test_decompose_roll_pass.py @@ -0,0 +1,181 @@ +# 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 pytest +import torch + +from executorch.backends.arm._passes import DecomposeRollPass +from executorch.backends.arm._passes.decompose_roll_pass import can_decompose_roll +from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec +from executorch.backends.arm.util._factory import create_partitioner +from executorch.backends.arm.vgf.compile_spec import VgfCompileSpec +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops + + +class Roll(torch.nn.Module): + def __init__(self, shifts: tuple[int, ...], dims: tuple[int, ...]) -> None: + super().__init__() + self.shifts = shifts + self.dims = dims + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.roll(x, self.shifts, self.dims) + + +roll_cases = ( + ((2, 5), (1,), (1,), 1), + ((2, 3, 5), (-2,), (-1,), 1), + ((2, 5), (12,), (1,), 1), + ((1, 8, 8, 4), (-2, -2), (1, 2), 2), + ((1, 5, 3), (1, 2), (1, 1), 2), + ((5, 4), (5, 1), (0, 1), 1), +) + + +@pytest.mark.parametrize( + "shape,shifts,dims,expected_cats", + roll_cases, + ids=( + "positive", + "negative_shift_and_dim", + "oversized_shift", + "multiple_dims", + "repeated_dim", + "mixed_zero_shift", + ), +) +def test_decompose_roll( + shape: tuple[int, ...], + shifts: tuple[int, ...], + dims: tuple[int, ...], + expected_cats: int, +) -> None: + model = Roll(shifts, dims) + inputs = (torch.randn(shape),) + eager_output = model(*inputs) + edge = to_edge( + torch.export.export(model, inputs, strict=True), + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + preserve_ops=[torch.ops.aten.roll.default], + ), + ) + + edge = edge.transform([DecomposeRollPass()]) + graph = edge.exported_program().graph + targets = [node.target for node in graph.nodes if node.op == "call_function"] + + assert exir_ops.edge.aten.roll.default not in targets + assert targets.count(exir_ops.edge.aten.slice_copy.Tensor) == 2 * expected_cats + assert targets.count(exir_ops.edge.aten.cat.default) == expected_cats + assert torch.equal(edge.exported_program().module()(*inputs), eager_output) + + +@pytest.mark.parametrize( + "shape,shifts,dims", + ( + ((2, 4), (1,), ()), + ((2, 4), (4,), (1,)), + ((2, 0), (1,), (1,)), + ((2, 4), (1, 2), (1,)), + ((2, 4), (1,), (2,)), + ), + ids=("flat", "no_op", "zero_size", "mismatched_args", "invalid_dim"), +) +def test_roll_decomposition_guards( + shape: tuple[int, ...], shifts: tuple[int, ...], dims: tuple[int, ...] +) -> None: + assert not can_decompose_roll(shape, shifts, dims) + + +@pytest.mark.parametrize( + "dtype,compile_spec", + ( + (torch.float32, TosaCompileSpec("TOSA-1.0+FP")), + (torch.float16, TosaCompileSpec("TOSA-1.0+FP")), + (torch.bfloat16, TosaCompileSpec("TOSA-1.0+FP+bf16")), + (torch.float32, VgfCompileSpec()), + ), + ids=("tosa_fp32", "tosa_fp16", "tosa_bf16", "vgf_fp32"), +) +def test_partitioner_preserves_supported_roll( + dtype: torch.dtype, compile_spec: TosaCompileSpec | VgfCompileSpec +) -> None: + model = Roll((-2, -2), (1, 2)) + exported_program = torch.export.export( + model, (torch.randn(1, 8, 8, 4, dtype=dtype),), strict=True + ) + partitioner = create_partitioner(compile_spec) + preserved_ops, filter_fn = partitioner.ops_to_not_decompose(exported_program) + roll_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.roll.default + ) + + assert torch.ops.aten.roll.default in preserved_ops + assert filter_fn is not None and filter_fn(roll_node) + + +def test_partitioner_does_not_preserve_flat_roll() -> None: + class FlatRoll(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.roll(x, 2) + + exported_program = torch.export.export( + FlatRoll(), (torch.randn(2, 4),), strict=True + ) + partitioner = create_partitioner(VgfCompileSpec()) + _, filter_fn = partitioner.ops_to_not_decompose(exported_program) + roll_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.roll.default + ) + + assert filter_fn is not None and not filter_fn(roll_node) + + +def test_partitioner_does_not_preserve_unquantized_tosa_int_roll() -> None: + model = Roll((1,), (1,)) + exported_program = torch.export.export(model, (torch.randn(2, 4),), strict=True) + partitioner = create_partitioner(TosaCompileSpec("TOSA-1.0+INT")) + _, filter_fn = partitioner.ops_to_not_decompose(exported_program) + roll_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.roll.default + ) + + assert filter_fn is not None and not filter_fn(roll_node) + + +@pytest.mark.parametrize( + "dtype,compile_spec", + ( + (torch.float64, VgfCompileSpec()), + (torch.int32, VgfCompileSpec()), + (torch.bool, VgfCompileSpec()), + (torch.bfloat16, TosaCompileSpec("TOSA-1.0+FP")), + ), + ids=("float64", "int32", "bool", "bf16_without_extension"), +) +def test_partitioner_does_not_preserve_unsupported_dtype_roll( + dtype: torch.dtype, compile_spec: TosaCompileSpec | VgfCompileSpec +) -> None: + model = Roll((1,), (1,)) + exported_program = torch.export.export( + model, (torch.zeros(2, 4, dtype=dtype),), strict=True + ) + partitioner = create_partitioner(compile_spec) + _, filter_fn = partitioner.ops_to_not_decompose(exported_program) + roll_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.roll.default + ) + + assert filter_fn is not None and not filter_fn(roll_node) diff --git a/backends/arm/test/passes/test_deduplicate_const_shapes_pass.py b/backends/arm/test/passes/test_deduplicate_const_shapes_pass.py new file mode 100644 index 00000000000..d14dbc998e2 --- /dev/null +++ b/backends/arm/test/passes/test_deduplicate_const_shapes_pass.py @@ -0,0 +1,58 @@ +# 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 torch +from executorch.backends.arm._passes import DeduplicateConstShapesPass +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import Graph, GraphModule + + +def _const_shape(graph: Graph, name: str, values: list[int]): + node = graph.call_function( + exir_ops.backend.tosa.CONST_SHAPE.default, + (values,), + ) + node.name = name + node.meta["val"] = values + return node + + +def test_deduplicate_identical_const_shapes(): + graph = Graph() + first = _const_shape(graph, "first", [2, 3]) + second = _const_shape(graph, "second", [2, 3]) + graph.output((first, second)) + graph_module = GraphModule(torch.nn.Module(), graph) + + result = DeduplicateConstShapesPass()(graph_module) + + assert result is not None + assert result.modified + const_shapes = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.backend.tosa.CONST_SHAPE.default + ] + assert [node.name for node in const_shapes] == ["first"] + assert graph_module.graph.output_node().args[0] == (first, first) + + +def test_keep_const_shapes_with_different_values(): + graph = Graph() + first = _const_shape(graph, "first", [2, 3]) + second = _const_shape(graph, "second", [3, 2]) + graph.output((first, second)) + graph_module = GraphModule(torch.nn.Module(), graph) + + result = DeduplicateConstShapesPass()(graph_module) + + assert result is not None + assert not result.modified + const_shapes = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.backend.tosa.CONST_SHAPE.default + ] + assert const_shapes == [first, second] diff --git a/backends/arm/test/passes/test_fuse_batchnorm_pass.py b/backends/arm/test/passes/test_fuse_batchnorm_pass.py index 1c4d862d356..4b09cdaa647 100644 --- a/backends/arm/test/passes/test_fuse_batchnorm_pass.py +++ b/backends/arm/test/passes/test_fuse_batchnorm_pass.py @@ -104,6 +104,47 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x +class MergeConvTransposeBN(torch.nn.Module): + ops_before_pass: ClassVar[Dict[str, int]] = { + "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default": 1, + "executorch_exir_dialects_edge__ops_aten_convolution_default": 1, + } + ops_after_pass: ClassVar[Dict[str, int]] = { + "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default": 0, + "executorch_exir_dialects_edge__ops_aten_convolution_default": 1, + } + + def __init__( + self, + groups: int = 1, + bias: bool = False, + affine: bool = True, + in_channels: int = 4, + out_channels: int = 6, + ) -> None: + super().__init__() + self.conv_transpose2d = torch.nn.ConvTranspose2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2, + stride=2, + groups=groups, + bias=bias, + ) + self.batch_norm2d = torch.nn.BatchNorm2d(out_channels, affine=affine) + self.batch_norm2d.running_mean = torch.rand(out_channels) + self.batch_norm2d.running_var = torch.rand(out_channels) + if affine: + self.batch_norm2d.weight = torch.nn.Parameter(torch.rand(out_channels)) + self.batch_norm2d.bias = torch.nn.Parameter(torch.rand(out_channels)) + + def get_inputs(self) -> input_t: + return (torch.randn(1, self.conv_transpose2d.in_channels, 8, 8),) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.batch_norm2d(self.conv_transpose2d(x)) + + class MergeMultipleUsersBN(torch.nn.Module): ops_before_pass: ClassVar[Dict[str, int]] = { "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default": 2, @@ -154,6 +195,20 @@ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: "merge_two_of_two_bn_affine": cast( ModuleWithBatchNormAttrs, MergeTwosOfTwoBN(True) ), + "merge_conv_transpose_bn": cast(ModuleWithBatchNormAttrs, MergeConvTransposeBN()), + "merge_grouped_conv_transpose_bn": cast( + ModuleWithBatchNormAttrs, MergeConvTransposeBN(groups=2) + ), + "merge_grouped_conv_transpose_bn_bias": cast( + ModuleWithBatchNormAttrs, MergeConvTransposeBN(groups=2, bias=True) + ), + "merge_grouped_conv_transpose_bn_no_affine": cast( + ModuleWithBatchNormAttrs, MergeConvTransposeBN(groups=2, affine=False) + ), + "merge_grouped_conv_transpose_bn_equal_channels": cast( + ModuleWithBatchNormAttrs, + MergeConvTransposeBN(groups=2, in_channels=4, out_channels=4), + ), "merge_multiple_users_bn_affine": cast( ModuleWithBatchNormAttrs, MergeMultipleUsersBN(True) ), diff --git a/backends/arm/test/passes/test_fuse_duplicate_users_pass.py b/backends/arm/test/passes/test_fuse_duplicate_users_pass.py index 893d9eefea5..aad803b98e7 100644 --- a/backends/arm/test/passes/test_fuse_duplicate_users_pass.py +++ b/backends/arm/test/passes/test_fuse_duplicate_users_pass.py @@ -8,6 +8,7 @@ import executorch.backends.arm.tosa.dialect # noqa: F401 import torch from executorch.backends.arm._passes import ( + DeduplicateConstShapesPass, EnsureUniqueOutputNodesPass, FuseDuplicateUsersPass, InsertRescalePass, @@ -21,6 +22,9 @@ TosaLoweringContext, TosaSpecification, ) +from executorch.backends.transforms.fuse_duplicate_users_pass import ( + DO_NOT_FUSE_DUPLICATE_META_KEY, +) from executorch.exir import EdgeCompileConfig, to_edge from executorch.exir.dialects._ops import ops as exir_ops from torch.export import export @@ -163,6 +167,22 @@ def test_fuse_duplicate_users_preserves_graph_order_for_representative(): assert len(_add_node_names(result.graph_module)) == 1 +def test_fuse_duplicate_users_honors_do_not_fuse_marker(): + graph_module = _graph_with_users_not_in_node_order() + marked_node = next( + node + for node in graph_module.graph.nodes + if node.target == torch.ops.aten.add.Tensor + ) + marked_node.meta[DO_NOT_FUSE_DUPLICATE_META_KEY] = True + + result = FuseDuplicateUsersPass()(graph_module) + + result.graph_module.graph.lint() + assert not result.modified + assert len(_add_node_names(result.graph_module)) == 2 + + def test_fuse_duplicate_users_keeps_identical_rescale_users(): graph_module = _graph_with_duplicate_rescale_users() @@ -222,8 +242,9 @@ def test_fuse_duplicate_users_runs_after_tosa_transformations(): for index, pass_type in enumerate(pass_types) if pass_type is RemoveNoopPass ) - assert pass_types[post_noop_index + 1 : post_noop_index + 4] == [ + assert pass_types[post_noop_index + 1 : post_noop_index + 5] == [ FuseDuplicateUsersPass, InsertRescalePass, + DeduplicateConstShapesPass, EnsureUniqueOutputNodesPass, ] diff --git a/backends/arm/test/passes/test_insert_rescale_i32_pass.py b/backends/arm/test/passes/test_insert_rescale_i32_pass.py index e685da11e05..b163db11887 100644 --- a/backends/arm/test/passes/test_insert_rescale_i32_pass.py +++ b/backends/arm/test/passes/test_insert_rescale_i32_pass.py @@ -1,16 +1,19 @@ -# 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. -from typing import Tuple +from typing import Callable, Tuple +import pytest import torch from executorch.backends.arm._passes import ( FoldAndAnnotateQParamsPass, InsertRescaleInt32Pass, ) +from executorch.backends.arm.common.annotation_meta import ArmAnnotationInfo from executorch.backends.arm.test.tester.test_pipeline import PassPipeline +from executorch.exir.dialects._ops import ops as exir_ops class MultipleOpsModel(torch.nn.Module): @@ -90,6 +93,55 @@ def test_insert_rescale_int32_tosa_INT_multiple_ops(): _test_model_with_f32_data(MultipleOpsModel()) +@pytest.mark.parametrize( + "binary_target", + (exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.mul.Tensor), +) +def test_insert_rescale_int32_preserves_partial_qdq( + binary_target: Callable[..., object], +) -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + y = graph.placeholder("y") + x_q = graph.call_function( + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + (x, 0.5, 0, -128, 127, torch.int8), + ) + x_dq = graph.call_function( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + (x_q, 0.5, 0, -128, 127, torch.int8), + ) + binary = graph.call_function(binary_target, (x_dq, y)) + binary.meta["custom"] = { + ArmAnnotationInfo.CUSTOM_META_KEY: ArmAnnotationInfo(quantized=True) + } + output_q = graph.call_function( + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + (binary, 0.5, 0, -128, 127, torch.int8), + ) + output_dq = graph.call_function( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + (output_q, 0.5, 0, -128, 127, torch.int8), + ) + graph.output(output_dq) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + fold_result = FoldAndAnnotateQParamsPass(preserve_partial_binary_tensor_qdq=True)( + graph_module + ) + assert fold_result is not None + rescale_result = InsertRescaleInt32Pass()(fold_result.graph_module) + assert rescale_result is not None + + assert binary.args == (x_dq, y) + assert output_q in binary.users + assert not rescale_result.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.backend.tosa.RESCALE.default, + sort=False, + ) + + def test_insert_rescale_int32_tosa_FP_dont_insert_rescales(): module = MultipleOpsModel() input_t = Tuple[torch.Tensor, torch.Tensor] diff --git a/backends/arm/test/passes/test_normalize_max_pool2d_input_rank_pass.py b/backends/arm/test/passes/test_normalize_max_pool2d_input_rank_pass.py new file mode 100644 index 00000000000..11aee89778c --- /dev/null +++ b/backends/arm/test/passes/test_normalize_max_pool2d_input_rank_pass.py @@ -0,0 +1,87 @@ +# 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._passes import ( + NormalizeMaxPool2dInputRankPass, + RemoveGetItemPass, +) +from executorch.backends.arm.test.tester.test_pipeline import PassPipeline +from executorch.backends.test.harness.stages import StageType +from executorch.exir.dialects._ops import ops as exir_ops + + +input_t = Tuple[torch.Tensor] + + +class MaxPool2d(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.max_pool2d( + x, + kernel_size=(3, 2), + stride=(2, 1), + padding=(1, 0), + dilation=(1, 1), + ceil_mode=True, + ) + + +def test_normalize_rank3_max_pool2d_input() -> None: + pipeline = PassPipeline[input_t]( + MaxPool2d(), + (torch.rand(3, 9, 11),), + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_max_pool2d_with_indices_default": 1, + }, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_aten_unsqueeze_copy_default": 1, + "executorch_exir_dialects_edge__ops_aten_max_pool2d_default": 1, + "executorch_exir_dialects_edge__ops_aten_squeeze_copy_dims": 1, + }, + pass_list=[RemoveGetItemPass, NormalizeMaxPool2dInputRankPass], + ) + pipeline.run() + + exported_program = pipeline.tester.get_artifact( + StageType.RUN_PASSES + ).exported_program() + pool_node = next( + node + for node in exported_program.graph.nodes + if node.target == exir_ops.edge.aten.max_pool2d.default + ) + unsqueeze_node = pool_node.args[0] + assert isinstance(unsqueeze_node, torch.fx.Node) + assert unsqueeze_node.target == exir_ops.edge.aten.unsqueeze_copy.default + assert unsqueeze_node.args[1] == 0 + assert tuple(pool_node.args[1]) == (3, 2) + assert tuple(pool_node.args[2]) == (2, 1) + assert tuple(pool_node.args[3]) == (1, 0) + assert tuple(pool_node.args[4]) == (1, 1) + assert pool_node.args[5] is True + + squeeze_node = next(iter(pool_node.users)) + assert squeeze_node.target == exir_ops.edge.aten.squeeze_copy.dims + assert squeeze_node.args == (pool_node, [0]) + + +def test_normalize_rank4_max_pool2d_input_is_noop() -> None: + PassPipeline[input_t]( + MaxPool2d(), + (torch.rand(1, 3, 9, 11),), + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_max_pool2d_with_indices_default": 1, + }, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_aten_max_pool2d_default": 1, + }, + ops_not_after_pass=[ + "executorch_exir_dialects_edge__ops_aten_unsqueeze_copy_default", + "executorch_exir_dialects_edge__ops_aten_squeeze_copy_dims", + ], + pass_list=[RemoveGetItemPass, NormalizeMaxPool2dInputRankPass], + ).run() diff --git a/backends/arm/test/passes/test_remove_data_layout_noops.py b/backends/arm/test/passes/test_remove_data_layout_noops.py index 87e592dcdea..7090cfbf008 100644 --- a/backends/arm/test/passes/test_remove_data_layout_noops.py +++ b/backends/arm/test/passes/test_remove_data_layout_noops.py @@ -8,6 +8,7 @@ import torch from executorch.backends.arm._passes import ( CanonicalizeViewCopyPermutePass, + DeduplicateConstShapesPass, EnsureUniqueOutputNodesPass, ExirToTosaPass, FuseDuplicateUsersPass, @@ -499,4 +500,5 @@ def test_data_layout_noop_cleanup_pipeline_order(): assert pass_types[pre_tosa_cleanup + 1] is CanonicalizeViewCopyPermutePass assert pass_types[post_tosa_cleanup + 1] is FuseDuplicateUsersPass assert pass_types[post_tosa_cleanup + 2] is InsertRescalePass - assert pass_types[post_tosa_cleanup + 3] is EnsureUniqueOutputNodesPass + assert pass_types[post_tosa_cleanup + 3] is DeduplicateConstShapesPass + assert pass_types[post_tosa_cleanup + 4] is EnsureUniqueOutputNodesPass diff --git a/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py b/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py index 864b6c669f9..c9ceccdef89 100644 --- a/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py +++ b/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py @@ -26,6 +26,8 @@ RESCALE_TARGET = exir_ops.backend.tosa.RESCALE.default MUL_TARGET = exir_ops.edge.aten.mul.Tensor ADD_TARGET = exir_ops.edge.aten.add.Tensor +SUB_TARGET = exir_ops.edge.aten.sub.Tensor +VIEW_TARGET = exir_ops.edge.aten.view_copy.default ERF_TARGET = exir_ops.edge.aten.erf.default @@ -150,6 +152,149 @@ def test_remove_permutes_around_rescale_tosa_INT() -> None: assert _count_nodes(result.graph_module, RESCALE_TARGET) == 1 +def test_terminal_sink_view_broadcast_is_optimized_tosa_INT() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.randn(1, 4, 1, 1) + direct = graph.placeholder("direct") + direct.meta["val"] = torch.randn(1, 8, 4) + + permute = graph.create_node("call_function", PERMUTE_TARGET, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.randn(1, 1, 1, 4) + mul = graph.create_node("call_function", MUL_TARGET, args=(permute, permute)) + mul.meta["val"] = torch.randn(1, 1, 1, 4) + sink = graph.create_node("call_function", VIEW_TARGET, args=(mul, [1, 1, 4])) + sink.meta["val"] = torch.randn(1, 1, 4) + rescale = graph.create_node( + "call_function", + RESCALE_TARGET, + args=(sink, torch.int8, [1.0], 0, 0), + ) + rescale.meta["val"] = torch.randn(1, 1, 4) + sub = graph.create_node("call_function", SUB_TARGET, args=(direct, rescale)) + sub.meta["val"] = torch.randn(1, 8, 4) + graph.output(sub) + + graph_module = torch.fx.GraphModule({}, graph) + with TosaLoweringContext(TOSA_INT_SPEC): + result = RemovePermutesAroundElementwiseTosaOps(_fake_exported_program()).call( + graph_module + ) + + assert result.modified + assert _count_nodes(result.graph_module, PERMUTE_TARGET) == 0 + assert sub.args == (direct, rescale) + + +def test_remove_permutes_around_singleton_view_and_gate_region_tosa_INT() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.randn(1, 8, 4) + gate_source = graph.placeholder("gate_source") + gate_source.meta["val"] = torch.randn(1, 1, 1, 4) + skip = graph.placeholder("skip") + skip.meta["val"] = torch.randn(1, 8, 4) + + layout_in = graph.create_node("call_function", PERMUTE_TARGET, args=(x, [0, 2, 1])) + layout_in.meta["val"] = torch.randn(1, 4, 8) + clamp = graph.create_node( + "call_function", exir_ops.edge.aten.clamp.default, args=(layout_in, 0, None) + ) + clamp.meta["val"] = torch.randn(1, 4, 8) + + pool_view = graph.create_node( + "call_function", VIEW_TARGET, args=(clamp, [1, 4, 8, 1]) + ) + pool_view.meta["val"] = torch.randn(1, 4, 8, 1) + pool_layout = graph.create_node( + "call_function", PERMUTE_TARGET, args=(pool_view, [0, 2, 3, 1]) + ) + pool_layout.meta["val"] = torch.randn(1, 8, 1, 4) + + gate = graph.create_node( + "call_function", VIEW_TARGET, args=(gate_source, [1, 4, 1]) + ) + gate.meta["val"] = torch.randn(1, 4, 1) + gated = graph.create_node("call_function", MUL_TARGET, args=(clamp, gate)) + gated.meta["val"] = torch.randn(1, 4, 8) + + skip_layout = graph.create_node( + "call_function", PERMUTE_TARGET, args=(skip, [0, 2, 1]) + ) + skip_layout.meta["val"] = torch.randn(1, 4, 8) + residual = graph.create_node("call_function", ADD_TARGET, args=(gated, skip_layout)) + residual.meta["val"] = torch.randn(1, 4, 8) + layout_out = graph.create_node( + "call_function", PERMUTE_TARGET, args=(residual, [0, 2, 1]) + ) + layout_out.meta["val"] = torch.randn(1, 8, 4) + graph.output((pool_layout, layout_out)) + + graph_module = torch.fx.GraphModule({}, graph) + inputs = ( + torch.randn(1, 8, 4), + torch.randn(1, 1, 1, 4), + torch.randn(1, 8, 4), + ) + expected = graph_module(*inputs) + + with TosaLoweringContext(TOSA_INT_SPEC): + result = RemovePermutesAroundElementwiseTosaOps(_fake_exported_program()).call( + graph_module + ) + + assert result.modified + assert _count_nodes(result.graph_module, PERMUTE_TARGET) == 0 + assert pool_view.args[1] == [1, 8, 1, 4] + assert gate.args[1] == [1, 1, 4] + actual = result.graph_module(*inputs) + torch.testing.assert_close(actual, expected) + + +def test_remove_permutes_around_amax_reduction_region_tosa_INT() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.randn(1, 8, 4) + + layout_in = graph.create_node("call_function", PERMUTE_TARGET, args=(x, [0, 2, 1])) + layout_in.meta["val"] = torch.randn(1, 4, 8) + clamp = graph.create_node( + "call_function", exir_ops.edge.aten.clamp.default, args=(layout_in, 0, None) + ) + clamp.meta["val"] = torch.randn(1, 4, 8) + + conv_layout = graph.create_node( + "call_function", PERMUTE_TARGET, args=(clamp, [0, 2, 1]) + ) + conv_layout.meta["val"] = torch.randn(1, 8, 4) + maximum = graph.create_node( + "call_function", exir_ops.edge.aten.amax.default, args=(clamp, 2, True) + ) + maximum.meta["val"] = torch.randn(1, 4, 1) + centered = graph.create_node("call_function", SUB_TARGET, args=(clamp, maximum)) + centered.meta["val"] = torch.randn(1, 4, 8) + stats_layout = graph.create_node( + "call_function", PERMUTE_TARGET, args=(centered, [0, 2, 1]) + ) + stats_layout.meta["val"] = torch.randn(1, 8, 4) + graph.output((conv_layout, stats_layout)) + + graph_module = torch.fx.GraphModule({}, graph) + inputs = (torch.randn(1, 8, 4),) + expected = graph_module(*inputs) + + with TosaLoweringContext(TOSA_INT_SPEC): + result = RemovePermutesAroundElementwiseTosaOps(_fake_exported_program()).call( + graph_module + ) + + assert result.modified + assert _count_nodes(result.graph_module, PERMUTE_TARGET) == 0 + assert maximum.args[1] == 1 + actual = result.graph_module(*inputs) + torch.testing.assert_close(actual, expected) + + def test_remove_permutes_around_gelu_with_folded_scalar_constants_tosa_FP() -> None: graph = torch.fx.Graph() x = graph.placeholder("x") diff --git a/backends/arm/test/passes/test_remove_safe_softmax_guard_pass.py b/backends/arm/test/passes/test_remove_safe_softmax_guard_pass.py index 772b5d39eb8..8064607de59 100644 --- a/backends/arm/test/passes/test_remove_safe_softmax_guard_pass.py +++ b/backends/arm/test/passes/test_remove_safe_softmax_guard_pass.py @@ -18,20 +18,24 @@ _get_tosa_operator_distribution, ArmTester, ) +from executorch.backends.arm.tosa.partitioner import TOSAPartitioner from executorch.backends.test.harness.stages import StageType -from executorch.exir import to_edge +from executorch.exir import to_edge, to_edge_transform_and_lower from executorch.exir.dialects._ops import ops as exir_ops from torch.export import export from torch.fx import GraphModule, Node class SDPA(torch.nn.Module): - def __init__(self, attn_mask: torch.Tensor | None = None) -> None: + def __init__( + self, attn_mask: torch.Tensor | None = None, is_causal: bool = False + ) -> None: super().__init__() if attn_mask is not None: self.register_buffer("attn_mask", attn_mask) else: self.attn_mask = None + self.is_causal = is_causal def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor @@ -41,6 +45,35 @@ def forward( key, value, attn_mask=self.attn_mask, + is_causal=self.is_causal, + ) + + +class DynamicMaskSDPA(torch.nn.Module): + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor, + ) -> torch.Tensor: + return torch.nn.functional.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attn_mask, + ) + + +class DropoutSDPA(torch.nn.Module): + def forward( + self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor + ) -> torch.Tensor: + return torch.nn.functional.scaled_dot_product_attention( + query, + key, + value, + dropout_p=0.1, ) @@ -217,14 +250,35 @@ def test_sdpa_safe_softmax_guard_preserve_keeps_guard_before_tosa_lowering(): assert counts["SELECT"] == 1 -def test_sdpa_safe_softmax_guard_remove_when_proven_keeps_guard(): +def test_auto_removal_runs_through_partitioner_hook(): compile_spec = common.get_tosa_compile_spec("TOSA-1.0+FP") compile_spec.set_pass_pipeline_config( - ArmPassPipelineConfig( - sdpa_safe_softmax_guard=(SDPASafeSoftmaxGuardPolicy.REMOVE_WHEN_PROVEN) - ) + ArmPassPipelineConfig(sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO) ) - tester = ArmTester(SDPA(), _sdpa_inputs(), compile_spec) + edge_program = to_edge_transform_and_lower( + export(SDPA(), _sdpa_inputs(), strict=True), + partitioner=[TOSAPartitioner(compile_spec)], + ) + graph_module = edge_program.exported_program().graph_module + counts = dict(_get_tosa_operator_distribution(graph_module)) + + assert counts.get("EQUAL", 0) == 0 + assert counts.get("LOGICAL_NOT", 0) == 0 + assert counts.get("REDUCE_ANY", 0) == 0 + assert counts.get("SELECT", 0) == 0 + assert counts["REDUCE_MAX"] == 1 + assert counts["EXP"] == 1 + assert counts["REDUCE_SUM"] == 1 + assert counts["RECIPROCAL"] == 1 + + +def test_sdpa_safe_softmax_guard_auto_keeps_dynamic_mask_guard(): + compile_spec = common.get_tosa_compile_spec("TOSA-1.0+FP") + compile_spec.set_pass_pipeline_config( + ArmPassPipelineConfig(sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO) + ) + attn_mask = torch.zeros(1, 1, 4, 4) + tester = ArmTester(DynamicMaskSDPA(), (*_sdpa_inputs(), attn_mask), compile_spec) tester.export().to_edge_transform_and_lower() graph_module = ( @@ -237,7 +291,49 @@ def test_sdpa_safe_softmax_guard_remove_when_proven_keeps_guard(): assert counts["EQUAL"] == 1 assert counts["LOGICAL_NOT"] == 2 assert counts["REDUCE_ANY"] == 1 - assert counts["SELECT"] == 1 + assert counts["SELECT"] >= 1 + + +def test_sdpa_safe_softmax_guard_auto_keeps_causal_guard(): + compile_spec = common.get_tosa_compile_spec("TOSA-1.0+FP") + compile_spec.set_pass_pipeline_config( + ArmPassPipelineConfig(sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO) + ) + tester = ArmTester(SDPA(is_causal=True), _sdpa_inputs(), compile_spec) + + tester.export().to_edge_transform_and_lower() + graph_module = ( + tester.get_artifact(StageType.TO_EDGE_TRANSFORM_AND_LOWER) + .exported_program() + .graph_module + ) + counts = dict(_get_tosa_operator_distribution(graph_module)) + + assert counts["EQUAL"] == 1 + assert counts["LOGICAL_NOT"] == 2 + assert counts["REDUCE_ANY"] == 1 + assert counts["SELECT"] >= 1 + + +def test_sdpa_safe_softmax_guard_auto_keeps_dropout_guard(): + compile_spec = common.get_tosa_compile_spec("TOSA-1.0+FP") + compile_spec.set_pass_pipeline_config( + ArmPassPipelineConfig(sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO) + ) + tester = ArmTester(DropoutSDPA(), _sdpa_inputs(), compile_spec) + + tester.export().to_edge_transform_and_lower() + graph_module = ( + tester.get_artifact(StageType.TO_EDGE_TRANSFORM_AND_LOWER) + .exported_program() + .graph_module + ) + counts = dict(_get_tosa_operator_distribution(graph_module)) + + assert counts["EQUAL"] == 1 + assert counts["LOGICAL_NOT"] == 2 + assert counts["REDUCE_ANY"] == 1 + assert counts["SELECT"] >= 1 def test_remove_safe_softmax_guard_pass_does_not_rewrite_regular_softmax(): diff --git a/backends/arm/test/passes/test_rewrite_conv_pass.py b/backends/arm/test/passes/test_rewrite_conv_pass.py index 31c4205f16f..3928aa81e75 100644 --- a/backends/arm/test/passes/test_rewrite_conv_pass.py +++ b/backends/arm/test/passes/test_rewrite_conv_pass.py @@ -27,7 +27,10 @@ from executorch.backends.arm.test.misc.test_dw_convs_with_shared_weights import ( DWConvsModule, ) -from executorch.backends.arm.test.tester.test_pipeline import PassPipeline +from executorch.backends.arm.test.tester.test_pipeline import ( + EthosU55PipelineINT, + PassPipeline, +) from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec from executorch.backends.arm.tosa.mapping import TosaSpecialDtype from executorch.backends.arm.tosa.partitioner import TOSAPartitioner @@ -36,6 +39,9 @@ TosaSpecification, ) from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner +from executorch.backends.transforms.fuse_duplicate_users_pass import ( + build_node_signature, +) from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower from executorch.exir.dialects._ops import ops as exir_ops from torch.export import Dim, export @@ -90,6 +96,37 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.conv(x) + x +class A16W8Conv1dInt32Consumer(nn.Module): + """Exercise a rank-three A16W8 convolution consumed only by an INT32 add.""" + + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv1d(4, 4, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Feed the convolution output directly to a residual addition.""" + return self.conv(x) + x + + +class A16W8Conv1dSharedConsumers(nn.Module): + """Exercise a rank-three A16W8 convolution read by two INT32 consumers. + + Mirrors the attention tail of an ECAPA-style model, where ``Softmax(dim=2)`` + decomposes into an ``amax`` reduction and a subtraction that both read the + convolution output. + + """ + + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv1d(4, 4, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Share the convolution output between reduction and subtraction.""" + y = self.conv(x) + return y - y.amax(dim=2, keepdim=True) + + class A16W8MixedConsumer(nn.Module): """Exercise a shared A16W8 convolution output with mixed consumers.""" @@ -236,7 +273,9 @@ def _get_expected_int32_scales( def _rewrite_a16w8_convs( - model: nn.Module, inputs: tuple[torch.Tensor, ...] + model: nn.Module, + inputs: tuple[torch.Tensor, ...], + tosa_spec: TosaSpecification | None = None, ) -> tuple[torch.fx.GraphModule, list[list[float]]]: """Run the passes needed to inspect rewritten A16W8 convolutions.""" exported_program = _export_quantized_a16w8(model, inputs) @@ -245,7 +284,7 @@ def _rewrite_a16w8_convs( ).exported_program() gm = _run_pre_rewrite_passes(edge_program) rewrite_pass = RewriteConvPass(edge_program) - with TosaLoweringContext(_compile_spec_int16().tosa_spec): + with TosaLoweringContext(tosa_spec or _compile_spec_int16().tosa_spec): rescale_result = InsertRescaleInt32Pass()(gm) assert rescale_result is not None expected_int32_scales = _get_expected_int32_scales( @@ -263,6 +302,29 @@ def _get_call_function_node(gm: torch.fx.GraphModule, target): raise AssertionError(f"Node with target {target} not found") +def _add_a16w8_rescale_head( + graph: torch.fx.Graph, + accumulator: torch.fx.Node, + positional_unsigned: tuple[bool, ...] = (), +) -> tuple[torch.fx.Node, torch.fx.Node]: + rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + args=( + accumulator, + torch.int16, + [1.0], + 0, + 0, + *positional_unsigned, + ), + ) + layout_permute = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + args=(rescale, [0, 3, 1, 2]), + ) + return rescale, layout_permute + + class ConvModule(torch.nn.Module): def __init__(self): super().__init__() @@ -296,6 +358,27 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.conv(x) +class Conv1dBiasModule(torch.nn.Module): + def __init__(self, depthwise: bool = False) -> None: + super().__init__() + groups = 4 if depthwise else 1 + out_channels = 8 if depthwise else 6 + self.conv = torch.nn.Conv1d( + 4, + out_channels, + kernel_size=3, + padding=1, + groups=groups, + bias=True, + ) + + def get_inputs(self) -> tuple[torch.Tensor]: + return (torch.randn(1, 4, 8),) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + class Conv3dBiasModule(torch.nn.Module): def __init__(self) -> None: super().__init__() @@ -458,6 +541,152 @@ def test_rewrite_conv_a16w8_mixed_consumers_restore_int16( ) +def test_rewrite_conv_rescale_signature_includes_positional_unsigned_flags() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + signed_rescale, _ = _add_a16w8_rescale_head(graph, accumulator, (False, False)) + unsigned_rescale, _ = _add_a16w8_rescale_head(graph, accumulator, (False, True)) + + assert build_node_signature( + signed_rescale, positional_arg_start=1 + ) != build_node_signature(unsigned_rescale, positional_arg_start=1) + + +def test_rewrite_conv_without_convolution_does_not_require_context() -> None: + inputs = (torch.randn(1, 4),) + edge_program = to_edge(export(nn.Identity(), inputs)).exported_program() + + result = RewriteConvPass(edge_program)(edge_program.graph_module) + + assert result is not None + assert not result.modified + + +def test_rewrite_conv_a16w8_unknown_accumulator_user_is_unchanged() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + _, first_permute = _add_a16w8_rescale_head(graph, accumulator) + _, second_permute = _add_a16w8_rescale_head(graph, accumulator) + unexpected_user = graph.call_function(torch.neg, args=(accumulator,)) + graph.output((first_permute, second_permute, unexpected_user)) + graph_module = torch.fx.GraphModule({}, graph) + nodes_before = list(graph.nodes) + users_before = {node: tuple(node.users) for node in graph.nodes} + node_order = {node: index for index, node in enumerate(graph.nodes)} + + result = RewriteConvPass._deduplicate_a16w8_output_rescales( + graph_module, accumulator, node_order + ) + + assert result is None + assert list(graph.nodes) == nodes_before + assert {node: tuple(node.users) for node in graph.nodes} == users_before + graph.lint() + + +def test_rewrite_conv_a16w8_multi_consumer_rescale_is_not_deduplicated() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + first_rescale, first_permute = _add_a16w8_rescale_head(graph, accumulator) + second_rescale, second_permute = _add_a16w8_rescale_head(graph, accumulator) + output = graph.output((first_permute, second_permute, second_rescale)) + graph_module = torch.fx.GraphModule({}, graph) + node_order = {node: index for index, node in enumerate(graph.nodes)} + + result = RewriteConvPass._deduplicate_a16w8_output_rescales( + graph_module, accumulator, node_order + ) + + assert result == [first_rescale, second_rescale] + assert set(second_rescale.users) == {second_permute, output} + graph.lint() + + +def test_rewrite_conv_a16w8_deduplication_uses_graph_order() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + temporary_input = graph.placeholder("temporary_input") + early_rescale, early_permute = _add_a16w8_rescale_head(graph, temporary_input) + late_rescale, late_permute = _add_a16w8_rescale_head(graph, accumulator) + early_rescale.replace_input_with(temporary_input, accumulator) + graph.output((early_permute, late_permute)) + graph_module = torch.fx.GraphModule({}, graph) + node_order = {node: index for index, node in enumerate(graph.nodes)} + + assert list(accumulator.users) == [late_rescale, early_rescale] + assert node_order[early_rescale] < node_order[late_rescale] + + result = RewriteConvPass._deduplicate_a16w8_output_rescales( + graph_module, accumulator, node_order + ) + + assert result == [early_rescale] + assert late_rescale not in graph.nodes + assert late_permute not in graph.nodes + graph.lint() + + +def test_rewrite_conv_a16w8_unknown_order_user_is_unchanged() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + _, first_permute = _add_a16w8_rescale_head(graph, accumulator) + node_order = {node: index for index, node in enumerate(graph.nodes)} + _, second_permute = _add_a16w8_rescale_head(graph, accumulator) + graph.output((first_permute, second_permute)) + graph_module = torch.fx.GraphModule({}, graph) + nodes_before = list(graph.nodes) + users_before = {node: tuple(node.users) for node in graph.nodes} + + result = RewriteConvPass._deduplicate_a16w8_output_rescales( + graph_module, accumulator, node_order + ) + + assert result is None + assert list(graph.nodes) == nodes_before + assert {node: tuple(node.users) for node in graph.nodes} == users_before + graph.lint() + + +def test_rewrite_conv_a16w8_u55_separates_distinct_output_rescales() -> None: + model = A16W8MixedConsumerChain(nn.Conv2d(4, 4, 1)) + inputs = (torch.randn(1, 4, 8, 8),) + generic_graph, _ = _rewrite_a16w8_convs(model, inputs) + u55_graph, _ = _rewrite_a16w8_convs( + model, + inputs, + TosaSpecification.create_from_string("TOSA-1.0+INT+int16+int4+u55"), + ) + + conv_targets = { + exir_ops.backend.tosa.CONV2D.default, + exir_ops.backend.tosa.DEPTHWISE_CONV2D.default, + } + generic_convs = [ + node for node in generic_graph.graph.nodes if node.target in conv_targets + ] + u55_convs = [node for node in u55_graph.graph.nodes if node.target in conv_targets] + + assert len(u55_convs) == len(generic_convs) + 1 + assert all(len(conv.users) == 1 for conv in u55_convs) + assert all( + next(iter(conv.users)).target == exir_ops.backend.tosa.RESCALE.default + for conv in u55_convs + ) + + +def test_rewrite_conv_a16w8_mixed_consumers_lowers_on_u55() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + pipeline = EthosU55PipelineINT[tuple[torch.Tensor]]( + A16W8MixedConsumerChain(nn.Conv2d(4, 4, 1)), + inputs, + aten_ops=[], + exir_ops=[], + run_on_fvp=False, + a16w8_quantization=True, + ) + pipeline.run() + + def test_rewrite_conv_a16w8_preserves_int32_for_int32_consumers() -> None: r"""Test that an exclusively INT32 consumer keeps the widened path. @@ -491,6 +720,69 @@ def test_rewrite_conv_a16w8_preserves_int32_for_int32_consumers() -> None: assert direct_int32_rescales[0].args[2] == pytest.approx(expected_int32_scales[0]) +def test_rewrite_conv1d_a16w8_narrows_instead_of_forking_int32() -> None: + """Test that a rank-three A16W8 convolution narrows instead of forking. + + Each widened INT32 branch carries its own boundary rescale and layout + permutation. Vela materialises the rank-three permutation as a full + transpose of the convolution output, so a second branch doubles that cost. A + rank-three convolution therefore narrows to its exported INT16 domain and + keeps a single layout boundary. + + """ + inputs = (torch.randn(1, 4, 8),) + gm, _ = _rewrite_a16w8_convs(A16W8Conv1dInt32Consumer(), inputs) + + conv = _get_call_function_node(gm, exir_ops.backend.tosa.CONV2D.default) + forked_int32_rescales = [ + node + for node in gm.graph.nodes + if node.op == "call_function" + and node.target == exir_ops.backend.tosa.RESCALE.default + and node.args[1] == torch.int32 + and node.all_input_nodes[0] is conv + ] + assert forked_int32_rescales == [] + + (boundary_rescale,) = tuple(conv.users) + assert boundary_rescale.target == exir_ops.backend.tosa.RESCALE.default + (squeeze_view,) = tuple(boundary_rescale.users) + assert squeeze_view.target == exir_ops.edge.aten.view_copy.default + assert squeeze_view.meta["val"].shape == torch.Size((1, 8, 4)) + (boundary_permute,) = tuple(squeeze_view.users) + assert boundary_permute.target == exir_ops.edge.aten.permute_copy.default + assert boundary_permute.meta["val"].shape == torch.Size((1, 4, 8)) + + +def test_rewrite_conv1d_a16w8_shares_one_layout_boundary() -> None: + """Test that consumers of a rank-three A16W8 convolution share one boundary. + + Forking a widened INT32 branch gives every consumer its own boundary rescale + and layout permutation. Vela materialises the rank-three permutation as a + full transpose of the convolution output, so a second branch doubles it. + + """ + inputs = (torch.randn(1, 4, 8),) + gm, _ = _rewrite_a16w8_convs(A16W8Conv1dSharedConsumers(), inputs) + + conv = _get_call_function_node(gm, exir_ops.backend.tosa.CONV2D.default) + boundary_rescales = [ + node + for node in conv.users + if node.target == exir_ops.backend.tosa.RESCALE.default + ] + assert len(boundary_rescales) == 1 + + output_permutes = [ + node + for node in gm.graph.nodes + if node.op == "call_function" + and node.target == exir_ops.edge.aten.permute_copy.default + and node.meta["val"].shape == torch.Size((1, 4, 8)) + ] + assert len(output_permutes) == 1 + + def test_rewrite_conv_a16w8_preserves_int32_after_permute() -> None: r"""Test that an indirect INT32 consumer keeps a widened branch. @@ -533,6 +825,62 @@ def test_rewrite_conv_a16w8_preserves_int32_after_permute() -> None: assert len(widened_paths) == 1 +@pytest.mark.parametrize( + "depthwise,target_op,expected_weight_shape,expected_output_shape", + [ + ( + False, + exir_ops.backend.tosa.CONV2D.default, + (6, 1, 3, 4), + (1, 6, 8), + ), + ( + True, + exir_ops.backend.tosa.DEPTHWISE_CONV2D.default, + (1, 3, 4, 2), + (1, 8, 8), + ), + ], +) +def test_rewrite_conv1d_emits_atomic_rank3_layout_boundaries( + depthwise: bool, + target_op, + expected_weight_shape: tuple[int, ...], + expected_output_shape: tuple[int, ...], +) -> None: + module = Conv1dBiasModule(depthwise).eval() + edge_program = to_edge(export(module, module.get_inputs())).exported_program() + + with TosaLoweringContext(_compile_spec().tosa_spec): + result = RewriteConvPass(edge_program)(edge_program.graph_module) + assert result is not None + graph_module = result.graph_module + + conv = _get_call_function_node(graph_module, target_op) + input_view = conv.args[0] + assert isinstance(input_view, torch.fx.Node) + assert input_view.target == exir_ops.edge.aten.view_copy.default + input_permute = input_view.args[0] + assert isinstance(input_permute, torch.fx.Node) + assert input_permute.target == exir_ops.edge.aten.permute_copy.default + assert input_permute.args[1] == [0, 2, 1] + assert input_view.meta["val"].shape == torch.Size((1, 1, 8, 4)) + + weight = conv.args[1] + assert isinstance(weight, torch.fx.Node) + assert weight.meta["val"].shape == torch.Size(expected_weight_shape) + + output_view = next( + node + for node in graph_module.graph.nodes + if node.target == exir_ops.edge.aten.view_copy.default and node.args[0] is conv + ) + output_permute = next(iter(output_view.users)) + assert output_permute.target == exir_ops.edge.aten.permute_copy.default + assert output_permute.args[1] == [0, 2, 1] + assert output_permute.meta["val"].shape == torch.Size(expected_output_shape) + + @pytest.mark.skipif(not _VGF_ENABLED, reason="VGF not enabled") def test_fold_and_annotate_q_params_vgf_quant_tracks_fused_relu_qparams() -> None: exported_program = _export_quantized(TinyConvReluCat()) diff --git a/backends/arm/test/passes/test_symbolic_value_range.py b/backends/arm/test/passes/test_symbolic_value_range.py index 99dfafc93a6..698be1a8bd2 100644 --- a/backends/arm/test/passes/test_symbolic_value_range.py +++ b/backends/arm/test/passes/test_symbolic_value_range.py @@ -9,6 +9,7 @@ evaluate_symbolic_expr_values, ) from torch.fx.experimental.symbolic_shapes import ShapeEnv +from torch.utils._sympy.functions import PythonMod def _make_shape_env( @@ -70,7 +71,7 @@ def test_evaluate_symbolic_expr_values_bails_out_for_large_symbol_ranges() -> No assert evaluate_symbolic_expr_values(symint, shape_env) is None -def test_evaluate_symbolic_expr_values_does_not_require_shape_env_bounds( +def test_evaluate_symbolic_expr_values_bails_out_on_recursive_bounds( monkeypatch, ) -> None: shape_env, symint = _make_shape_env(hint=3, compiler_min=2, compiler_max=6) @@ -81,3 +82,24 @@ def raise_recursion(_expr): monkeypatch.setattr(shape_env, "bound_sympy", raise_recursion) assert evaluate_symbolic_expr_values(symint, shape_env) == {2, 3, 4, 5, 6} + + +def test_evaluate_symbolic_expr_values_handles_python_mod() -> None: + shape_env, symint = _make_shape_env(hint=3, compiler_min=2, compiler_max=6) + + assert evaluate_symbolic_expr_values( + PythonMod(16 * symint.node.expr - 7, 4), shape_env + ) == {1} + + +def test_evaluate_symbolic_expr_values_handles_python_floordiv() -> None: + class PythonFloorDiv(sympy.Function): + _torch_handler_name = "python_floordiv" + is_integer = True + nargs = (2,) + + shape_env, symint = _make_shape_env(hint=3, compiler_min=2, compiler_max=6) + + assert evaluate_symbolic_expr_values( + PythonFloorDiv(symint.node.expr, 2), shape_env + ) == {1, 2, 3} diff --git a/backends/arm/test/pytest.ini b/backends/arm/test/pytest.ini index 09b26752421..301b06d4917 100644 --- a/backends/arm/test/pytest.ini +++ b/backends/arm/test/pytest.ini @@ -1,6 +1,7 @@ [pytest] timeout = 1800 addopts = --strict-markers +pythonpath = ../../.. markers = slow: Tests that take long time xlarge: Tests that are known to use a lot of memory diff --git a/backends/arm/test/quantizer/test_generic_annotater.py b/backends/arm/test/quantizer/test_generic_annotater.py index e3d32a1372a..48cfdfa0110 100644 --- a/backends/arm/test/quantizer/test_generic_annotater.py +++ b/backends/arm/test/quantizer/test_generic_annotater.py @@ -145,6 +145,17 @@ def test_flip_tosa_INT(): ) +def test_roll_tosa_INT(): + check_annotation( + SingleOpModel( + torch.roll, + (torch.randn(2, 4),), + shifts=(1, 2), + dims=(0, 1), + ), + ) + + def test_concat_tosa_INT(): check_annotation( SingleOpModel( diff --git a/backends/arm/test/recipes/test_arm_recipes.py b/backends/arm/test/recipes/test_arm_recipes.py new file mode 100644 index 00000000000..1544863d86d --- /dev/null +++ b/backends/arm/test/recipes/test_arm_recipes.py @@ -0,0 +1,491 @@ +# 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. +"""Tests for the Arm ExportRecipe provider. + +Building a recipe only assembles a compile spec, a quantizer and a partitioner, +so none of these tests need Vela, the model converter or an FVP. Class and +method names route each test to exactly one of the existing target-less, TOSA +and VGF suites; see the ``-k`` filters in +``backends/arm/test/test_arm_backend.sh``. + +""" + +# pyre-strict + +import unittest +from typing import Any, Optional + +import torch + +from executorch.backends.arm.recipes.arm_recipe_provider import ArmRecipeProvider +from executorch.backends.arm.recipes.arm_recipe_types import ARM_BACKEND, ArmRecipeType +from executorch.export import ExportRecipe, recipe_registry, StageType +from executorch.export.export import ExportSession + + +_PROVIDER_LOGGER = "executorch.backends.arm.recipes.arm_recipe_provider" + +try: + import ethosu.vela.architecture_features # type: ignore # noqa: F401 + + _VELA_INSTALLED = True +except ImportError: + # The target-less CI job installs the Arm deps without Vela, and a recipe + # has to build there; only the accelerator check needs it. + _VELA_INSTALLED = False + + +class _AddModule(torch.nn.Module): + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x + y + + +class _ConvReluModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, kernel_size=3, padding=1) + self.relu = torch.nn.ReLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.relu(self.conv(x)) + + +def _compile_spec_value(partitioner: Any, key: str) -> Optional[str]: + for spec in partitioner.delegation_spec.compile_specs: + if spec.key == key: + value = spec.value + return value.decode() if isinstance(value, (bytes, bytearray)) else value + return None + + +def _backend_id(recipe: ExportRecipe) -> str: + return _first_partitioner(recipe).delegation_spec.backend_id + + +def _first_partitioner(recipe: ExportRecipe) -> Any: + assert recipe.lowering_recipe is not None + # Arm recipes partition every method the same way, so the list form rather + # than the per-method dict `LoweringRecipe` also accepts. + partitioners = recipe.lowering_recipe.partitioners + assert isinstance(partitioners, list) and partitioners + return partitioners[0] + + +def _global_config(recipe: ExportRecipe) -> Any: + assert recipe.quantization_recipe is not None + assert recipe.quantization_recipe.quantizers is not None + return recipe.quantization_recipe.quantizers[0].global_config # type: ignore[attr-defined] + + +def _input_activation_dtype(recipe: ExportRecipe) -> Optional[torch.dtype]: + config = _global_config(recipe) + if config is None or config.input_activation is None: + return None + return config.input_activation.dtype + + +class _ArmRecipeTestCase(unittest.TestCase): + """Re-registers the Arm provider before each test. + + Registration happens when ``backends.arm.recipes`` is imported, so it cannot + repeat: the module is already loaded. Other suites in the same process clear + the singleton registry in teardown, and under ``pytest -n`` their tests can + interleave with these. + + """ + + def setUp(self) -> None: + super().setUp() + recipe_registry.register_backend_recipe_provider(ArmRecipeProvider()) + + +class TestArmRecipeRegistration(_ArmRecipeTestCase): + def test_backend_registered(self) -> None: + self.assertIn(ARM_BACKEND, recipe_registry.list_backends()) + + def test_supported_recipes_match_enum(self) -> None: + # Catches an enum member added but never wired into a target table. + supported = recipe_registry.get_supported_recipes(ARM_BACKEND) + self.assertEqual(set(supported), set(ArmRecipeType)) + + def test_unknown_recipe_returns_none(self) -> None: + from executorch.export import RecipeType + + class _StubRecipeType(RecipeType): + FOO = "stub_foo" + + @classmethod + def get_backend_name(cls) -> str: + return "stub" + + self.assertIsNone(ArmRecipeProvider().create_recipe(_StubRecipeType.FOO)) + + +class TestTosaRecipes(_ArmRecipeTestCase): + def test_tosa_construction(self) -> None: + cases = [ + (ArmRecipeType.TOSA_FP, "arm_tosa_fp", "TOSA-1.0+FP", None), + (ArmRecipeType.TOSA_INT8, "arm_tosa_int8", "TOSA-1.0+INT", torch.int8), + ( + ArmRecipeType.TOSA_A16W8, + "arm_tosa_a16w8", + "TOSA-1.0+INT+int16", + torch.int16, + ), + ] + for recipe_type, expected_name, expected_spec, expected_act_dtype in cases: + with self.subTest(recipe_type=recipe_type): + recipe = ExportRecipe.get_recipe(recipe_type) + self.assertEqual(recipe.name, expected_name) + # A VGF spec here would emit a container TOSA cannot consume. + self.assertEqual(_backend_id(recipe), "TOSABackend") + self.assertEqual( + _compile_spec_value(_first_partitioner(recipe), "tosa_spec"), + expected_spec, + ) + if expected_act_dtype is None: + self.assertIsNone(recipe.quantization_recipe) + else: + self.assertEqual( + _input_activation_dtype(recipe), expected_act_dtype + ) + + def test_weights_are_per_channel(self) -> None: + for recipe_type in (ArmRecipeType.TOSA_INT8, ArmRecipeType.TOSA_A16W8): + with self.subTest(recipe_type=recipe_type): + weight = _global_config(ExportRecipe.get_recipe(recipe_type)).weight + self.assertEqual(weight.qscheme, torch.per_channel_symmetric) + + def test_unexpected_kwarg_warns(self) -> None: + with self.assertLogs(_PROVIDER_LOGGER, level="WARNING"): + ExportRecipe.get_recipe(ArmRecipeType.TOSA_INT8, foo=1) + + +class TestVgfRecipes(_ArmRecipeTestCase): + """Named without the ``_vgf_`` token that routes tests to the VKML suite: + + constructing a compile spec needs no model converter, so these belong in the + target-less suite that runs on every PR. + + """ + + def test_construction(self) -> None: + cases = [ + (ArmRecipeType.VGF_FP, "arm_vgf_fp", "TOSA-1.0+FP", None), + (ArmRecipeType.VGF_INT8, "arm_vgf_int8", "TOSA-1.0+INT", torch.int8), + ] + for recipe_type, expected_name, expected_spec, expected_act_dtype in cases: + with self.subTest(recipe_type=recipe_type): + recipe = ExportRecipe.get_recipe(recipe_type) + self.assertEqual(recipe.name, expected_name) + self.assertEqual( + _compile_spec_value(_first_partitioner(recipe), "tosa_spec"), + expected_spec, + ) + # A TOSA spec here would emit a flatbuffer VKML cannot load. + self.assertEqual(_backend_id(recipe), "VgfBackend") + self.assertEqual( + _compile_spec_value(_first_partitioner(recipe), "output_format"), + "vgf", + ) + if expected_act_dtype is None: + self.assertIsNone(recipe.quantization_recipe) + else: + self.assertEqual( + _input_activation_dtype(recipe), expected_act_dtype + ) + + def test_keeps_quantized_decomposed_ops(self) -> None: + # VGF consumes the quantized_decomposed QDQ ops, so ReplaceQuantNodesPass + # must not run; see _apply_replace_quant_nodes in aot_arm_compiler.py. + for recipe_type in (ArmRecipeType.VGF_INT8, ArmRecipeType.VGF_FP): + with self.subTest(recipe_type=recipe_type): + recipe = ExportRecipe.get_recipe(recipe_type) + assert recipe.lowering_recipe is not None + self.assertIsNone(recipe.lowering_recipe.edge_manager_transform_passes) + + +class TestEthosURecipes(_ArmRecipeTestCase): + def test_ethos_recipes_carry_their_pass_pipeline_config(self) -> None: + # Building the partitioner before the config is materialised silently + # drops this entry. Only the U55 subset has a non-default config. + for recipe_type in ( + ArmRecipeType.ETHOS_U55_INT8, + ArmRecipeType.ETHOS_U65_INT8, + ): + with self.subTest(recipe_type=recipe_type): + partitioner = _first_partitioner(ExportRecipe.get_recipe(recipe_type)) + self.assertIsNotNone( + _compile_spec_value(partitioner, "transform_pipeline_config") + ) + self.assertIsNone( + _compile_spec_value( + _first_partitioner( + ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U85_INT8) + ), + "transform_pipeline_config", + ) + ) + + def test_default_macs(self) -> None: + cases = [ + (ArmRecipeType.ETHOS_U55_INT8, "ethos-u55-128"), + (ArmRecipeType.ETHOS_U65_INT8, "ethos-u65-256"), + (ArmRecipeType.ETHOS_U85_INT8, "ethos-u85-256"), + ] + for recipe_type, expected_target in cases: + with self.subTest(recipe_type=recipe_type): + recipe = ExportRecipe.get_recipe(recipe_type) + self.assertEqual(recipe.name, recipe_type.value) + self.assertEqual(_input_activation_dtype(recipe), torch.int8) + partitioner = _first_partitioner(recipe) + self.assertEqual( + _compile_spec_value(partitioner, "target"), expected_target + ) + + def test_custom_macs(self) -> None: + cases = [ + (ArmRecipeType.ETHOS_U55_INT8, 32, "ethos-u55-32"), + (ArmRecipeType.ETHOS_U55_INT8, 256, "ethos-u55-256"), + (ArmRecipeType.ETHOS_U65_INT8, 512, "ethos-u65-512"), + (ArmRecipeType.ETHOS_U85_INT8, 128, "ethos-u85-128"), + (ArmRecipeType.ETHOS_U85_INT8, 2048, "ethos-u85-2048"), + ] + for recipe_type, macs, expected_target in cases: + with self.subTest(recipe_type=recipe_type, macs=macs): + recipe = ExportRecipe.get_recipe(recipe_type, macs=macs) + partitioner = _first_partitioner(recipe) + self.assertEqual( + _compile_spec_value(partitioner, "target"), expected_target + ) + + @unittest.skipUnless(_VELA_INSTALLED, "accelerator configs come from Vela") + def test_invalid_macs_raises_u55(self) -> None: + cases = [ + (ArmRecipeType.ETHOS_U55_INT8, 512), + (ArmRecipeType.ETHOS_U65_INT8, 128), + (ArmRecipeType.ETHOS_U85_INT8, 64), + (ArmRecipeType.ETHOS_U55_INT8, 999), + ] + for recipe_type, macs in cases: + with self.subTest(recipe_type=recipe_type, macs=macs): + with self.assertRaises(ValueError): + ExportRecipe.get_recipe(recipe_type, macs=macs) + + def test_pass_through_kwargs(self) -> None: + recipe = ExportRecipe.get_recipe( + ArmRecipeType.ETHOS_U55_INT8, + macs=128, + system_config="Custom_System", + memory_mode="Custom_Memory", + extra_flags=["--user-flag"], + config_ini="custom/vela.ini", + ) + partitioner = _first_partitioner(recipe) + flags = _compile_spec_value(partitioner, "compile_flags") or "" + # Vela takes the last occurrence of a repeated flag, so the defaults + # have to come first for a caller override to win. + self.assertTrue( + flags.startswith("--verbose-operators --verbose-cycle-estimate"), + f"default flags must be prepended, got {flags}", + ) + self.assertLess(flags.index("--verbose-operators"), flags.index("--user-flag")) + self.assertIn("--system-config=Custom_System", flags) + self.assertIn("--memory-mode=Custom_Memory", flags) + self.assertIn("--verbose-operators", flags) + self.assertIn("--verbose-cycle-estimate", flags) + self.assertIn("--user-flag", flags) + self.assertIn("--config=custom/vela.ini", flags) + + def test_default_vela_flags(self) -> None: + # test_pass_through_kwargs supplies every kwarg, so the defaults would + # otherwise never be built. A wrong default config path only surfaces + # when Vela runs. + partitioner = _first_partitioner( + ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U55_INT8) + ) + flags = _compile_spec_value(partitioner, "compile_flags") or "" + self.assertIn("--config=Arm/vela.ini", flags) + self.assertIn("--verbose-operators", flags) + self.assertIn("--verbose-cycle-estimate", flags) + + def test_documented_kwargs_do_not_warn(self) -> None: + # Every one of these is honoured, so reporting it as ignored would be + # a lie about what the recipe did. + with self.assertNoLogs(_PROVIDER_LOGGER, level="WARNING"): + ExportRecipe.get_recipe( + ArmRecipeType.ETHOS_U55_INT8, + macs=128, + system_config="Custom_System", + memory_mode="Custom_Memory", + extra_flags=["--user-flag"], + config_ini="custom/vela.ini", + ) + + def test_unexpected_kwarg_warns(self) -> None: + # Flags typos like `mac=128` (instead of `macs=128`), which would + # otherwise silently produce a default-target binary. + with self.assertLogs(_PROVIDER_LOGGER, level="WARNING"): + ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U55_INT8, mac=128) + + def test_extra_flags_must_be_a_list(self) -> None: + # A bare string is iterable, so it would reach Vela as one flag per + # character instead of failing. + with self.assertRaisesRegex(ValueError, "extra_flags must be a list"): + ExportRecipe.get_recipe( + ArmRecipeType.ETHOS_U55_INT8, extra_flags="--enable-debug-db" + ) + with self.assertRaisesRegex(ValueError, "extra_flags must be a list"): + ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U55_INT8, extra_flags=[1]) + # Not iterable at all: checking the elements first raises TypeError out + # of `all` and the caller never sees the real complaint. + with self.assertRaisesRegex(ValueError, "extra_flags must be a list"): + ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U55_INT8, extra_flags=7) + + def test_macs_must_be_an_int(self) -> None: + with self.assertRaisesRegex(ValueError, "macs must be an int"): + ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U55_INT8, macs="128") + + def test_program_config_matches_the_cli(self) -> None: + # The Arm runtime has only ever been run against an inline delegate + # payload, and quantized Arm graphs do not survive the edge verifier. + recipe = ExportRecipe.get_recipe(ArmRecipeType.TOSA_INT8) + assert recipe.executorch_backend_config is not None + self.assertFalse(recipe.executorch_backend_config.extract_delegate_segments) + assert recipe.lowering_recipe is not None + assert recipe.lowering_recipe.edge_compile_config is not None + self.assertFalse(recipe.lowering_recipe.edge_compile_config._check_ir_validity) + + def test_fp_recipes_run_no_post_partition_transform(self) -> None: + # No QDQ ops to rewrite, so the extra stage must not be scheduled. + recipe = ExportRecipe.get_recipe(ArmRecipeType.TOSA_FP) + assert recipe.lowering_recipe is not None + self.assertIsNone(recipe.lowering_recipe.edge_manager_transform_passes) + + def test_recipes_do_not_share_config_objects(self) -> None: + # _combine_recipes compares configs by value on the assumption that + # each provider hands out a fresh one. + first = ExportRecipe.get_recipe(ArmRecipeType.TOSA_INT8) + second = ExportRecipe.get_recipe(ArmRecipeType.TOSA_INT8) + assert first.lowering_recipe is not None + assert second.lowering_recipe is not None + self.assertIsNot( + first.lowering_recipe.edge_compile_config, + second.lowering_recipe.edge_compile_config, + ) + self.assertIsNot( + first.executorch_backend_config, second.executorch_backend_config + ) + + +class TestQuantizedRecipeLowering(_ArmRecipeTestCase): + """Every quantized recipe has to rewrite the QDQ ops left outside the + delegate, and can only do so from a stage that runs after partitioning. + """ + + QUANTIZED_RECIPES = ( + ArmRecipeType.TOSA_INT8, + ArmRecipeType.TOSA_A16W8, + ArmRecipeType.ETHOS_U55_INT8, + ArmRecipeType.ETHOS_U65_INT8, + ArmRecipeType.ETHOS_U85_INT8, + ) + + def test_replace_quant_nodes_is_wired(self) -> None: + for recipe_type in self.QUANTIZED_RECIPES: + with self.subTest(recipe_type=recipe_type): + recipe = ExportRecipe.get_recipe(recipe_type) + assert recipe.lowering_recipe is not None + self.assertTrue( + recipe.lowering_recipe.edge_manager_transform_passes, + "quantized recipes must run ReplaceQuantNodesPass", + ) + + def test_session_schedules_the_stage(self) -> None: + session = ExportSession( + model=_ConvReluModule(), + example_inputs=[(torch.randn(1, 3, 8, 8),)], + export_recipe=ExportRecipe.get_recipe(ArmRecipeType.TOSA_INT8), + ) + stages = session._pipeline_stages + self.assertIn(StageType.EDGE_PROGRAM_MANAGER_TRANSFORM, stages) + self.assertGreater( + stages.index(StageType.EDGE_PROGRAM_MANAGER_TRANSFORM), + stages.index(StageType.TO_EDGE_TRANSFORM_AND_LOWER), + ) + + +class TestTosaAOTRoundTrip(_ArmRecipeTestCase): + """End-to-end exports through the recipe pipeline. + + Ethos-U and VGF round-trips need a real compiler and are deferred to an FVP- + bearing follow-up. + + """ + + def _export( + self, + recipe: ExportRecipe, + model: torch.nn.Module, + example_inputs: tuple, + ): + from executorch.export import export + + session = export( + model=model, + example_inputs=[example_inputs], + export_recipe=recipe, + ) + return session.get_executorch_program() + + def _instruction_kinds(self, program) -> tuple[list, list]: + from executorch.exir.schema import DelegateCall, KernelCall + + instructions = program.execution_plan[0].chains[0].instructions + assert instructions is not None + operators = program.execution_plan[0].operators + delegate_calls = [ + i for i in instructions if isinstance(i.instr_args, DelegateCall) + ] + kernel_op_names = [ + operators[i.instr_args.op_index].name + for i in instructions + if isinstance(i.instr_args, KernelCall) + ] + return delegate_calls, kernel_op_names + + def test_tosa_fp_export(self) -> None: + # FP path: no quant ops, expect full delegation (Add is supported by TOSA). + program = self._export( + ExportRecipe.get_recipe(ArmRecipeType.TOSA_FP), + _AddModule(), + (torch.randn(2, 3), torch.randn(2, 3)), + ) + delegates, kernels = self._instruction_kinds(program) + self.assertEqual(len(delegates), 1, "Add should produce one TOSA delegate") + self.assertEqual( + kernels, [], f"Expected full delegation, got kernels {kernels}" + ) + + def test_tosa_int8_export(self) -> None: + # INT8 path: boundary quantize/dequantize remain outside the delegate + # and ReplaceQuantNodesPass rewrites them to cortex_m::*. + program = self._export( + ExportRecipe.get_recipe(ArmRecipeType.TOSA_INT8), + _ConvReluModule(), + (torch.randn(1, 3, 8, 8),), + ) + delegates, kernels = self._instruction_kinds(program) + self.assertGreaterEqual(len(delegates), 1, "Conv+ReLU should delegate") + for op_name in kernels: + self.assertTrue( + op_name.startswith("cortex_m::"), + f"Non-delegate kernels must be cortex_m boundary ops; got {op_name}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/arm/test/runtime/test_vgf_tensor_buffer_runtime.py b/backends/arm/test/runtime/test_vgf_tensor_buffer_runtime.py index 0baea8b832e..c24bbf6e11f 100644 --- a/backends/arm/test/runtime/test_vgf_tensor_buffer_runtime.py +++ b/backends/arm/test/runtime/test_vgf_tensor_buffer_runtime.py @@ -6,6 +6,7 @@ import sys from pathlib import Path +import pytest import torch import torch.nn.functional as F @@ -162,6 +163,11 @@ def test_two_input_add_buffer_shader_executes(tmp_path): # Covers the two-input storage-buffer shader path when both inputs are the same tensor. # Checks runtime execution matches eager output for the duplicated-input add case. +@pytest.mark.xfail( + sys.platform == "darwin", + reason="Model Converter drops duplicated custom-shader descriptor binding", + strict=True, +) @common.SkipIfNoModelConverter def test_two_input_add_buffer_shader_with_duplicated_input_executes(tmp_path): x = torch.randn(256) diff --git a/backends/arm/test/targets.bzl b/backends/arm/test/targets.bzl index 1b08a0ec4ef..a38743526cc 100644 --- a/backends/arm/test/targets.bzl +++ b/backends/arm/test/targets.bzl @@ -19,6 +19,8 @@ def define_arm_tests(): "ops/test_avg_pool2d.py", "ops/test_cat.py", "ops/test_conv2d.py", + "ops/test_isinf.py", + "ops/test_isnan.py", "ops/test_linear.py", "ops/test_log10.py", "ops/test_max_pool1d.py", @@ -48,6 +50,11 @@ def define_arm_tests(): "ops/test_split.py", ] + # Export recipes + test_files += [ + "recipes/test_arm_recipes.py", + ] + # Quantization test_files += [ "quantizer/test_generic_annotater.py", @@ -61,11 +68,12 @@ def define_arm_tests(): "misc/test_external_vela_blocks.py", # "misc/test_evaluate_model.py", "misc/test_pass_pipeline_config.py", + "misc/test_tosa_constant_pool.py", "misc/tosa_dialect/test_tosa_dialect_cast_to_block_scaled.py", "misc/tosa_dialect/test_tosa_dialect_mxfp_conv2d.py", "misc/tosa_dialect/test_tosa_dialect_mxfp_linear.py", "misc/tosa_dialect/test_tosa_resize.py", - "misc/test_tosa_spec.py", + "misc/tosa_dialect/test_tosa_spec.py", "misc/test_bn_relu_folding_qat.py", "misc/test_custom_partition.py", "misc/test_debug_hook.py", @@ -131,6 +139,7 @@ def define_arm_tests(): "//executorch/backends/arm/test/misc:dw_convs_shared_weights_module", "//executorch/backends/arm:ao_ext", "//executorch/backends/arm:ethosu", + "//executorch/backends/arm/recipes:recipes", "//executorch/backends/arm/tosa:compile_spec", "//executorch/backends/arm/tosa:partitioner", "//executorch/backends/arm:vgf", @@ -179,3 +188,18 @@ def define_arm_tests(): "fbsource//third-party/vulkan-headers-1.4.343/v1.4.343/src:vulkan-headers", ], ) + + if not runtime.is_oss and _ENABLE_VGF: + runtime.cxx_test( + name = "vgf_vulkan_features_test", + srcs = ["vgf_vulkan_features_test.cpp"], + compiler_flags = [ + "-DUSE_VULKAN_WRAPPER", + "-DUSE_VULKAN_VOLK", + ], + deps = [ + "//executorch/backends/arm/runtime:vgf_backend", + "fbsource//third-party/vulkan-headers-1.4.343/v1.4.343/src:volk_arm", + "fbsource//third-party/vulkan-headers-1.4.343/v1.4.343/src:vulkan-headers", + ], + ) diff --git a/backends/arm/test/test_arm_backend.sh b/backends/arm/test/test_arm_backend.sh index 0a209c8e356..7e2e28d9d51 100755 --- a/backends/arm/test/test_arm_backend.sh +++ b/backends/arm/test/test_arm_backend.sh @@ -12,6 +12,9 @@ script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) et_root_dir=$(cd ${script_dir}/../../.. && pwd) cd "${et_root_dir}" pwd + +# Cap pytest-xdist's `auto` workers to the container's CPU quota. +source .ci/scripts/pytest-parallelism.sh 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." diff --git a/backends/arm/test/test_arm_ootb.sh b/backends/arm/test/test_arm_ootb.sh index 7e3d110855b..d61606ba707 100755 --- a/backends/arm/test/test_arm_ootb.sh +++ b/backends/arm/test/test_arm_ootb.sh @@ -173,140 +173,7 @@ run_deit_e2e_ethos_u() { } run_mobilesam_e2e_ethos_u() { - echo "$FUNCNAME: Export, build, and run the MobileSAM e2e test" - - local example_dir="${et_root_dir}/examples/arm/mobilesam_prompt_segmentation_example_ethos_u" - local work_root="${et_root_dir}/arm_test/mobilesam_ootb_smoke" - local export_dir="${work_root}/export" - local artifact_dir="${work_root}/artifacts" - local debug_dir="${work_root}/debug" - local et_build_dir="${work_root}/cmake-out-arm" - local quantized_aot_build_dir="${work_root}/quantized_ops_aot" - local build_dir="${work_root}/runtime" - local mobile_sam_source="${work_root}/mobile_sam/source" - local image_path="${et_root_dir}/examples/models/dinov2/dog.jpg" - local pte_path="${export_dir}/mobilesam_prompt_smoke.pte" - local metadata_path="${export_dir}/mobilesam_prompt_smoke.json" - local fvp_log="${work_root}/fvp.log" - local toolchain_file="${et_root_dir}/examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake" - local input_size=448 - local fvp_timelimit="${FVP_TIMELIMIT:-300}" - echo "${FUNCNAME}: Work directory: ${work_root}; existing artifacts will be reused if present" - - mkdir -p "${export_dir}" "${artifact_dir}" "${debug_dir}" "${build_dir}" - - setup_path_script=${et_root_dir}/examples/arm/arm-scratch/setup_path.sh - source ${setup_path_script} - - source ${et_root_dir}/backends/arm/scripts/utils.sh - local n_proc="$(get_parallel_jobs)" - - echo "${FUNCNAME}: Building ExecuTorch (if needed)" - cmake --preset arm-baremetal -B "${et_build_dir}" - cmake --build "${et_build_dir}" --target install -j"$n_proc" - - echo "${FUNCNAME}: Building host quantized AOT library" - local python_executable - python_executable="$(python3 -c 'import sys; print(sys.executable)')" - cmake \ - -S "${et_root_dir}" \ - -B "${quantized_aot_build_dir}" \ - -DCMAKE_BUILD_TYPE=Release \ - -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \ - -DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON \ - -DEXECUTORCH_BUILD_XNNPACK=OFF \ - -DPYTHON_EXECUTABLE="${python_executable}" - cmake --build "${quantized_aot_build_dir}" --target quantized_ops_aot_lib -j"$n_proc" - - local quantized_ops_library - quantized_ops_library="$( - find "${quantized_aot_build_dir}/kernels/quantized" \ - -name 'libquantized_ops_aot_lib.*' \ - -type f \ - -print \ - -quit - )" - [[ -n "${quantized_ops_library}" ]] || { - echo "${FUNCNAME}: Missing quantized AOT library under ${quantized_aot_build_dir}" - return 1 - } - - echo "${FUNCNAME}: Installing example requirements" - pip install -r "${example_dir}/requirements.txt" - - echo "${FUNCNAME}: Preparing pinned MobileSAM source" - python3 "${example_dir}/model_export/prepare_mobilesam.py" \ - --source-dir "${mobile_sam_source}" - - echo "${FUNCNAME}: Exporting quantized MobileSAM PTE" - env EXECUTORCH_QUANTIZED_OPS_AOT_LIBRARY="${quantized_ops_library}" \ - python3 "${example_dir}/model_export/export_mobilesam.py" \ - --output-path "${pte_path}" \ - --calibration-image "${image_path}" \ - --eval-image "${image_path}" \ - --mobile-sam-source "${mobile_sam_source}" \ - --input-size "${input_size}" \ - --point 219 193 \ - --num-calibration-samples 1 \ - --num-eval-samples 1 \ - --num-debug-samples 1 \ - --minimum-fp32-quantized-iou 0.9 \ - --artifact-dir "${artifact_dir}" \ - --debug-output-dir "${debug_dir}" - - for artifact in \ - "${pte_path}" \ - "${metadata_path}" \ - "${export_dir}/mobilesam_prompt_smoke_delegation.txt" \ - "${export_dir}/mobilesam_prompt_smoke_metrics.json"; do - [[ -f "${artifact}" ]] || { - echo "${FUNCNAME}: Missing export artifact ${artifact}" - return 1 - } - done - - echo "${FUNCNAME}: Configuring the MobileSAM application" - cmake \ - -U "LIB_*" \ - -U executorch_DIR \ - -S "${example_dir}/runtime" \ - -B "${build_dir}" \ - -DCMAKE_TOOLCHAIN_FILE="${toolchain_file}" \ - -DET_PTE_FILE_PATH="${pte_path}" \ - -DMODEL_METADATA_PATH="${metadata_path}" \ - -DIMAGE_PATH="${image_path}" \ - -DMASK_THRESHOLD=0.0 \ - -DET_SEGMENTATION_DUMP_MASK=ON \ - -DPYTHON_EXECUTABLE="${python_executable}" \ - -DET_BUILD_DIR_PATH="${et_build_dir}" - - echo "${FUNCNAME}: Building mobilesam_prompt_segmentation_example" - cmake --build "${build_dir}" -j"$n_proc" --target mobilesam_prompt_segmentation_example - - local elf="${build_dir}/mobilesam_prompt_segmentation_example" - - echo "${FUNCNAME}: Running on FVP" - backends/arm/scripts/run_fvp.sh \ - --elf="${elf}" \ - --target=ethos-u85-256 \ - --timeout="${fvp_timelimit}" \ - --semihosting-cwd="${build_dir}" \ - --fast | tee "${fvp_log}" - - grep -q "Model executed successfully." "${fvp_log}" || { - echo "${FUNCNAME}: FVP run did not report successful execution" - return 1 - } - - python3 "${example_dir}/runtime/visualize_fvp_output.py" \ - --fvp-log "${fvp_log}" \ - --input-image "${image_path}" \ - --metadata "${metadata_path}" \ - --reference-mask "${debug_dir}/dog/quantized_mask.png" \ - --minimum-iou 0.9 \ - --output-dir "${work_root}/fvp_visual" - - echo "${FUNCNAME}: PASS" + "${et_root_dir}/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/run.sh" } run_swin2sr_e2e_vgf() { diff --git a/backends/arm/test/tester/arm_tester.py b/backends/arm/test/tester/arm_tester.py index ba84e30c809..c7e09129f27 100644 --- a/backends/arm/test/tester/arm_tester.py +++ b/backends/arm/test/tester/arm_tester.py @@ -192,7 +192,6 @@ def __init__( transform_passes: Optional[ Union[Sequence[PassType], Dict[str, Sequence[PassType]]] ] = None, - compile_spec: Optional[ArmCompileSpec] = None, ): super().__init__( default_partitioner_cls=None, @@ -232,18 +231,13 @@ def run( class ToExecutorch(BaseStages.ToExecutorch): def run_artifact(self, inputs): with TosaReferenceModelDispatch(): - # Check if the model has mutable buffers. These are not delegated to the backend - # and are handled by core ExecuTorch as I/O. In other words, the mutable buffer - # is outputted and re-inputted into the model. As we are calling the graph module - # directly, we need to ensure we handle these extra mutable inputs. - if ( - len(self.artifact.exported_program().graph_signature.buffers_to_mutate) - > 0 - ): - buffers = list(self.artifact.exported_program().buffers()) - buffers.extend(inputs) - - return self.artifact.exported_program().graph_module(*buffers) + program = self.artifact.exported_program() + # Mutable inputs and other parameters become inputs to the graph + # so we need to input these in the correct order. + # Also, execute the raw graph to preserve mutation outputs for comparison. + if program.graph_signature.buffers_to_mutate: + flat_inputs = program._graph_module_flat_inputs(inputs, {}) + return program.graph_module(*flat_inputs) else: return super().run_artifact(inputs) @@ -474,7 +468,6 @@ def to_edge_transform_and_lower( edge_compile_config, constant_methods=self.constant_methods, transform_passes=self.transform_passes, - compile_spec=self.compile_spec, ) else: if partitioners is not None: diff --git a/backends/arm/test/tester/quantize.py b/backends/arm/test/tester/quantize.py index ae3a216b528..f1b57a8e237 100644 --- a/backends/arm/test/tester/quantize.py +++ b/backends/arm/test/tester/quantize.py @@ -27,6 +27,7 @@ def __init__( is_qat: Optional[bool] = False, set_global: bool = True, fold_quantize: bool = True, + dynamic_shapes: Optional[Tuple[Any, ...]] = None, ): super().__init__( quantizer, @@ -37,6 +38,7 @@ def __init__( set_global, ) self.fold_quantize = fold_quantize + self.dynamic_shapes = dynamic_shapes def run( self, artifact: torch.nn.Module, inputs: Optional[Tuple[torch.Tensor]] @@ -44,7 +46,9 @@ def run( assert inputs is not None if self.is_qat: artifact.train() - captured_graph = export(artifact, inputs, strict=True).module() + captured_graph = export( + artifact, inputs, dynamic_shapes=self.dynamic_shapes, strict=True + ).module() if not isinstance(self.quantizer, TOSAQuantizer): raise ValueError("ArmQuantizer can only run with TOSAQuantizer.") diff --git a/backends/arm/test/tester/test_quantize.py b/backends/arm/test/tester/test_quantize.py new file mode 100644 index 00000000000..62f84db25fa --- /dev/null +++ b/backends/arm/test/tester/test_quantize.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. + +import torch + +from executorch.backends.arm.quantizer import ( + get_symmetric_quantization_config, + TOSAQuantizer, +) +from executorch.backends.arm.test.tester.quantize import ArmQuantize +from executorch.backends.arm.tosa import TosaSpecification + + +class Add(torch.nn.Module): + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x + y + + +def _quantize( + module: torch.nn.Module, + inputs: tuple[torch.Tensor, ...], + dynamic_shapes: tuple[dict[int, torch.export.Dim], ...], +) -> torch.fx.GraphModule: + quantization_config = get_symmetric_quantization_config() + quantizer = TOSAQuantizer(TosaSpecification.create_from_string("TOSA-1.0+INT")) + stage = ArmQuantize( + quantizer, + quantization_config, + dynamic_shapes=dynamic_shapes, + ) + + stage.run(module, inputs) # type: ignore[arg-type] + + return stage.artifact + + +def test_arm_quantize_preserves_dynamic_input_shape() -> None: + inputs = (torch.randn(2, 4),) + batch = torch.export.Dim("batch", min=1, max=4) + + graph_module = _quantize(torch.nn.ReLU(), inputs, ({0: batch},)) + + placeholder = next( + node for node in graph_module.graph.nodes if node.op == "placeholder" + ) + assert isinstance(placeholder.meta["val"].shape[0], torch.SymInt) + assert graph_module(torch.randn(3, 4)).shape == (3, 4) + + +def test_arm_quantize_preserves_shared_dynamic_input_shape() -> None: + inputs = (torch.randn(2, 4), torch.randn(2, 4)) + batch = torch.export.Dim("batch", min=1, max=4) + + graph_module = _quantize(Add(), inputs, ({0: batch}, {0: batch})) + + placeholders = [ + node for node in graph_module.graph.nodes if node.op == "placeholder" + ] + batch_sizes = [node.meta["val"].shape[0] for node in placeholders] + assert all(isinstance(size, torch.SymInt) for size in batch_sizes) + assert batch_sizes[0] == batch_sizes[1] + assert graph_module(torch.randn(3, 4), torch.randn(3, 4)).shape == (3, 4) diff --git a/backends/arm/test/vgf_vulkan_features_test.cpp b/backends/arm/test/vgf_vulkan_features_test.cpp new file mode 100644 index 00000000000..149e4dcd9a9 --- /dev/null +++ b/backends/arm/test/vgf_vulkan_features_test.cpp @@ -0,0 +1,48 @@ +/* + * 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 + +#include + +namespace executorch { +namespace backends { +namespace vgf { +namespace { + +TEST(VgfVulkanFeaturesTest, EnablesDataGraphShaderModule) { + VkPhysicalDeviceTensorFeaturesARM next{}; + next.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_TENSOR_FEATURES_ARM; + + const auto features = make_vgf_data_graph_features(&next); + + EXPECT_EQ( + features.sType, + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DATA_GRAPH_FEATURES_ARM); + EXPECT_EQ(features.pNext, &next); + EXPECT_EQ(features.dataGraph, VK_TRUE); + EXPECT_EQ(features.dataGraphShaderModule, VK_TRUE); +} + +// cppcheck-suppress syntaxError +TEST(VgfVulkanFeaturesTest, RequiresDataGraphShaderModuleSupport) { + VkPhysicalDeviceDataGraphFeaturesARM available{}; + available.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DATA_GRAPH_FEATURES_ARM; + + EXPECT_FALSE(vgf_data_graph_features_supported(available)); + + available.dataGraph = VK_TRUE; + EXPECT_FALSE(vgf_data_graph_features_supported(available)); + + available.dataGraphShaderModule = VK_TRUE; + EXPECT_TRUE(vgf_data_graph_features_supported(available)); +} + +} // namespace +} // namespace vgf +} // namespace backends +} // namespace executorch diff --git a/backends/arm/tosa/BUCK b/backends/arm/tosa/BUCK index b7073c97a15..7ae9919ce5c 100644 --- a/backends/arm/tosa/BUCK +++ b/backends/arm/tosa/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target(_kind = runtime.python_library, name = "schemas", srcs = ["schemas/__init__.py"], @@ -21,6 +23,16 @@ fbcode_target(_kind = runtime.python_library, ], ) +fbcode_target(_kind = runtime.python_library, + name = "constant_pool", + srcs = [ + "constant_pool.py", + ], + deps = [ + "fbsource//third-party/tosa_tools:serializer", + ], +) + fbcode_target(_kind = runtime.python_library, name = "mapping", srcs = [ @@ -93,6 +105,7 @@ fbcode_target(_kind = runtime.python_library, ], deps = [ ":compile_spec", + ":constant_pool", "//executorch/backends/arm:constants", "//executorch/backends/arm:process_node", "//executorch/backends/arm/debug:schema", diff --git a/backends/arm/tosa/backend.py b/backends/arm/tosa/backend.py index 6ec7d078674..6efe71baec8 100644 --- a/backends/arm/tosa/backend.py +++ b/backends/arm/tosa/backend.py @@ -31,6 +31,7 @@ process_placeholder, ) from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec +from executorch.backends.arm.tosa.constant_pool import TosaSerializerWithConstantPool from executorch.backends.arm.tosa.mapping import ( TOSA_CONTROL_FLOW_REGION_NAME_META, TOSA_CONTROL_FLOW_SOURCE_NODE_META, @@ -152,7 +153,7 @@ def _preprocess( # noqa: C901 artifact_path = "" version = tosa_spec.version - tosa_graph = ts.TosaSerializer( + tosa_graph = TosaSerializerWithConstantPool( artifact_path, targetMajor=version.major, targetMinor=version.minor, @@ -251,7 +252,7 @@ def _preprocess_module( # noqa: C901 graph_module: GraphModule, edge_program: ExportedProgram, compile_spec: TosaCompileSpec, - tosa_graph: ts.TosaSerializer, + tosa_graph: TosaSerializerWithConstantPool, debug_hook: DebugHook | None, submodule_name: str | None = None, containing_graph_module: GraphModule | None = None, @@ -262,9 +263,12 @@ def _preprocess_module( # noqa: C901 graph_module (GraphModule): Module to lower recursively. edge_program (ExportedProgram): Original exported program. compile_spec (TosaCompileSpec): Backend options with TOSA settings. - tosa_graph (ts.TosaSerializer): Serializer receiving operators. + tosa_graph (TosaSerializerWithConstantPool): Serializer receiving + operators. debug_hook (DebugHook | None): Optional debug instrumentation. submodule_name (str | None): Name used when visiting nested blocks. + containing_graph_module (GraphModule | None): Parent graph module for + nested control flow. Raises: RuntimeError: If an FX node with an unsupported op kind is found. diff --git a/backends/arm/tosa/constant_pool.py b/backends/arm/tosa/constant_pool.py new file mode 100644 index 00000000000..92cdc65df60 --- /dev/null +++ b/backends/arm/tosa/constant_pool.py @@ -0,0 +1,59 @@ +# 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 + +import tosa_serializer as ts + + +_ConstantKey = tuple[Any, tuple[int, ...], bytes | None] + + +def _constant_key(shape, dtype, values) -> _ConstantKey: + if dtype == ts.DType.SHAPE: + if len(shape) > 1: + raise ValueError(f"CONST_SHAPE expects rank metadata, got {shape}") + rank = 0 if len(shape) == 0 else shape[0] + constant = ts.TosaSerializerShape("", rank, values) + else: + constant = ts.TosaSerializerTensor("", shape, dtype, values) + + data = None if constant.data is None else bytes(constant.data) + return constant.dtype, tuple(constant.shape), data + + +class TosaSerializerWithConstantPool(ts.TosaSerializer): + """Pool generated constants independently within each TOSA basic block.""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + # Native tensor wrappers retain their serializer. Cache names to avoid + # an ownership cycle that Python's garbage collector cannot release. + self._block_pools: dict[Any, dict[_ConstantKey, str]] = {} + + def addConst(self, shape, dtype, vals=None, name=""): + """Return a matching constant in the current block or add a new one.""" + block = self.currRegion.currBasicBlock + pool = self._block_pools.setdefault(block, {}) + key = _constant_key(shape, dtype, vals) + if key not in pool: + constant = super().addConst(shape, dtype, vals, name) + pool[key] = constant.name + return constant + + # Resolve the cached name to the object expected by callers. TOSA stores + # shape constants separately from tensor constants. + cached_name = pool[key] + if dtype == ts.DType.SHAPE: + constant = block.getShapeByName(cached_name) + else: + constant = block.getTensorByName(cached_name) + return constant + + def addUnpooledConst(self, shape, dtype, vals=None, name=""): + """Add a constant without pooling so its requested name remains + addressable. + """ + return super().addConst(shape, dtype, vals, name) diff --git a/backends/arm/tosa/dialect/BUCK b/backends/arm/tosa/dialect/BUCK index b3a484bb8d9..3c1a070f82a 100644 --- a/backends/arm/tosa/dialect/BUCK +++ b/backends/arm/tosa/dialect/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target(_kind = runtime.python_library, name = "core", srcs = [ diff --git a/backends/arm/tosa/dialect/__init__.py b/backends/arm/tosa/dialect/__init__.py index 504323365ac..ee6e56b39fb 100644 --- a/backends/arm/tosa/dialect/__init__.py +++ b/backends/arm/tosa/dialect/__init__.py @@ -9,6 +9,7 @@ avg_pool2d, avg_pool2d_adaptive, binary_elementwise, + comparison, conv2d, conv2d_block_scaled, conv3d, diff --git a/backends/arm/tosa/dialect/ops/_common.py b/backends/arm/tosa/dialect/ops/_common.py index daeef30b097..28ac9ac84ef 100644 --- a/backends/arm/tosa/dialect/ops/_common.py +++ b/backends/arm/tosa/dialect/ops/_common.py @@ -29,6 +29,18 @@ def require_same_dtype(input1: torch.Tensor, input2: torch.Tensor, op: str) -> N ) +def binary_meta( + input1: torch.Tensor, + input2: torch.Tensor, + op: str, + *, + output_dtype: torch.dtype | None = None, +) -> torch.Tensor: + require_same_dtype(input1, input2, op) + output_shape = broadcast_shape(input1, input2, op) + return torch.empty(output_shape, dtype=output_dtype or input1.dtype) + + def validate_nan_mode(nan_mode: str, op: str) -> None: if nan_mode not in _VALID_NAN_MODES: raise TosaValueError( diff --git a/backends/arm/tosa/dialect/ops/binary_elementwise.py b/backends/arm/tosa/dialect/ops/binary_elementwise.py index 1a3f7222419..5c181ae0a80 100644 --- a/backends/arm/tosa/dialect/ops/binary_elementwise.py +++ b/backends/arm/tosa/dialect/ops/binary_elementwise.py @@ -6,8 +6,7 @@ import torch from executorch.backends.arm.tosa.dialect.lib import TosaValueError from executorch.backends.arm.tosa.dialect.ops._common import ( - broadcast_shape, - require_same_dtype, + binary_meta, validate_nan_mode, ) from executorch.backends.arm.tosa.dialect.ops_registration import register_fake_tosa_op @@ -37,18 +36,6 @@ def _raise_unsupported_profile(dtype: torch.dtype, op: str) -> None: ) -def _binary_meta( - input1: torch.Tensor, - input2: torch.Tensor, - op: str, - *, - output_dtype: torch.dtype | None = None, -) -> torch.Tensor: - require_same_dtype(input1, input2, op) - output_shape = broadcast_shape(input1, input2, op) - return torch.empty(output_shape, dtype=output_dtype or input1.dtype) - - def _require_int_profile_support(dtype: torch.dtype, op: str) -> None: if not get_context_spec().support_integer(): _raise_unsupported_profile(dtype, op) @@ -140,7 +127,7 @@ def _validate_and_infer_mul_output_dtype(dtype: torch.dtype) -> torch.dtype: # ) def ADD(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_add_sub_dtype(input1.dtype, "ADD") - return _binary_meta(input1, input2, "ADD") + return binary_meta(input1, input2, "ADD") @register_fake_tosa_op( @@ -154,7 +141,7 @@ def ARITHMETIC_RIGHT_SHIFT( round: bool = False, ) -> torch.Tensor: _validate_any_profile_int_dtype(input1.dtype, "ARITHMETIC_RIGHT_SHIFT") - return _binary_meta(input1, input2, "ARITHMETIC_RIGHT_SHIFT") + return binary_meta(input1, input2, "ARITHMETIC_RIGHT_SHIFT") @register_fake_tosa_op( @@ -163,7 +150,7 @@ def ARITHMETIC_RIGHT_SHIFT( ) def BITWISE_AND(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_bitwise_and_dtype(input1.dtype) - return _binary_meta(input1, input2, "BITWISE_AND") + return binary_meta(input1, input2, "BITWISE_AND") @register_fake_tosa_op( @@ -172,7 +159,7 @@ def BITWISE_AND(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def BITWISE_OR(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_int_dtype(input1.dtype, "BITWISE_OR") - return _binary_meta(input1, input2, "BITWISE_OR") + return binary_meta(input1, input2, "BITWISE_OR") @register_fake_tosa_op( @@ -181,34 +168,7 @@ def BITWISE_OR(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def BITWISE_XOR(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_int_dtype(input1.dtype, "BITWISE_XOR") - return _binary_meta(input1, input2, "BITWISE_XOR") - - -@register_fake_tosa_op( - "EQUAL(Tensor input1, Tensor input2) -> Tensor", - TosaSpecification.all_versions_and_profiles(), -) -def EQUAL(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: - _validate_profile_int32_or_fp_dtype(input1.dtype, "EQUAL") - return _binary_meta(input1, input2, "EQUAL", output_dtype=torch.bool) - - -@register_fake_tosa_op( - "GREATER(Tensor input1, Tensor input2) -> Tensor", - TosaSpecification.all_versions_and_profiles(), -) -def GREATER(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: - _validate_profile_int32_or_fp_dtype(input1.dtype, "GREATER") - return _binary_meta(input1, input2, "GREATER", output_dtype=torch.bool) - - -@register_fake_tosa_op( - "GREATER_EQUAL(Tensor input1, Tensor input2) -> Tensor", - TosaSpecification.all_versions_and_profiles(), -) -def GREATER_EQUAL(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: - _validate_profile_int32_or_fp_dtype(input1.dtype, "GREATER_EQUAL") - return _binary_meta(input1, input2, "GREATER_EQUAL", output_dtype=torch.bool) + return binary_meta(input1, input2, "BITWISE_XOR") @register_fake_tosa_op( @@ -217,7 +177,7 @@ def GREATER_EQUAL(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def INTDIV(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_int32_dtype(input1.dtype, "INTDIV") - return _binary_meta(input1, input2, "INTDIV") + return binary_meta(input1, input2, "INTDIV") @register_fake_tosa_op( @@ -226,7 +186,7 @@ def INTDIV(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def LOGICAL_AND(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_bool_dtype(input1.dtype, "LOGICAL_AND") - return _binary_meta(input1, input2, "LOGICAL_AND") + return binary_meta(input1, input2, "LOGICAL_AND") @register_fake_tosa_op( @@ -235,7 +195,7 @@ def LOGICAL_AND(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def LOGICAL_LEFT_SHIFT(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_any_profile_int_dtype(input1.dtype, "LOGICAL_LEFT_SHIFT") - return _binary_meta(input1, input2, "LOGICAL_LEFT_SHIFT") + return binary_meta(input1, input2, "LOGICAL_LEFT_SHIFT") @register_fake_tosa_op( @@ -244,7 +204,7 @@ def LOGICAL_LEFT_SHIFT(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tens ) def LOGICAL_RIGHT_SHIFT(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_any_profile_int_dtype(input1.dtype, "LOGICAL_RIGHT_SHIFT") - return _binary_meta(input1, input2, "LOGICAL_RIGHT_SHIFT") + return binary_meta(input1, input2, "LOGICAL_RIGHT_SHIFT") @register_fake_tosa_op( @@ -253,7 +213,7 @@ def LOGICAL_RIGHT_SHIFT(input1: torch.Tensor, input2: torch.Tensor) -> torch.Ten ) def LOGICAL_OR(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_bool_dtype(input1.dtype, "LOGICAL_OR") - return _binary_meta(input1, input2, "LOGICAL_OR") + return binary_meta(input1, input2, "LOGICAL_OR") @register_fake_tosa_op( @@ -262,7 +222,7 @@ def LOGICAL_OR(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def LOGICAL_XOR(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_bool_dtype(input1.dtype, "LOGICAL_XOR") - return _binary_meta(input1, input2, "LOGICAL_XOR") + return binary_meta(input1, input2, "LOGICAL_XOR") @register_fake_tosa_op( @@ -277,7 +237,7 @@ def MAXIMUM( ) -> torch.Tensor: validate_nan_mode(nan_mode, "MAXIMUM") _validate_profile_int32_or_fp_dtype(input1.dtype, "MAXIMUM") - return _binary_meta(input1, input2, "MAXIMUM") + return binary_meta(input1, input2, "MAXIMUM") @register_fake_tosa_op( @@ -292,7 +252,7 @@ def MINIMUM( ) -> torch.Tensor: validate_nan_mode(nan_mode, "MINIMUM") _validate_profile_int32_or_fp_dtype(input1.dtype, "MINIMUM") - return _binary_meta(input1, input2, "MINIMUM") + return binary_meta(input1, input2, "MINIMUM") @register_fake_tosa_op( @@ -315,7 +275,7 @@ def MUL( op="MUL", ) - return _binary_meta(input1, input2, "MUL", output_dtype=output_dtype) + return binary_meta(input1, input2, "MUL", output_dtype=output_dtype) @register_fake_tosa_op( @@ -324,7 +284,7 @@ def MUL( ) def POW(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_fp_dtype(input1.dtype, "POW") - return _binary_meta(input1, input2, "POW") + return binary_meta(input1, input2, "POW") @register_fake_tosa_op( @@ -333,4 +293,4 @@ def POW(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def SUB(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_add_sub_dtype(input1.dtype, "SUB") - return _binary_meta(input1, input2, "SUB") + return binary_meta(input1, input2, "SUB") diff --git a/backends/arm/tosa/dialect/ops/comparison.py b/backends/arm/tosa/dialect/ops/comparison.py new file mode 100644 index 00000000000..f01b9e0180b --- /dev/null +++ b/backends/arm/tosa/dialect/ops/comparison.py @@ -0,0 +1,78 @@ +# 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 torch +from executorch.backends.arm.tosa.dialect.lib import TosaValueError +from executorch.backends.arm.tosa.dialect.ops._common import binary_meta +from executorch.backends.arm.tosa.dialect.ops_registration import register_fake_tosa_op +from executorch.backends.arm.tosa.specification import ( + get_context_spec, + TosaSpecification, +) + +FP_DTYPES = (torch.float16, torch.float32) + + +def _dtype_name(dtype: torch.dtype) -> str: + return str(dtype).removeprefix("torch.") + + +def _raise_unsupported_dtype(dtype: torch.dtype, op: str) -> None: + raise TosaValueError(f"Unsupported dtype {dtype} for {op}", op=op) + + +def _raise_unsupported_profile(dtype: torch.dtype, op: str) -> None: + raise TosaValueError( + f"TOSA spec {get_context_spec()} doesn't support {_dtype_name(dtype)} for {op}", + op=op, + ) + + +def _validate_comparison_dtype(dtype: torch.dtype, op: str) -> None: + tosa_spec = get_context_spec() + + if dtype == torch.int32: + if not tosa_spec.support_integer(): + _raise_unsupported_profile(dtype, op) + return + + if dtype in FP_DTYPES: + if not tosa_spec.support_float(): + _raise_unsupported_profile(dtype, op) + return + + if dtype == torch.bfloat16: + if not (tosa_spec.support_float() and tosa_spec.support_extension("bf16")): + _raise_unsupported_profile(dtype, op) + return + + _raise_unsupported_dtype(dtype, op) + + +@register_fake_tosa_op( + "EQUAL(Tensor input1, Tensor input2) -> Tensor", + TosaSpecification.all_versions_and_profiles(), +) +def EQUAL(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: + _validate_comparison_dtype(input1.dtype, "EQUAL") + return binary_meta(input1, input2, "EQUAL", output_dtype=torch.bool) + + +@register_fake_tosa_op( + "GREATER(Tensor input1, Tensor input2) -> Tensor", + TosaSpecification.all_versions_and_profiles(), +) +def GREATER(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: + _validate_comparison_dtype(input1.dtype, "GREATER") + return binary_meta(input1, input2, "GREATER", output_dtype=torch.bool) + + +@register_fake_tosa_op( + "GREATER_EQUAL(Tensor input1, Tensor input2) -> Tensor", + TosaSpecification.all_versions_and_profiles(), +) +def GREATER_EQUAL(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: + _validate_comparison_dtype(input1.dtype, "GREATER_EQUAL") + return binary_meta(input1, input2, "GREATER_EQUAL", output_dtype=torch.bool) diff --git a/backends/arm/tosa/dialect/ops/resize.py b/backends/arm/tosa/dialect/ops/resize.py index 18ebe6c6210..35db1ae66f5 100644 --- a/backends/arm/tosa/dialect/ops/resize.py +++ b/backends/arm/tosa/dialect/ops/resize.py @@ -96,6 +96,8 @@ def RESIZE( validation_error = get_tosa_resize_output_hw_validation_error(output_hw) if validation_error is not None: raise TosaValueError(validation_error, op="RESIZE") + OH: int | torch.SymInt + OW: int | torch.SymInt if output_hw is None: scale_y_n, scale_y_d, scale_x_n, scale_x_d = scale offset_y, offset_x = offset diff --git a/backends/arm/tosa/partitioner.py b/backends/arm/tosa/partitioner.py index 57aa1db2777..87e207d5450 100644 --- a/backends/arm/tosa/partitioner.py +++ b/backends/arm/tosa/partitioner.py @@ -21,6 +21,7 @@ from typing import Callable, cast, List, Mapping, Optional, Sequence, Tuple import torch +from executorch.backends.arm._passes.arm_pass_manager import ArmPassManager from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor from executorch.backends.arm._passes.convert_expand_copy_to_repeat import ( calculate_multiples, @@ -28,13 +29,16 @@ from executorch.backends.arm._passes.decompose_large_stride_maxpool2d_pass import ( can_decompose_large_stride_maxpool2d, ) +from executorch.backends.arm._passes.decompose_roll_pass import can_decompose_roll from executorch.backends.arm._passes.decompose_unsupported_bilinear_resize_pass import ( is_exact_tosa_boundary_bilinear_downscale, ) +from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec from executorch.backends.arm.common.type import ensure_type -from executorch.backends.arm.constants import DQ_OPS, Q_OPS +from executorch.backends.arm.constants import DQ_OPS, MAX_RANK, Q_OPS from executorch.backends.arm.operator_support.tosa_supported_operators import ( + is_quantized, tosa_support_factory, ) from executorch.backends.arm.tosa.backend import TOSABackend @@ -50,6 +54,7 @@ from executorch.exir.graph_module import get_cond_while_submodules from torch.export.exported_program import ExportedProgram from torch.fx import GraphModule +from torch.fx.experimental.symbolic_shapes import statically_known_true from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition from torch.fx.passes.operator_support import any_chain, OperatorSupportBase @@ -120,6 +125,51 @@ def is_node_supported( return is_exact_tosa_boundary_bilinear_downscale(node, self.tosa_spec) +def _is_decomposable_roll_node( + node: torch.fx.Node, tosa_spec: TosaSpecification +) -> bool: + """Return whether backend preprocessing can decompose a roll node.""" + if node.target not in { + torch.ops.aten.roll.default, + exir_ops.edge.aten.roll.default, + }: + return False + if ( + tosa_spec.support_integer() + and not tosa_spec.support_float() + and not is_quantized(node) + ): + return False + input_node = ensure_type(torch.fx.Node, node.args[0]) + input_tensor = get_first_fake_tensor(input_node) + if not 0 < len(input_tensor.shape) <= MAX_RANK: + return False + if input_tensor.dtype not in {torch.float16, torch.float32} and not ( + input_tensor.dtype == torch.bfloat16 and tosa_spec.support_extension("bf16") + ): + return False + + dims = node.args[2] if len(node.args) > 2 else () + return can_decompose_roll(input_tensor.shape, node.args[1], dims) + + +class DecomposableRollSupported(OperatorSupportBase): + """Accept static rolls that backend preprocessing can decompose.""" + + def __init__(self, tosa_spec: TosaSpecification) -> None: + """Initialize the check with the active TOSA specification.""" + self.tosa_spec = tosa_spec + + def is_node_supported( + self, + submodules: Mapping[str, torch.nn.Module], + node: torch.fx.Node, + ) -> bool: + """Return True when backend preprocessing can decompose the roll.""" + del submodules + return _is_decomposable_roll_node(node, self.tosa_spec) + + def _is_custom_partition_op( custom_ops: set[torch._ops.OpOverload], target: object ) -> bool: @@ -151,10 +201,23 @@ def _is_noop_as_strided_copy(node: torch.fx.Node) -> bool: 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 - and input_tensor.stride() == output_tensor.stride() - and input_tensor.storage_offset() == output_tensor.storage_offset() + return bool( + len(input_tensor.shape) == len(output_tensor.shape) + and all( + statically_known_true(input_dim == output_dim) + for input_dim, output_dim in zip( + input_tensor.shape, output_tensor.shape + ) + ) + and all( + statically_known_true(input_stride == output_stride) + for input_stride, output_stride in zip( + input_tensor.stride(), output_tensor.stride() + ) + ) + and statically_known_true( + input_tensor.storage_offset() == output_tensor.storage_offset() + ) ) @@ -180,7 +243,7 @@ def _is_noop_squeeze(node: torch.fx.Node) -> bool: 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 + return bool(input_tensor.shape == output_tensor.shape) def _is_noop_flip(node: torch.fx.node.Node) -> bool: @@ -347,6 +410,8 @@ class TOSAPartitioner(Partitioner): """ + compile_spec: ArmCompileSpec + def __init__( self, compile_spec: TosaCompileSpec, @@ -367,12 +432,34 @@ def __init__( self.delegation_spec = DelegationSpec( TOSABackend.__name__, compile_spec._to_list() ) + self.compile_spec = compile_spec self.tosa_spec = compile_spec.tosa_spec self.additional_checks = additional_checks self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec) self._custom_partition_ops: set[torch._ops.OpOverload] = set() self.intermediate_path = compile_spec._get_intermediate_path() + def transform_for_pre_decomposition( + self, exported_program: ExportedProgram + ) -> ExportedProgram: + """Apply required Arm passes before default ATen decompositions. + + EXIR invokes this backend extension hook automatically through + ``to_edge_transform_and_lower``. Model export users should not call it + directly. + + Args: + exported_program (ExportedProgram): The ATen-dialect program to + transform. + + Returns: + ExportedProgram: The transformed ATen-dialect program. + + """ + return ArmPassManager( + self.compile_spec + ).transform_for_pre_decomposition_pipeline(exported_program) + def register_custom_partition_op(self, op: torch._ops.OpOverload) -> None: """Register a custom op to be considered supported.""" self._custom_partition_ops.add(op) @@ -653,6 +740,7 @@ def _create_operator_support( additional_positive_checks=[ self._decomposable_resize_support, DecomposableLargeStrideMaxPool2dForU55Supported(self.tosa_spec), + DecomposableRollSupported(self.tosa_spec), ], ) @@ -742,11 +830,31 @@ def ops_to_not_decompose( # noqa: C901 ops_to_not_decompose_always = { torch.ops.aten.logit.default, } + ops_to_not_decompose_conditionally = { + torch.ops.aten.roll.default, + } ops_to_not_decompose_if_integer = { torch.ops.aten.eye.default, torch.ops.aten.linspace.default, torch.ops.aten.silu.default, } + ops_to_not_decompose = ( + ops_to_not_decompose_always + | ops_to_not_decompose_if_quant_op + | ops_to_not_decompose_if_fp + | ops_to_not_decompose_if_integer + | ops_to_not_decompose_conditionally + ) + + if not self.tosa_spec.is_U55_subset: + # Tosa operator "RESIZE" is not supported on U55. Since + # upsample_bilinear2d and upsample_nearest2d decompose into that it + # will not be possible to delegate those operators on U55. If we + # have said here to not decompose them there will be an error saying + # the operator was not decomposed. It will not be possible for it + # to end up on either CPU or NPU. + ops_to_not_decompose.add(torch.ops.aten.upsample_nearest2d.vec) + ops_to_not_decompose.add(torch.ops.aten.upsample_bilinear2d.vec) def filter_fn(node: torch.fx.Node) -> bool: """Return True if an op should not be decomposed. @@ -763,6 +871,13 @@ def filter_fn(node: torch.fx.Node) -> bool: """ if _is_custom_partition_op(self._custom_partition_ops, node.target): return True + if ( + node.target in ops_to_not_decompose + and get_first_fake_tensor(node).dtype == torch.float64 + ): + return False + if node.target in ops_to_not_decompose_conditionally: + return _is_decomposable_roll_node(node, self.tosa_spec) if ( self.tosa_spec.support_float() and node.target in ops_to_not_decompose_if_fp @@ -833,21 +948,5 @@ def filter_fn(node: torch.fx.Node) -> bool: return True return False - ops_to_not_decompose = list( - ops_to_not_decompose_always - | ops_to_not_decompose_if_quant_op - | ops_to_not_decompose_if_fp - | ops_to_not_decompose_if_integer - ) - ops_to_not_decompose.extend(self._custom_partition_ops) - - if not self.tosa_spec.is_U55_subset: - # Tosa operator "RESIZE" is not supported on U55. Since upsample_bilinear2d - # and upsample_nearest2d decompose into that it will not be possible to - # delegate those operators on U55. If we have said here to not decompose - # them there will be an error saying the operator was not decomposed. It - # will not be possible for it to end up on either CPU or NPU. - ops_to_not_decompose.append(torch.ops.aten.upsample_nearest2d.vec) - ops_to_not_decompose.append(torch.ops.aten.upsample_bilinear2d.vec) - - return (ops_to_not_decompose, filter_fn) + ops_to_not_decompose.update(self._custom_partition_ops) + return (list(ops_to_not_decompose), filter_fn) diff --git a/backends/arm/vgf/partitioner.py b/backends/arm/vgf/partitioner.py index 8ed8d8941e3..be7de275af4 100644 --- a/backends/arm/vgf/partitioner.py +++ b/backends/arm/vgf/partitioner.py @@ -35,6 +35,7 @@ def __init__( self.delegation_spec = DelegationSpec( VgfBackend.__name__, compile_spec._to_list() ) + self.compile_spec = compile_spec self.additional_checks = additional_checks self.tosa_spec = compile_spec.tosa_spec self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec) diff --git a/backends/cadence/aot/ops_registrations.py b/backends/cadence/aot/ops_registrations.py index da82a1ea3ec..6789d17f16b 100644 --- a/backends/cadence/aot/ops_registrations.py +++ b/backends/cadence/aot/ops_registrations.py @@ -489,6 +489,13 @@ def register_fake( "rope_rotate_stacked_halves.out(Tensor input, Tensor sin_tensor, Tensor cos_tensor, Tensor? pos, *, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "quantized_rope_rotate_stacked_halves(Tensor input, Tensor sin_tensor, Tensor cos_tensor, Tensor? pos, float in_scale, int in_zero_point, float out_scale, int out_zero_point) -> (Tensor out)" +) +lib.define( + "quantized_rope_rotate_stacked_halves.out(Tensor input, Tensor sin_tensor, Tensor cos_tensor, Tensor? pos, float in_scale, int in_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor(a!)" +) + lib.define( "quantized_softmax(Tensor input, Tensor mask, int dim, int mask_type, Tensor pos, Tensor in_scale, Tensor in_zero_point, Tensor out_scale, Tensor out_zero_point) -> (Tensor out)" ) @@ -3136,6 +3143,20 @@ def rope_rotate_stacked_halves_meta( return input.new_empty(input.shape, dtype=input.dtype) +@register_fake("cadence::quantized_rope_rotate_stacked_halves") +def quantized_rope_rotate_stacked_halves_meta( + input: torch.Tensor, + sin_tensor: torch.Tensor, + cos_tensor: torch.Tensor, + pos: Optional[torch.Tensor], + in_scale: float, + in_zero_point: int, + out_scale: float, + out_zero_point: int, +) -> torch.Tensor: + return rope_rotate_stacked_halves_meta(input, sin_tensor, cos_tensor, pos) + + @register_fake("cadence::idma_copy") def copy_idma_copy_impl( src: torch.Tensor, diff --git a/backends/cadence/aot/ref_implementations.py b/backends/cadence/aot/ref_implementations.py index d3a5c853a4a..16768a8b68e 100644 --- a/backends/cadence/aot/ref_implementations.py +++ b/backends/cadence/aot/ref_implementations.py @@ -2265,6 +2265,38 @@ def rope_rotate_stacked_halves( return rotated.view(original_shape) +@impl_tracked(m, "quantized_rope_rotate_stacked_halves") +def quantized_rope_rotate_stacked_halves( + input_tensor: torch.Tensor, + sin_tensor: torch.Tensor, + cos_tensor: torch.Tensor, + pos: torch.Tensor | None, + in_scale: float, + in_zero_point: int, + out_scale: float, + out_zero_point: int, +) -> torch.Tensor: + dtype = input_tensor.dtype + dtype_limits = torch.iinfo(dtype) + dequantized = dequantize_per_tensor_common( + input_tensor, + in_scale, + in_zero_point, + dtype_limits.min, + dtype_limits.max, + dtype, + ) + rotated = rope_rotate_stacked_halves(dequantized, sin_tensor, cos_tensor, pos) + return quantize_per_tensor_common( + rotated, + out_scale, + out_zero_point, + dtype_limits.min, + dtype_limits.max, + dtype, + ) + + @impl_tracked(m, "im2row") def im2row( input_tensor: torch.Tensor, diff --git a/backends/cadence/aot/reorder_ops.py b/backends/cadence/aot/reorder_ops.py index 3d7ae7ac3d1..b105988a251 100644 --- a/backends/cadence/aot/reorder_ops.py +++ b/backends/cadence/aot/reorder_ops.py @@ -1179,17 +1179,23 @@ class PropagateSlice(RemoveOrReplacePassInterface): Handles any slice dim and any step size. """ - def __init__(self) -> None: + def __init__( + self, + additional_unary_targets: Optional[list[EdgeOpOverload]] = None, + additional_binary_targets: Optional[list[EdgeOpOverload]] = None, + ) -> None: super().__init__() - elementwise_targets = [ + unary_targets = [ exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, exir_ops.edge.cadence.quantize_per_tensor.default, exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, exir_ops.edge.cadence.dequantize_per_tensor.default, + *(additional_unary_targets or []), ] binary_targets = [ exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.mul.Tensor, + *(additional_binary_targets or []), ] self._dispatch: dict[ EdgeOpOverload, @@ -1198,7 +1204,7 @@ def __init__(self) -> None: Callable[[torch.fx.Node, torch.fx.Node], bool], ], ] = {} - for t in elementwise_targets: + for t in unary_targets: self._dispatch[t] = ( self._should_swap_elementwise, self._swap_elementwise_slice, @@ -1224,7 +1230,8 @@ def _should_swap_elementwise( def _swap_elementwise_slice( self, op_node: torch.fx.Node, slice_node: torch.fx.Node ) -> bool: - op_input = get_arg(op_node, "input", torch.fx.Node) + op_input = op_node.args[0] + assert isinstance(op_input, torch.fx.Node) graph = slice_node.graph slice_dim = get_arg(slice_node, "dim", int) @@ -1290,17 +1297,28 @@ def _swap_binary_elementwise_slice( slice_step = get_arg(slice_node, "step", int) output_shape = op_node.meta["val"].shape + output_dim = slice_dim % len(output_shape) new_args = list(op_node.args) with graph.inserting_before(op_node): for i, inp in enumerate([lhs, rhs]): - if inp.meta["val"].shape[slice_dim] == output_shape[slice_dim]: + input_shape = inp.meta["val"].shape + # Broadcasting aligns operand dimensions to the right of the output. + input_dim = output_dim - (len(output_shape) - len(input_shape)) + if ( + input_dim >= 0 + and input_shape[input_dim] == output_shape[output_dim] + ): new_slice = graph.call_function( exir_ops.edge.aten.slice_copy.Tensor, - args=(inp, slice_dim, slice_start, slice_end, slice_step), + args=(inp, input_dim, slice_start, slice_end, slice_step), ) new_slice.meta["val"] = exir_ops.edge.aten.slice_copy.Tensor( - inp.meta["val"], slice_dim, slice_start, slice_end, slice_step + inp.meta["val"], + input_dim, + slice_start, + slice_end, + slice_step, ) new_args[i] = new_slice diff --git a/backends/cadence/aot/tests/test_reorder_ops_passes.py b/backends/cadence/aot/tests/test_reorder_ops_passes.py index fee34dabaa5..f0083ff8a35 100644 --- a/backends/cadence/aot/tests/test_reorder_ops_passes.py +++ b/backends/cadence/aot/tests/test_reorder_ops_passes.py @@ -1358,6 +1358,114 @@ def test_unsupported_parent_not_swapped(self) -> None: self.assertFalse(result.modified) + def test_swap_additional_unary_target(self) -> None: + x_data = torch.randn(4, 60, 1, 1) + builder = GraphBuilder() + x = builder.placeholder("x", x_data) + relu = builder.call_operator(exir_ops.edge.aten.relu.default, args=(x,)) + sliced = builder.call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + args=(relu, 0, 0, 4, 2), + ) + builder.output([sliced]) + gm = builder.get_graph_module() + + result = transform_and_check_numerics( + gm, + (x_data,), + PropagateSlice(additional_unary_targets=[exir_ops.edge.aten.relu.default]), + ) + + self.assertTrue(result.modified) + slice_nodes = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor + ) + self.assertEqual(len(slice_nodes), 1) + relu_nodes = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.relu.default + ) + self.assertEqual(len(relu_nodes), 1) + self.assertIs(relu_nodes[0].args[0], slice_nodes[0]) + self.assertEqual(list(relu_nodes[0].meta["val"].shape), [2, 60, 1, 1]) + + def test_swap_additional_binary_target(self) -> None: + lhs_data = torch.randn(1, 60, 1, 1) + rhs_data = torch.randn(4, 60, 1, 1) + builder = GraphBuilder() + lhs = builder.placeholder("lhs", lhs_data) + rhs = builder.placeholder("rhs", rhs_data) + sub = builder.call_operator( + exir_ops.edge.aten.sub.Tensor, + args=(lhs, rhs), + ) + sliced = builder.call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + args=(sub, 0, 0, 4, 2), + ) + builder.output([sliced]) + gm = builder.get_graph_module() + + result = transform_and_check_numerics( + gm, + (lhs_data, rhs_data), + PropagateSlice(additional_binary_targets=[exir_ops.edge.aten.sub.Tensor]), + ) + + self.assertTrue(result.modified) + slice_nodes = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor + ) + self.assertEqual(len(slice_nodes), 1) + self.assertEqual(slice_nodes[0].args[0].name, "rhs") + sub_nodes = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.sub.Tensor + ) + self.assertEqual(len(sub_nodes), 1) + self.assertIs(sub_nodes[0].args[0], lhs.node) + self.assertIs(sub_nodes[0].args[1], slice_nodes[0]) + self.assertEqual(list(sub_nodes[0].meta["val"].shape), [2, 60, 1, 1]) + + def test_swap_additional_binary_target_with_mismatched_ranks(self) -> None: + lhs_data = torch.randn(2, 3, 4) + rhs_data = torch.randn(3, 4) + builder = GraphBuilder() + lhs = builder.placeholder("lhs", lhs_data) + rhs = builder.placeholder("rhs", rhs_data) + sub = builder.call_operator( + exir_ops.edge.aten.sub.Tensor, + args=(lhs, rhs), + ) + sliced = builder.call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + args=(sub, 1, 0, 2, 1), + ) + builder.output([sliced]) + gm = builder.get_graph_module() + + result = transform_and_check_numerics( + gm, + (lhs_data, rhs_data), + PropagateSlice(additional_binary_targets=[exir_ops.edge.aten.sub.Tensor]), + ) + + self.assertTrue(result.modified) + slice_nodes = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor + ) + self.assertEqual(len(slice_nodes), 2) + lhs_slice, rhs_slice = slice_nodes + self.assertIs(lhs_slice.args[0], lhs.node) + self.assertEqual(lhs_slice.args[1], 1) + self.assertEqual(list(lhs_slice.meta["val"].shape), [2, 2, 4]) + self.assertIs(rhs_slice.args[0], rhs.node) + self.assertEqual(rhs_slice.args[1], 0) + self.assertEqual(list(rhs_slice.meta["val"].shape), [2, 4]) + sub_nodes = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.sub.Tensor + ) + self.assertEqual(len(sub_nodes), 1) + self.assertEqual(list(sub_nodes[0].meta["val"].shape), [2, 2, 4]) + def test_swap_broadcast_mul_slice_on_broadcast_dim(self) -> None: """[1,60,1,1] * [4,1,1,1] → [4,60,1,1] → slice(dim=0, step=2) Only the [4,1,1,1] input should be sliced.""" diff --git a/backends/cadence/build_cadence_fusionG3.sh b/backends/cadence/build_cadence_fusionG3.sh index 47a0f9ff9bb..a95721833f4 100644 --- a/backends/cadence/build_cadence_fusionG3.sh +++ b/backends/cadence/build_cadence_fusionG3.sh @@ -55,7 +55,7 @@ if $STEPWISE_BUILD; then -DFLATCC_ALLOW_WERROR=OFF \ -Bcmake-out/backends/cadence \ backends/cadence - cmake --build cmake-out/backends/cadence -j8 + cmake --build cmake-out/backends/cadence -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) else echo "Building Cadence toolchain with ExecuTorch packages" cmake_prefix_path="${PWD}/cmake-out/lib/cmake/ExecuTorch;${PWD}/cmake-out/third-party/gflags" @@ -80,7 +80,7 @@ else -DHAVE_FNMATCH_H=OFF \ -DFLATCC_ALLOW_WERROR=OFF \ -Bcmake-out - cmake --build cmake-out --target install --config Release -j8 + cmake --build cmake-out --target install --config Release -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) fi echo "Run simple model to verify cmake build" diff --git a/backends/cadence/build_cadence_hifi4.sh b/backends/cadence/build_cadence_hifi4.sh index 22775af7082..cbac5e5b5d6 100644 --- a/backends/cadence/build_cadence_hifi4.sh +++ b/backends/cadence/build_cadence_hifi4.sh @@ -54,7 +54,7 @@ if $STEPWISE_BUILD; then -DFLATCC_ALLOW_WERROR=OFF \ -Bcmake-out/backends/cadence \ backends/cadence - cmake --build cmake-out/backends/cadence -j8 + cmake --build cmake-out/backends/cadence -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) else echo "Building Cadence toolchain with ExecuTorch packages" cmake_prefix_path="${PWD}/cmake-out/lib/cmake/ExecuTorch;${PWD}/cmake-out/third-party/gflags" @@ -78,7 +78,7 @@ else -DHAVE_FNMATCH_H=OFF \ -DFLATCC_ALLOW_WERROR=OFF \ -Bcmake-out - cmake --build cmake-out --target install --config Release -j8 + cmake --build cmake-out --target install --config Release -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) fi echo "Run simple model to verify cmake build" diff --git a/backends/cadence/build_cadence_runner.sh b/backends/cadence/build_cadence_runner.sh index 82968b196b3..57ce5a339e5 100755 --- a/backends/cadence/build_cadence_runner.sh +++ b/backends/cadence/build_cadence_runner.sh @@ -27,7 +27,7 @@ main() { -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \ -DEXECUTORCH_ENABLE_LOGGING=ON \ -Bcmake-out . - cmake --build cmake-out --target install --config Release -j16 + cmake --build cmake-out --target install --config Release -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) local example_dir=backends/cadence local build_dir="cmake-out/${example_dir}" @@ -46,7 +46,7 @@ main() { -DPYTHON_EXECUTABLE="$(which python3)" \ -B"${build_dir}" \ "${example_dir}" - cmake --build "${build_dir}" --config Release -j16 + cmake --build "${build_dir}" --config Release -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) local runner="${PWD}/${build_dir}/cadence_runner" if [[ ! -f "${runner}" ]]; then diff --git a/backends/cadence/build_cadence_vision.sh b/backends/cadence/build_cadence_vision.sh index b3972db4f31..fe2a07974d1 100755 --- a/backends/cadence/build_cadence_vision.sh +++ b/backends/cadence/build_cadence_vision.sh @@ -54,7 +54,7 @@ if $STEPWISE_BUILD; then -DFLATCC_ALLOW_WERROR=OFF \ -Bcmake-out/backends/cadence \ backends/cadence - cmake --build cmake-out/backends/cadence -j8 + cmake --build cmake-out/backends/cadence -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) else echo "Building Cadence toolchain with ExecuTorch packages" cmake_prefix_path="${PWD}/cmake-out/lib/cmake/ExecuTorch;${PWD}/cmake-out/third-party/gflags" @@ -78,7 +78,7 @@ else -DHAVE_FNMATCH_H=OFF \ -DFLATCC_ALLOW_WERROR=OFF \ -Bcmake-out - cmake --build cmake-out --target install --config Release -j8 + cmake --build cmake-out --target install --config Release -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) fi echo "Run simple model to verify cmake build" diff --git a/backends/cadence/runtime/executor_main.sh b/backends/cadence/runtime/executor_main.sh index 7d6cba09b87..5e630bdb524 100644 --- a/backends/cadence/runtime/executor_main.sh +++ b/backends/cadence/runtime/executor_main.sh @@ -24,7 +24,7 @@ cmake_install_executorch_devtools_lib() { -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \ -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \ -Bcmake-out . - cmake --build cmake-out -j9 --target install --config Release + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install --config Release } test_cmake_devtools_example_runner() { @@ -40,7 +40,7 @@ test_cmake_devtools_example_runner() { ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config Release + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release echo 'Running devtools/example_runner' ${build_dir}/example_runner --bundled_program_path="./CadenceDemoModel.bpte" diff --git a/backends/cortex_m/README.md b/backends/cortex_m/README.md index f077814d8a5..caead2cb7d3 100644 --- a/backends/cortex_m/README.md +++ b/backends/cortex_m/README.md @@ -5,7 +5,17 @@ ## Overview -The Cortex-M backend is implemented as an operator dialect/library based on [CMSIS-NN](https://github.com/ARM-software/CMSIS-NN), together with the `CortexMQuantizer` which targets supported ops, and the `CortexMPassManager` which modifies the exported program to use Cortex-M operators where possible. It is intended for use with **channels-last input** since this is what the accelerated kernels are using. +The Cortex-M backend is implemented as an operator dialect/library based on [CMSIS-NN](https://github.com/ARM-software/CMSIS-NN), together with the `CortexMQuantizer` which targets supported ops, and the `CortexMPassManager` which modifies the exported program to use Cortex-M operators where possible. + +The default AOT path retains the established channels-last input and dim-order contract. An experimental explicit-layout path accepts ordinary contiguous inputs, inserts graph-visible NHWC copies, and lowers spatial kernels to the `cortex_m::*_nhwc` operator family. Enable it with `--cortex-m-explicit-layout`; the two modes do not fall back to or mix with each other. Explicit-layout compilation fails when a spatial operator is not eligible for NHWC lowering, so models using unsupported configurations must use the legacy mode. + +### Explicit-layout migration + +The `use_explicit_layout=True` modes on `CortexMQuantizer` and `CortexMPassManager` are temporary staging APIs. They keep the experimental path isolated while the default modes and `CortexMTester` continue to exercise the legacy path. + +When explicit layout becomes the default, its support table and pass list will become the defaults in `CortexMQuantizer` and `CortexMPassManager`. The existing legacy support table and pass list will remain temporarily behind an opt-out AOT flag. This keeps the public Python entry points stable and switches `CortexMTester` to explicit layout without changing its callers. The direct NHWC kernel tests and explicit-only model tests will then be folded into the normal operator and model suites. + +After the legacy AOT compatibility period, the legacy support table, pass list, input conversion, opt-out flag, and remaining dual-mode tests will be removed. Legacy runtime operators will remain registered so programs serialized by the old AOT path continue to load. For a detailed example of the full lowering flow, see `examples/arm/cortex_m_mv2_example.ipynb`. diff --git a/backends/cortex_m/ops/BUCK b/backends/cortex_m/ops/BUCK index e538e9e1dd6..3c232ba048d 100644 --- a/backends/cortex_m/ops/BUCK +++ b/backends/cortex_m/ops/BUCK @@ -32,6 +32,7 @@ fbcode_target(_kind = runtime.python_library, "fbcode//caffe2:torch", "//executorch/backends/cortex_m/passes:passes_utils", "//executorch/backends/cortex_m/quantizer:quantization_configs", + "//executorch/exir:_warnings", ], ) diff --git a/backends/cortex_m/ops/cortex_m_ops_common.h b/backends/cortex_m/ops/cortex_m_ops_common.h index 2e3f49dd861..bcaed0a1bc7 100644 --- a/backends/cortex_m/ops/cortex_m_ops_common.h +++ b/backends/cortex_m/ops/cortex_m_ops_common.h @@ -14,6 +14,7 @@ #include #include +#include #include #include @@ -36,6 +37,11 @@ using KernelRuntimeContext = torch::executor::KernelRuntimeContext; // 16-byte alignment for MVE vector operations. constexpr size_t kCortexMMveAlignment = 16; +enum class ActivationLayout { + NCHWLogical, + NHWCLogical, +}; + // Basic tensor type / layout validation and dimension order checking inline void validate_cmsis_nn_tensor_requirements( const Tensor& input1, @@ -203,7 +209,7 @@ inline bool prepare_cmsis_pool2d_config( int64_t activation_min, int64_t activation_max, CmsisPool2DConfig& config, - bool require_channels_last = true, + ActivationLayout layout, bool allow_ceil_mode = false) { if (input.dim() != 4 || output.dim() != 4) { ET_LOG(Error, "%s: tensors must be 4-D", op_name); @@ -218,7 +224,9 @@ inline bool prepare_cmsis_pool2d_config( return false; } - if (input.size(0) != output.size(0) || input.size(1) != output.size(1)) { + const int64_t channel_dim = layout == ActivationLayout::NHWCLogical ? 3 : 1; + if (input.size(0) != output.size(0) || + input.size(channel_dim) != output.size(channel_dim)) { ET_LOG( Error, "%s: batch and channel dimensions must match between input and output", @@ -227,13 +235,21 @@ inline bool prepare_cmsis_pool2d_config( return false; } - if (require_channels_last) { - if (!is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { - ET_LOG( - Error, "%s: tensors must use channels_last dimension order", op_name); + if (layout == ActivationLayout::NHWCLogical) { + if (!executorch::runtime::is_contiguous_dim_order( + input.dim_order().data(), input.dim_order().size()) || + !executorch::runtime::is_contiguous_dim_order( + output.dim_order().data(), output.dim_order().size())) { + ET_LOG(Error, "%s: tensors must use contiguous dimension order", op_name); context.fail(Error::InvalidArgument); return false; } + } else if ( + !is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { + ET_LOG( + Error, "%s: tensors must use channels_last dimension order", op_name); + context.fail(Error::InvalidArgument); + return false; } auto check_tuple_len = [&](const Int64ArrayRef& arr, @@ -312,19 +328,29 @@ inline bool prepare_cmsis_pool2d_config( return false; } + const int64_t height_dim = layout == ActivationLayout::NHWCLogical ? 1 : 2; + const int64_t width_dim = layout == ActivationLayout::NHWCLogical ? 2 : 3; int32_t batch, channels, input_h, input_w, output_h, output_w; if (!check_int32_within_range( context, op_name, input.size(0), "input batch", batch) || !check_int32_within_range( - context, op_name, input.size(1), "input channels", channels) || + context, + op_name, + input.size(channel_dim), + "input channels", + channels) || !check_int32_within_range( - context, op_name, input.size(2), "input height", input_h) || + context, op_name, input.size(height_dim), "input height", input_h) || !check_int32_within_range( - context, op_name, input.size(3), "input width", input_w) || + context, op_name, input.size(width_dim), "input width", input_w) || !check_int32_within_range( - context, op_name, output.size(2), "output height", output_h) || + context, + op_name, + output.size(height_dim), + "output height", + output_h) || !check_int32_within_range( - context, op_name, output.size(3), "output width", output_w)) { + context, op_name, output.size(width_dim), "output width", output_w)) { return false; } diff --git a/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp b/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp index 39b6432c45a..ba8b51d18fd 100644 --- a/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp +++ b/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp @@ -1,5 +1,7 @@ /* * Copyright 2025-2026 Arm Limited and/or its affiliates. + * 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. @@ -66,7 +68,7 @@ bool validate_avg_pool2d_output_size( } // namespace // cppcheck-suppress unusedFunction -Tensor& quantized_avg_pool2d_out( +static Tensor& quantized_avg_pool2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Int64ArrayRef kernel_size, @@ -77,6 +79,7 @@ Tensor& quantized_avg_pool2d_out( const int64_t multiplier, const int64_t shift, const Tensor& scratch, + ActivationLayout layout, Tensor& out) { constexpr int32_t activation_min = std::numeric_limits::min(); constexpr int32_t activation_max = std::numeric_limits::max(); @@ -97,7 +100,7 @@ Tensor& quantized_avg_pool2d_out( activation_min, activation_max, pool_config, - true, + layout, true)) { return out; } @@ -153,5 +156,61 @@ Tensor& quantized_avg_pool2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_avg_pool2d_out( + KernelRuntimeContext& context, + const Tensor& input, + const Int64ArrayRef kernel_size, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const bool ceil_mode, + const int64_t zero_point, + const int64_t multiplier, + const int64_t shift, + const Tensor& scratch, + Tensor& out) { + return quantized_avg_pool2d_out_impl( + context, + input, + kernel_size, + stride, + padding, + ceil_mode, + zero_point, + multiplier, + shift, + scratch, + ActivationLayout::NCHWLogical, + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_avg_pool2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Int64ArrayRef kernel_size, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const bool ceil_mode, + const int64_t zero_point, + const int64_t multiplier, + const int64_t shift, + const Tensor& scratch, + Tensor& out) { + return quantized_avg_pool2d_out_impl( + context, + input, + kernel_size, + stride, + padding, + ceil_mode, + zero_point, + multiplier, + shift, + scratch, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_conv2d.cpp b/backends/cortex_m/ops/op_quantized_conv2d.cpp index 204a2b8369b..91cc893fba7 100644 --- a/backends/cortex_m/ops/op_quantized_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_conv2d.cpp @@ -1,10 +1,14 @@ /* * Copyright 2025-2026 Arm Limited and/or its affiliates. + * 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. */ +#include + #include "cortex_m_ops_common.h" namespace cortex_m { @@ -25,7 +29,8 @@ bool validate_conv2d_arguments( const Int64ArrayRef& padding, const Int64ArrayRef& dilation, const Tensor& requantize_multipliers, - const Tensor& requantize_shifts) { + const Tensor& requantize_shifts, + ActivationLayout layout) { if (input.dim() != kConvDim || weight.dim() != kConvDim || output.dim() != kConvDim) { ET_LOG(Error, "quantized_conv2d_out: tensors must be 4-D"); @@ -33,20 +38,22 @@ bool validate_conv2d_arguments( return false; } - // Check for channels_last dim_order (NHWC: 0, 2, 3, 1) - // Skip check if channels == 1, as dim_order is ambiguous in that case - if (input.size(1) > 1 && !is_channels_last_tensor(input)) { - ET_LOG( - Error, - "quantized_conv2d_out: input must have channels_last dim_order (NHWC)"); - context.fail(Error::InvalidArgument); - return false; - } - - if (output.size(1) > 1 && !is_channels_last_tensor(output)) { + if (layout == ActivationLayout::NHWCLogical) { + if (!executorch::runtime::is_contiguous_dim_order( + input.dim_order().data(), input.dim_order().size()) || + !executorch::runtime::is_contiguous_dim_order( + output.dim_order().data(), output.dim_order().size())) { + ET_LOG( + Error, + "quantized_conv2d_nhwc_out: input and output must have contiguous dim_order"); + context.fail(Error::InvalidArgument); + return false; + } + } else if ( + !is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { ET_LOG( Error, - "quantized_conv2d_out: output must have channels_last dim_order (NHWC)"); + "quantized_conv2d_out: input and output must have channels_last dim_order"); context.fail(Error::InvalidArgument); return false; } @@ -78,7 +85,8 @@ bool validate_conv2d_arguments( return false; } - const int64_t out_channels = output.size(1); + const int64_t out_channels = + output.size(layout == ActivationLayout::NHWCLogical ? 3 : 1); if (requantize_multipliers.size(0) != out_channels || requantize_shifts.size(0) != out_channels) { ET_LOG( @@ -94,7 +102,7 @@ bool validate_conv2d_arguments( } // namespace // cppcheck-suppress unusedFunction -Tensor& quantized_conv2d_out( +static Tensor& quantized_conv2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Tensor& weight, @@ -109,6 +117,7 @@ Tensor& quantized_conv2d_out( const int64_t activation_min, const int64_t activation_max, const Tensor& scratch, + ActivationLayout layout, Tensor& out) { if (!validate_conv2d_arguments( context, @@ -120,23 +129,30 @@ Tensor& quantized_conv2d_out( padding, dilation, requantize_multipliers, - requantize_shifts)) { + requantize_shifts, + layout)) { return out; } const int32_t batch = static_cast(input.size(0)); - const int32_t input_channels = static_cast(input.size(1)); - const int32_t input_height = static_cast(input.size(2)); - const int32_t input_width = static_cast(input.size(3)); + const int32_t input_channels = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t input_height = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t input_width = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); const int32_t kernel_output_channels = static_cast(weight.size(0)); const int32_t kernel_height = static_cast(weight.size(1)); const int32_t kernel_width = static_cast(weight.size(2)); const int32_t kernel_input_channels = static_cast(weight.size(3)); - const int32_t output_channels = static_cast(out.size(1)); - const int32_t output_height = static_cast(out.size(2)); - const int32_t output_width = static_cast(out.size(3)); + const int32_t output_channels = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t output_height = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t output_width = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); const int32_t input_offset_val = static_cast(input_offset); const int32_t output_offset_val = static_cast(output_offset); @@ -228,5 +244,77 @@ Tensor& quantized_conv2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_conv2d_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + Tensor& out) { + return quantized_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ActivationLayout::NCHWLogical, + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_conv2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + Tensor& out) { + return quantized_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp b/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp index 0793606de44..296fda24b56 100644 --- a/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp @@ -1,10 +1,14 @@ /* * Copyright 2025-2026 Arm Limited and/or its affiliates. + * 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. */ +#include + #include "cortex_m_ops_common.h" namespace cortex_m { @@ -26,7 +30,8 @@ bool validate_depthwise_conv2d_arguments( const Int64ArrayRef& dilation, const int64_t depth_multiplier, const Tensor& requantize_multipliers, - const Tensor& requantize_shifts) { + const Tensor& requantize_shifts, + ActivationLayout layout) { if (input.dim() != kConvDim || weight.dim() != kConvDim || output.dim() != kConvDim) { ET_LOG(Error, "quantized_depthwise_conv2d_out: tensors must be 4-D"); @@ -55,7 +60,8 @@ bool validate_depthwise_conv2d_arguments( } const int64_t weight_output_channels = weight.size(3); - const int64_t output_channels = output.size(1); + const int64_t output_channels = + output.size(layout == ActivationLayout::NHWCLogical ? 3 : 1); if (weight_output_channels != output_channels) { ET_LOG( Error, @@ -66,16 +72,22 @@ bool validate_depthwise_conv2d_arguments( return false; } - if (!is_channels_last_tensor(input)) { - ET_LOG( - Error, "quantized_depthwise_conv2d_out: input must be channels_last"); - context.fail(Error::InvalidArgument); - return false; - } - - if (!is_channels_last_tensor(output)) { + if (layout == ActivationLayout::NHWCLogical) { + if (!executorch::runtime::is_contiguous_dim_order( + input.dim_order().data(), input.dim_order().size()) || + !executorch::runtime::is_contiguous_dim_order( + output.dim_order().data(), output.dim_order().size())) { + ET_LOG( + Error, + "quantized_depthwise_conv2d_nhwc_out: input and output must have contiguous dim_order"); + context.fail(Error::InvalidArgument); + return false; + } + } else if ( + !is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { ET_LOG( - Error, "quantized_depthwise_conv2d_out: output must be channels_last"); + Error, + "quantized_depthwise_conv2d_out: input and output must be channels_last"); context.fail(Error::InvalidArgument); return false; } @@ -108,7 +120,8 @@ bool validate_depthwise_conv2d_arguments( return false; } - const int64_t input_channels = input.size(1); + const int64_t input_channels = + input.size(layout == ActivationLayout::NHWCLogical ? 3 : 1); // output_channels already extracted above for weight validation if (output_channels != input_channels * depth_multiplier) { ET_LOG( @@ -136,7 +149,7 @@ bool validate_depthwise_conv2d_arguments( } // namespace // cppcheck-suppress unusedFunction -Tensor& quantized_depthwise_conv2d_out( +static Tensor& quantized_depthwise_conv2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Tensor& weight, @@ -152,6 +165,7 @@ Tensor& quantized_depthwise_conv2d_out( const int64_t activation_min, const int64_t activation_max, const Tensor& scratch, + ActivationLayout layout, Tensor& out) { if (!validate_depthwise_conv2d_arguments( context, @@ -164,23 +178,30 @@ Tensor& quantized_depthwise_conv2d_out( dilation, depth_multiplier, requantize_multipliers, - requantize_shifts)) { + requantize_shifts, + layout)) { return out; } const int32_t batch = static_cast(input.size(0)); - const int32_t input_channels = static_cast(input.size(1)); - const int32_t input_height = static_cast(input.size(2)); - const int32_t input_width = static_cast(input.size(3)); + const int32_t input_channels = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t input_height = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t input_width = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); // Weight is in IHWO layout after permutation in the pass: [1, H, W, C_OUT] // For depthwise conv, this matches CMSIS-NN's expected format const int32_t kernel_height = static_cast(weight.size(1)); const int32_t kernel_width = static_cast(weight.size(2)); - const int32_t output_channels = static_cast(out.size(1)); - const int32_t output_height = static_cast(out.size(2)); - const int32_t output_width = static_cast(out.size(3)); + const int32_t output_channels = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t output_height = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t output_width = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); const int32_t depth_multiplier_val = static_cast(depth_multiplier); @@ -272,5 +293,81 @@ Tensor& quantized_depthwise_conv2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_depthwise_conv2d_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const int64_t depth_multiplier, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + Tensor& out) { + return quantized_depthwise_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + dilation, + depth_multiplier, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ActivationLayout::NCHWLogical, + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_depthwise_conv2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const int64_t depth_multiplier, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + Tensor& out) { + return quantized_depthwise_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + dilation, + depth_multiplier, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_max_pool2d.cpp b/backends/cortex_m/ops/op_quantized_max_pool2d.cpp index ca1b00ff340..75986e99d9c 100644 --- a/backends/cortex_m/ops/op_quantized_max_pool2d.cpp +++ b/backends/cortex_m/ops/op_quantized_max_pool2d.cpp @@ -1,5 +1,7 @@ /* - * Copyright 2026 Arm Limited and/or its affiliates. + * Copyright 2026 Arm Limited + * 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. @@ -11,7 +13,7 @@ namespace cortex_m { namespace native { // cppcheck-suppress unusedFunction -Tensor& quantized_max_pool2d_out( +static Tensor& quantized_max_pool2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Int64ArrayRef kernel_size, @@ -23,6 +25,7 @@ Tensor& quantized_max_pool2d_out( const int64_t output_zero_point, const int64_t activation_min, const int64_t activation_max, + ActivationLayout layout, Tensor& out) { CmsisPool2DConfig pool_config; if (!prepare_cmsis_pool2d_config( @@ -37,7 +40,8 @@ Tensor& quantized_max_pool2d_out( ceil_mode, activation_min, activation_max, - pool_config)) { + pool_config, + layout)) { return out; } @@ -95,5 +99,65 @@ Tensor& quantized_max_pool2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_max_pool2d_out( + KernelRuntimeContext& context, + const Tensor& input, + const Int64ArrayRef kernel_size, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const bool ceil_mode, + const int64_t input_zero_point, + const int64_t output_zero_point, + const int64_t activation_min, + const int64_t activation_max, + Tensor& out) { + return quantized_max_pool2d_out_impl( + context, + input, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + input_zero_point, + output_zero_point, + activation_min, + activation_max, + ActivationLayout::NCHWLogical, + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_max_pool2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Int64ArrayRef kernel_size, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const bool ceil_mode, + const int64_t input_zero_point, + const int64_t output_zero_point, + const int64_t activation_min, + const int64_t activation_max, + Tensor& out) { + return quantized_max_pool2d_out_impl( + context, + input, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + input_zero_point, + output_zero_point, + activation_min, + activation_max, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp b/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp index 04d57d4c693..4ac9b2338e6 100644 --- a/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp @@ -24,7 +24,8 @@ bool validate_transpose_conv2d_arguments( const std::optional& bias, const Tensor& output, const Tensor& requantize_multipliers, - const Tensor& requantize_shifts) { + const Tensor& requantize_shifts, + ActivationLayout layout) { if (input.dim() != kConvTransposeDim || weight.dim() != kConvTransposeDim || output.dim() != kConvTransposeDim) { ET_LOG(Error, "quantized_transpose_conv2d_out: tensors must be 4-D"); @@ -32,16 +33,22 @@ bool validate_transpose_conv2d_arguments( return false; } - if (!is_channels_last_tensor(input)) { + if (layout == ActivationLayout::NHWCLogical) { + if (!executorch::runtime::is_contiguous_dim_order( + input.dim_order().data(), input.dim_order().size()) || + !executorch::runtime::is_contiguous_dim_order( + output.dim_order().data(), output.dim_order().size())) { + ET_LOG( + Error, + "quantized_transpose_conv2d_nhwc_out: input and output must have contiguous dim_order"); + context.fail(Error::InvalidArgument); + return false; + } + } else if ( + !is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { ET_LOG( - Error, "quantized_transpose_conv2d_out: input must be channels_last"); - context.fail(Error::InvalidArgument); - return false; - } - - if (!is_channels_last_tensor(output)) { - ET_LOG( - Error, "quantized_transpose_conv2d_out: output must be channels_last"); + Error, + "quantized_transpose_conv2d_out: input and output must be channels_last"); context.fail(Error::InvalidArgument); return false; } @@ -68,7 +75,8 @@ bool validate_transpose_conv2d_arguments( return false; } - const int64_t out_channels = output.size(1); + const int64_t out_channels = + output.size(layout == ActivationLayout::NHWCLogical ? 3 : 1); if (requantize_multipliers.size(0) != out_channels || requantize_shifts.size(0) != out_channels) { ET_LOG( @@ -84,7 +92,7 @@ bool validate_transpose_conv2d_arguments( } // namespace // cppcheck-suppress unusedFunction -Tensor& quantized_transpose_conv2d_out( +static Tensor& quantized_transpose_conv2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Tensor& weight, @@ -101,6 +109,7 @@ Tensor& quantized_transpose_conv2d_out( const int64_t activation_max, const Tensor& scratch, const Tensor& output_scratch, + ActivationLayout layout, Tensor& out) { if (!validate_transpose_conv2d_arguments( context, @@ -109,23 +118,30 @@ Tensor& quantized_transpose_conv2d_out( bias, out, requantize_multipliers, - requantize_shifts)) { + requantize_shifts, + layout)) { return out; } const int32_t batch = static_cast(input.size(0)); - const int32_t input_channels = static_cast(input.size(1)); - const int32_t input_height = static_cast(input.size(2)); - const int32_t input_width = static_cast(input.size(3)); + const int32_t input_channels = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t input_height = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t input_width = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); const int32_t kernel_output_channels = static_cast(weight.size(0)); const int32_t kernel_height = static_cast(weight.size(1)); const int32_t kernel_width = static_cast(weight.size(2)); const int32_t kernel_input_channels = static_cast(weight.size(3)); - const int32_t output_channels = static_cast(out.size(1)); - const int32_t output_height = static_cast(out.size(2)); - const int32_t output_width = static_cast(out.size(3)); + const int32_t output_channels = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t output_height = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t output_width = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); if (kernel_output_channels != output_channels) { ET_LOG( @@ -246,5 +262,85 @@ Tensor& quantized_transpose_conv2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_transpose_conv2d_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef output_padding, + const Int64ArrayRef dilation, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + const Tensor& output_scratch, + Tensor& out) { + return quantized_transpose_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + output_padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + output_scratch, + ActivationLayout::NCHWLogical, + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_transpose_conv2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef output_padding, + const Int64ArrayRef dilation, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + const Tensor& output_scratch, + Tensor& out) { + return quantized_transpose_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + output_padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + output_scratch, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index 44e47087c11..56af4cbe9fb 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -23,6 +23,7 @@ CMSIS_SOFTMAX_SCALE, CMSIS_SOFTMAX_ZERO_POINT, ) +from executorch.exir._warnings import experimental from executorch.exir.dialects._ops import ops as exir_ops # To provide the implementation of the operators @@ -31,6 +32,11 @@ # New operator library with a custom namespace to allow fusion etc. lib = Library("cortex_m", "DEF") +_EXPLICIT_LAYOUT_EXPERIMENTAL = ( + "This explicit-layout Cortex-M operator may change while the legacy " + "dim-order operators remain supported." +) + SOFTMAX_INPUT_INTEGER_BITS = 5 @@ -915,6 +921,92 @@ def quantized_conv2d_impl( return result.to(torch.int8, memory_format=torch.channels_last) +lib.define( + "quantized_conv2d_nhwc(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] dilation, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch) -> Tensor" +) +lib.define( + "quantized_conv2d_nhwc.out(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] dilation, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("cortex_m::quantized_conv2d_nhwc") # type: ignore[misc] +@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL) # type: ignore[misc] +def quantized_conv2d_nhwc_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_conv2d_meta( + input.permute(0, 3, 1, 2), + weight, + bias, + stride, + padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_conv2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_conv2d_nhwc_impl( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_conv2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + weight, + bias, + stride, + padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED DEPTHWISE CONV2D OPERATION DEFINITION # =================================================================== @@ -1062,6 +1154,96 @@ def quantized_depthwise_conv2d_impl( return result.to(torch.int8, memory_format=torch.channels_last) +lib.define( + "quantized_depthwise_conv2d_nhwc(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] dilation, int depth_multiplier, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch) -> Tensor" +) +lib.define( + "quantized_depthwise_conv2d_nhwc.out(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] dilation, int depth_multiplier, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("cortex_m::quantized_depthwise_conv2d_nhwc") # type: ignore[misc] +@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL) # type: ignore[misc] +def quantized_depthwise_conv2d_nhwc_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + depth_multiplier: int, + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_depthwise_conv2d_meta( + input.permute(0, 3, 1, 2), + weight, + bias, + stride, + padding, + dilation, + depth_multiplier, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_depthwise_conv2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_depthwise_conv2d_nhwc_impl( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + depth_multiplier: int, + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_depthwise_conv2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + weight, + bias, + stride, + padding, + dilation, + depth_multiplier, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED TRANSPOSE_CONV2D OPERATION DEFINITION # =================================================================== @@ -1270,6 +1452,101 @@ def quantized_transpose_conv2d_impl( return result.to(torch.int8).to(memory_format=torch.channels_last) +lib.define( + "quantized_transpose_conv2d_nhwc(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] output_padding, int[] dilation, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch, " + "Tensor output_scratch) -> Tensor" +) +lib.define( + "quantized_transpose_conv2d_nhwc.out(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] output_padding, int[] dilation, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch, " + "Tensor output_scratch, *, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("cortex_m::quantized_transpose_conv2d_nhwc") # type: ignore[misc] +@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL) # type: ignore[misc] +def quantized_transpose_conv2d_nhwc_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + output_padding: Sequence[int], + dilation: Sequence[int], + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, + output_scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_transpose_conv2d_meta( + input.permute(0, 3, 1, 2), + weight, + bias, + stride, + padding, + output_padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + output_scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_transpose_conv2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_transpose_conv2d_nhwc_impl( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + output_padding: Sequence[int], + dilation: Sequence[int], + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, + output_scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_transpose_conv2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + weight, + bias, + stride, + padding, + output_padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + output_scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED AVG_POOL2D OPERATION DEFINITION # =================================================================== @@ -1365,6 +1642,73 @@ def quantized_avg_pool2d_impl( return output.to(torch.int8) +lib.define( + "quantized_avg_pool2d_nhwc(" + "Tensor input, int[] kernel_size, int[] stride, int[] padding, " + "bool ceil_mode, int zero_point, int multiplier, int shift, " + "Tensor scratch) -> Tensor" +) +lib.define( + "quantized_avg_pool2d_nhwc.out(" + "Tensor input, int[] kernel_size, int[] stride, int[] padding, " + "bool ceil_mode, int zero_point, int multiplier, int shift, " + "Tensor scratch, *, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("cortex_m::quantized_avg_pool2d_nhwc") # type: ignore[misc] +@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL) # type: ignore[misc] +def quantized_avg_pool2d_nhwc_meta( + input: torch.Tensor, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: Sequence[int], + ceil_mode: bool, + zero_point: int, + multiplier: int, + shift: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_avg_pool2d_meta( + input.permute(0, 3, 1, 2), + kernel_size, + stride, + padding, + ceil_mode, + zero_point, + multiplier, + shift, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_avg_pool2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_avg_pool2d_nhwc_impl( + input: torch.Tensor, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: Sequence[int], + ceil_mode: bool, + zero_point: int, + multiplier: int, + shift: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_avg_pool2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + kernel_size, + stride, + padding, + ceil_mode, + zero_point, + multiplier, + shift, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED MAX POOL2D OPERATION DEFINITION # =================================================================== @@ -1520,3 +1864,75 @@ def quantized_max_pool2d_impl( ) result = torch.clamp(result, activation_min, activation_max) return result.to(torch.int8).contiguous(memory_format=torch.channels_last) + + +lib.define( + "quantized_max_pool2d_nhwc(" + "Tensor input, int[] kernel_size, int[] stride, int[] padding, " + "int[] dilation, bool ceil_mode, int input_zero_point, " + "int output_zero_point, int activation_min, int activation_max) -> Tensor" +) +lib.define( + "quantized_max_pool2d_nhwc.out(" + "Tensor input, int[] kernel_size, int[] stride, int[] padding, " + "int[] dilation, bool ceil_mode, int input_zero_point, " + "int output_zero_point, int activation_min, int activation_max, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("cortex_m::quantized_max_pool2d_nhwc") # type: ignore[misc] +@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL) # type: ignore[misc] +def quantized_max_pool2d_nhwc_meta( + input: torch.Tensor, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + ceil_mode: bool, + input_zero_point: int, + output_zero_point: int, + activation_min: int, + activation_max: int, +) -> torch.Tensor: + nchw = quantized_max_pool2d_meta( + input.permute(0, 3, 1, 2), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + input_zero_point, + output_zero_point, + activation_min, + activation_max, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_max_pool2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_max_pool2d_nhwc_impl( + input: torch.Tensor, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + ceil_mode: bool, + input_zero_point: int, + output_zero_point: int, + activation_min: int, + activation_max: int, +) -> torch.Tensor: + nchw = quantized_max_pool2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + input_zero_point, + output_zero_point, + activation_min, + activation_max, + ) + return nchw.permute(0, 2, 3, 1).contiguous() diff --git a/backends/cortex_m/ops/operators.yaml b/backends/cortex_m/ops/operators.yaml index 2c85325f854..ebe7590c2f4 100644 --- a/backends/cortex_m/ops/operators.yaml +++ b/backends/cortex_m/ops/operators.yaml @@ -83,6 +83,12 @@ - arg_meta: null kernel_name: cortex_m::quantized_conv2d_out +- func: cortex_m::quantized_conv2d_nhwc.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_conv2d_nhwc_out + - func: cortex_m::quantized_depthwise_conv2d.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int depth_multiplier, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) variants: function @@ -90,23 +96,48 @@ - arg_meta: null kernel_name: cortex_m::quantized_depthwise_conv2d_out +- func: cortex_m::quantized_depthwise_conv2d_nhwc.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int depth_multiplier, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_depthwise_conv2d_nhwc_out + - func: cortex_m::quantized_transpose_conv2d.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] output_padding, int[] dilation, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, Tensor output_scratch, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: - arg_meta: null kernel_name: cortex_m::quantized_transpose_conv2d_out +- func: cortex_m::quantized_transpose_conv2d_nhwc.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] output_padding, int[] dilation, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, Tensor output_scratch, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_transpose_conv2d_nhwc_out + - func: cortex_m::quantized_avg_pool2d.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, bool ceil_mode, int zero_point, int multiplier, int shift, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: - arg_meta: null kernel_name: cortex_m::quantized_avg_pool2d_out + +- func: cortex_m::quantized_avg_pool2d_nhwc.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, bool ceil_mode, int zero_point, int multiplier, int shift, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_avg_pool2d_nhwc_out + - func: cortex_m::quantized_max_pool2d.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, int[] dilation, bool ceil_mode, int input_zero_point, int output_zero_point, int activation_min, int activation_max, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: - arg_meta: null kernel_name: cortex_m::quantized_max_pool2d_out +- func: cortex_m::quantized_max_pool2d_nhwc.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, int[] dilation, bool ceil_mode, int input_zero_point, int output_zero_point, int activation_min, int activation_max, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_max_pool2d_nhwc_out + - func: cortex_m::quantized_batch_matmul.out(Tensor lhs, int lhs_zero_point, Tensor rhs_transposed, int rhs_zero_point, int output_zero_point, int output_multiplier, int output_shift, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index f33ddbf9cf3..4d758703cb4 100644 --- a/backends/cortex_m/passes/BUCK +++ b/backends/cortex_m/passes/BUCK @@ -35,6 +35,7 @@ fbcode_target(_kind = runtime.python_library, "cortex_m_pass_manager.py", "decompose_hardswish_pass.py", "decompose_mean_pass.py", + "explicit_layout_pass.py", "matmul_to_bmm_pass.py", "quantized_clamp_activation_pass.py", "scratch_buffer_sizes.py", @@ -48,8 +49,15 @@ fbcode_target(_kind = runtime.python_library, "//executorch/backends/cortex_m/passes:passes_utils", "//executorch/backends/cortex_m/passes:replace_quant_nodes_pass", "//executorch/backends/transforms:aten_to_dialect_pass", + "//executorch/backends/transforms:canonicalize_view_copy_permute_pass", + "//executorch/backends/transforms:channels_last_layout", + "//executorch/backends/transforms:channels_last_ops", + "//executorch/backends/transforms:convert_conv1d_to_conv2d_pass", "//executorch/backends/transforms:remove_getitem_op", + "//executorch/backends/transforms:remove_permutes_around_elementwise_ops", "//executorch/backends/transforms:replace_scalar_with_tensor", + "//executorch/backends/transforms:replace_ops_with_channels_last_variants", + "//executorch/backends/transforms:replace_squeeze_unsqueeze_with_view", "//executorch/backends/transforms:utils", "//executorch/exir:lib", "//executorch/exir:pass_base", diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py index af60df686f5..5799896e17c 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py @@ -10,6 +10,7 @@ from typing import cast, Optional import executorch.backends.cortex_m.ops.operators # noqa +import executorch.backends.transforms.channels_last_ops # noqa: F401 import executorch.exir as exir import torch import torch.fx @@ -75,6 +76,10 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: ) for node in result.graph_module.graph.nodes: + if getattr(node.target, "namespace", None) == "channels_last": + raise RuntimeError( + f"Cortex-M lowering left {node.target} in the graph." + ) self._initialize_alloc_node_size(node) return PassResult(result.graph_module, result.modified or max_pool_modified) @@ -448,6 +453,9 @@ def _get_linear_replacement( return DialectNodeSpec(exir_ops.edge.cortex_m.quantized_linear.default, args) +@AtenToCortexMPass.register_dialect_substitution( + exir_ops.edge.channels_last.convolution.default +) @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.convolution.default) def _get_convolution_replacement( node: Node, dialect_pass: AtenToDialectPass @@ -455,6 +463,8 @@ def _get_convolution_replacement( if not _has_qparams(node): return None + explicit_nhwc = node.target == exir_ops.edge.channels_last.convolution.default + exported_program = dialect_pass.exported_program conv_args = node.args ( @@ -612,7 +622,11 @@ def _get_convolution_replacement( scratch, ) return DialectNodeSpec( - exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default, + ( + exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default + ), depthwise_args, ) @@ -634,7 +648,14 @@ def _get_convolution_replacement( output_qmax, scratch, ) - return DialectNodeSpec(exir_ops.edge.cortex_m.quantized_conv2d.default, conv2d_args) + return DialectNodeSpec( + ( + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.quantized_conv2d.default + ), + conv2d_args, + ) def _get_transpose_conv2d_replacement( @@ -646,6 +667,7 @@ def _get_transpose_conv2d_replacement( if not _has_qparams(node): return None + explicit_nhwc = node.target == exir_ops.edge.channels_last.convolution.default exported_program = dialect_pass.exported_program conv_t_args = node.args ( @@ -758,7 +780,12 @@ def _get_transpose_conv2d_replacement( output_scratch, ) return DialectNodeSpec( - exir_ops.edge.cortex_m.quantized_transpose_conv2d.default, new_args + ( + exir_ops.edge.cortex_m.quantized_transpose_conv2d_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.quantized_transpose_conv2d.default + ), + new_args, ) @@ -825,12 +852,16 @@ def _get_bmm_replacement( @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.avg_pool2d.default) +@AtenToCortexMPass.register_dialect_substitution( + exir_ops.edge.channels_last.avg_pool2d.default +) def _get_avg_pool2d_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: if not _has_qparams(node): return None + explicit_nhwc = node.target == exir_ops.edge.channels_last.avg_pool2d.default exported_program = dialect_pass.exported_program pool_args = node.args kernel_size = cast(list[int], pool_args[1]) @@ -851,8 +882,11 @@ def _get_avg_pool2d_replacement( avg_padding = padding if count_include_pad: pad_h, pad_w = padding - input_tensor = get_first_fake_tensor(input_node) - pre_pad = post_pad = to_physical_order([0, 0, pad_h, pad_w], input_tensor) + if explicit_nhwc: + pre_pad = post_pad = [0, pad_h, pad_w, 0] + else: + input_tensor = get_first_fake_tensor(input_node) + pre_pad = post_pad = to_physical_order([0, 0, pad_h, pad_w], input_tensor) with node.graph.inserting_before(node): input_node = node.graph.create_node( "call_function", @@ -875,7 +909,12 @@ def _get_avg_pool2d_replacement( scratch, ) return DialectNodeSpec( - exir_ops.edge.cortex_m.quantized_avg_pool2d.default, new_args + ( + exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.quantized_avg_pool2d.default + ), + new_args, ) @@ -1108,10 +1147,14 @@ def _get_softmax_replacement( @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.max_pool2d.default) +@AtenToCortexMPass.register_dialect_substitution( + exir_ops.edge.channels_last.max_pool2d.default +) def _get_max_pool2d_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: del dialect_pass + explicit_nhwc = node.target == exir_ops.edge.channels_last.max_pool2d.default input_qparams = node.meta.get("input_qparams", {}).get(0) cortex_m_meta = node.meta.get("custom", {}).get("cortex_m", {}) if input_qparams is None or cortex_m_meta.get("skip_quantized_max_pool2d", False): @@ -1169,6 +1212,12 @@ def _get_max_pool2d_replacement( activation_min, activation_max, ) + if explicit_nhwc: + quantized_op = getattr( + exir_ops.edge.cortex_m, "quantized_max_pool2d_nhwc", None + ) + if quantized_op is None: + return None return DialectNodeSpec(quantized_op.default, args) @@ -1194,6 +1243,14 @@ def _get_maximum_replacement( return DialectNodeSpec(exir_ops.edge.cortex_m.maximum.default, node.args) +def _transpose_spec(node: Node, input_tensor) -> DialectNodeSpec: + rank = len(input_tensor.shape) + perms = [p % rank for p in cast(tuple[int, ...], node.args[1])] + return DialectNodeSpec( + exir_ops.edge.cortex_m.transpose.default, (node.args[0], perms) + ) + + @AtenToCortexMPass.register_dialect_substitution( exir_ops.edge.aten.permute_copy.default ) @@ -1202,14 +1259,19 @@ def _get_permute_replacement( ) -> DialectNodeSpec | None: del dialect_pass input_tensor = _get_input_tensor_data(node) - if input_tensor.dtype != torch.int8: + if input_tensor.dtype != torch.int8 or not 1 <= input_tensor.dim() <= 4: return None + return _transpose_spec(node, input_tensor) - rank = len(input_tensor.shape) - perms = [p % rank for p in cast(tuple[int, ...], node.args[1])] - return DialectNodeSpec( - exir_ops.edge.cortex_m.transpose.default, (node.args[0], perms) - ) + +@AtenToCortexMPass.register_dialect_substitution( + exir_ops.edge.channels_last.permute_copy.default +) +def _get_layout_permute_replacement( + node: Node, dialect_pass: AtenToDialectPass +) -> DialectNodeSpec | None: + del dialect_pass + return _transpose_spec(node, _get_input_tensor_data(node)) @AtenToCortexMPass.register_dialect_substitution( diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index 892baf136ed..a15a7262660 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -13,10 +13,19 @@ ScalarsToAttributePass, ) from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig +from executorch.backends.transforms.convert_conv1d_to_conv2d_pass import ( + ConvertConv1dToConv2dPass, +) from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass +from executorch.backends.transforms.remove_permutes_around_elementwise_ops import ( + RemovePermutesAroundElementwiseOps, +) from executorch.backends.transforms.replace_scalar_with_tensor import ( ReplaceScalarWithTensorArgPass, ) +from executorch.backends.transforms.replace_squeeze_unsqueeze_with_view import ( + ReplaceSqueezeAndUnsqueezeWithViewPass, +) from executorch.exir.pass_base import ExportPass from executorch.exir.pass_manager import PassManager from executorch.exir.program._program import _transform, lift_constant_tensor_pass @@ -27,6 +36,11 @@ from .clamp_hardswish_pass import ClampHardswishPass from .decompose_hardswish_pass import DecomposeHardswishPass from .decompose_mean_pass import DecomposeMeanPass +from .explicit_layout_pass import ( + CortexMCanonicalizeViewCopyPermutePass, + CortexMReplaceOpsWithChannelsLastVariants, + ValidateCortexMExplicitLayoutPass, +) from .matmul_to_bmm_pass import MatmulToBmmPass from .quantized_clamp_activation_pass import QuantizedClampActivationPass from .replace_quant_nodes_pass import ReplaceQuantNodesPass @@ -35,7 +49,7 @@ class CortexMPassManager(PassManager): - pass_list: list[PassClass] = [ + legacy_pass_list: list[PassClass] = [ # Run before folding so qparams attach to max_pool2d values, not tuple + getitem. RemoveGetItemPass, FoldAndAnnotateQParamsPass, @@ -47,6 +61,26 @@ class CortexMPassManager(PassManager): AtenToCortexMPass, ] + explicit_layout_pass_list: list[PassClass] = [ + RemoveGetItemPass, + FoldAndAnnotateQParamsPass, + ReplaceScalarWithTensorArgPass, + ActivationFusionPass, + QuantizedClampActivationPass, + DecomposeHardswishPass, + ConvertConv1dToConv2dPass, + CortexMReplaceOpsWithChannelsLastVariants, + ReplaceSqueezeAndUnsqueezeWithViewPass, + CortexMCanonicalizeViewCopyPermutePass, + RemovePermutesAroundElementwiseOps, + CortexMCanonicalizeViewCopyPermutePass, + ValidateCortexMExplicitLayoutPass, + ReplaceQuantNodesPass, + AtenToCortexMPass, + ] + + pass_list = legacy_pass_list + pass_list_transform_for_annotation: list[PassClass] = [ ScalarsToAttributePass, ReplaceScalarWithTensorArgPass, @@ -61,6 +95,7 @@ def __init__( exported_program: ExportedProgram | None, passes: Optional[list[PassClass]] = None, target_config: Optional[CortexMTargetConfig] = None, + use_explicit_layout: bool = False, ) -> None: """Initialize the Cortex-M pass manager. @@ -69,17 +104,25 @@ def __init__( before calling ``transform()``; may be ``None`` for callers that only use ``transform_for_annotation()``. passes: Optional override of the pass list. Defaults to - ``CortexMPassManager.pass_list``. + the legacy or explicit-layout pass list selected by + ``use_explicit_layout``. target_config: Compilation target for passes that need it. Defaults to ``CortexMTargetConfig(cpu=CortexM.M55)``, which resolves through cmsis_nn to the MVE backend — matching the pre-config historical behaviour. + use_explicit_layout: Select the experimental explicit-layout pass + sequence. Legacy lowering remains the default. """ super().__init__(passes=[]) self.exported_program = exported_program # PassManager.passes is typed as callables; this manager stores pass classes which are initialized at transform time with the exported_program. + default_passes = ( + self.explicit_layout_pass_list + if use_explicit_layout + else self.legacy_pass_list + ) self.passes: list[PassClass] = ( # type: ignore[assignment] - passes if passes is not None else self.pass_list # type: ignore[assignment] + passes if passes is not None else default_passes # type: ignore[assignment] ) self.target_config: CortexMTargetConfig = target_config or CortexMTargetConfig( cpu=CortexM.M55 diff --git a/backends/cortex_m/passes/explicit_layout_pass.py b/backends/cortex_m/passes/explicit_layout_pass.py new file mode 100644 index 00000000000..19481bd0314 --- /dev/null +++ b/backends/cortex_m/passes/explicit_layout_pass.py @@ -0,0 +1,156 @@ +# 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. + +import executorch.backends.transforms.channels_last_ops # noqa: F401 + +import torch +from executorch.backends.cortex_m.passes.passes_utils import ( + coerce_int_pair, + skips_quantized_max_pool2d, +) +from executorch.backends.transforms.canonicalize_view_copy_permute_pass import ( + CanonicalizeViewCopyPermutePass, +) +from executorch.backends.transforms.channels_last_layout import ( + LAYOUT_PERMUTE_COPY, + PERMUTE_COPY_TARGETS, +) +from executorch.backends.transforms.replace_ops_with_channels_last_variants import ( + ChannelsLastOpSpec, + ReplaceOpsWithChannelsLastVariants, +) +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx import GraphModule, Node +from torch.fx.node import Target + + +def _is_rank4(node: Node) -> bool: + return len(node.meta["val"].shape) == 4 + + +def _has_per_tensor_qparam(node: Node, key: str, index: int) -> bool: + qparam = node.meta.get(key, {}).get(index) + return qparam is not None and not getattr(qparam, "per_channel", False) + + +def _has_input_and_output_qparams(node: Node) -> bool: + return _has_per_tensor_qparam(node, "input_qparams", 0) and _has_per_tensor_qparam( + node, "output_qparams", 0 + ) + + +def _supports_avg_pool2d(node: Node) -> bool: + divisor_override = node.args[6] if len(node.args) > 6 else None + return ( + _is_rank4(node) + and _has_input_and_output_qparams(node) + and divisor_override is None + ) + + +def _supports_max_pool2d(node: Node) -> bool: + if not _is_rank4(node) or not _has_per_tensor_qparam(node, "input_qparams", 0): + return False + if skips_quantized_max_pool2d(node): + return False + + dilation = coerce_int_pair(node.args[4] if len(node.args) > 4 else None, (1, 1)) + ceil_mode = bool(node.args[5]) if len(node.args) > 5 else False + if dilation != (1, 1) or ceil_mode: + return False + + input_qparams = node.meta["input_qparams"][0] + output_qparams = node.meta.get("output_qparams", {}).get(0) + return output_qparams is None or ( + not getattr(output_qparams, "per_channel", False) + and abs(float(input_qparams.scale) - float(output_qparams.scale)) <= 1e-6 + and int(input_qparams.zp) == int(output_qparams.zp) + ) + + +_EXPLICIT_LAYOUT_OP_MAP: dict[Target, ChannelsLastOpSpec] = { + exir_ops.edge.aten.convolution.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.convolution.default, + input_indices=[0], + output_indices=[0], + filter_fn=lambda node: _is_rank4(node) and _has_input_and_output_qparams(node), + ), + exir_ops.edge.aten.avg_pool2d.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.avg_pool2d.default, + input_indices=[0], + output_indices=[0], + filter_fn=_supports_avg_pool2d, + ), + exir_ops.edge.aten.max_pool2d.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.max_pool2d.default, + input_indices=[0], + output_indices=[0], + filter_fn=_supports_max_pool2d, + ), +} + + +class CortexMReplaceOpsWithChannelsLastVariants(ReplaceOpsWithChannelsLastVariants): + def __init__(self, exported_program: ExportedProgram) -> None: + super().__init__(exported_program, op_map=dict(_EXPLICIT_LAYOUT_OP_MAP)) + + +class CortexMCanonicalizeViewCopyPermutePass(CanonicalizeViewCopyPermutePass): + def __init__(self) -> None: + super().__init__(permute_targets=PERMUTE_COPY_TARGETS) + + def _set_node_op(self, node, target, input_node, arg) -> None: + super()._set_node_op(node, target, input_node, arg) + if node.target != LAYOUT_PERMUTE_COPY: + return + input_value = input_node.meta.get("val") + if isinstance(input_value, torch.Tensor): + dims = [int(dim) % input_value.dim() for dim in arg] + node.meta["val"] = input_value.new_empty( + tuple(input_value.shape[dim] for dim in dims) + ) + + +class ValidateCortexMExplicitLayoutPass(ExportPass): + def call(self, graph_module: GraphModule) -> PassResult: + for node in graph_module.graph.nodes: + if node.target in _EXPLICIT_LAYOUT_OP_MAP: + raise RuntimeError( + "Cortex-M explicit layout requires every quantized spatial " + f"operator to be NHWC-eligible, but {node.target} was not. " + "Use the legacy layout pipeline for this model." + ) + + if node.target != LAYOUT_PERMUTE_COPY: + continue + input_node = node.args[0] + dims = node.args[1] if len(node.args) > 1 else None + input_value = ( + input_node.meta.get("val") if isinstance(input_node, Node) else None + ) + if ( + not isinstance(input_value, torch.Tensor) + or input_value.dtype != torch.int8 + ): + raise RuntimeError( + f"Cortex-M layout copy {node.name} must move an int8 tensor." + ) + rank = input_value.dim() + if ( + not 1 <= rank <= 4 + or not isinstance(dims, (list, tuple)) + or len(dims) != rank + or not all(isinstance(dim, int) for dim in dims) + or sorted(dim % rank for dim in dims) != list(range(rank)) + ): + raise RuntimeError( + f"Cortex-M layout copy {node.name} has invalid permutation " + f"{dims!r} for rank {rank}." + ) + + return PassResult(graph_module, False) diff --git a/backends/cortex_m/passes/scratch_buffer_sizes.py b/backends/cortex_m/passes/scratch_buffer_sizes.py index b247e2be944..65a3a178757 100644 --- a/backends/cortex_m/passes/scratch_buffer_sizes.py +++ b/backends/cortex_m/passes/scratch_buffer_sizes.py @@ -4,6 +4,7 @@ # LICENSE file in the root directory of this source tree. from collections.abc import Callable +from functools import partial from typing import Any, cast import executorch.backends.cortex_m.ops.operators # noqa @@ -37,6 +38,7 @@ def _shape_from_node(node: torch.fx.Node) -> torch.Size: def _get_common_conv_buffer_size_inputs( conv_node: torch.fx.Node, *, + nhwc_logical: bool = False, stride_arg_idx: int = 3, padding_arg_idx: int = 4, dilation_arg_idx: int = 5, @@ -54,13 +56,14 @@ def _get_common_conv_buffer_size_inputs( padding = cast(list[int], conv_node.args[padding_arg_idx]) dilation = cast(list[int], conv_node.args[dilation_arg_idx]) - # Input is NCHW (PyTorch); CMSIS-NN wants NHWC dims. - n, c_in, height, width = _shape_from_node(x) - weight_shape = _shape_from_node(weight) - # Output is NCHW; convert to NHWC dims. - out_n, out_c, out_h, out_w = _shape_from_node(conv_node) + if nhwc_logical: + n, height, width, c_in = _shape_from_node(x) + out_n, out_h, out_w, out_c = _shape_from_node(conv_node) + else: + n, c_in, height, width = _shape_from_node(x) + out_n, out_c, out_h, out_w = _shape_from_node(conv_node) input_nhwc = [n, height, width, c_in] output_nhwc = [out_n, out_h, out_w, out_c] @@ -81,6 +84,7 @@ def _get_common_conv_buffer_size_inputs( def cmsis_nn_conv_buffer_size( backend: cmsis_nn.Backend, conv_node: torch.fx.Node, + nhwc_logical: bool = False, ) -> list[int]: ( input_nhwc, @@ -89,7 +93,9 @@ def cmsis_nn_conv_buffer_size( stride_hw, padding_hw, dilation_hw, - ) = _get_common_conv_buffer_size_inputs(conv_node=conv_node) + ) = _get_common_conv_buffer_size_inputs( + conv_node=conv_node, nhwc_logical=nhwc_logical + ) input_offset = cast(int, conv_node.args[6]) output_offset = cast(int, conv_node.args[7]) output_qmin = cast(int, conv_node.args[10]) @@ -122,6 +128,7 @@ def cmsis_nn_conv_buffer_size( def cmsis_nn_depthwise_conv_buffer_size( backend: cmsis_nn.Backend, conv_node: torch.fx.Node, + nhwc_logical: bool = False, ) -> list[int]: ( input_nhwc, @@ -130,7 +137,9 @@ def cmsis_nn_depthwise_conv_buffer_size( stride_hw, padding_hw, dilation_hw, - ) = _get_common_conv_buffer_size_inputs(conv_node=conv_node) + ) = _get_common_conv_buffer_size_inputs( + conv_node=conv_node, nhwc_logical=nhwc_logical + ) depth_multiplier = cast(int, conv_node.args[6]) input_offset = cast(int, conv_node.args[7]) output_offset = cast(int, conv_node.args[8]) @@ -185,6 +194,7 @@ def cmsis_nn_batch_matmul_buffer_size( def cmsis_nn_transpose_conv_buffer_size( backend: cmsis_nn.Backend, conv_node: torch.fx.Node, + nhwc_logical: bool = False, ) -> list[int]: ( input_nhwc, @@ -195,6 +205,7 @@ def cmsis_nn_transpose_conv_buffer_size( dilation_hw, ) = _get_common_conv_buffer_size_inputs( conv_node=conv_node, + nhwc_logical=nhwc_logical, stride_arg_idx=3, padding_arg_idx=4, dilation_arg_idx=6, @@ -248,13 +259,16 @@ def cmsis_nn_transpose_conv_buffer_size( def cmsis_nn_avgpool_buffer_size( backend: cmsis_nn.Backend, pool_node: torch.fx.Node, + nhwc_logical: bool = False, ) -> list[int]: x = cast(torch.fx.Node, pool_node.args[0]) - # Input is NCHW (PyTorch); CMSIS-NN's avgpool buffer sizer only needs the - # input channel count and output width. - _, c_in, _, _ = _shape_from_node(x) - _, _, _, out_w = _shape_from_node(pool_node) + if nhwc_logical: + _, _, _, c_in = _shape_from_node(x) + _, _, out_w, _ = _shape_from_node(pool_node) + else: + _, c_in, _, _ = _shape_from_node(x) + _, _, _, out_w = _shape_from_node(pool_node) return [ int( @@ -270,10 +284,22 @@ def cmsis_nn_avgpool_buffer_size( _target_to_buffer_sizes_registry: dict[Any, BufferSizeFunction] = { exir_ops.edge.cortex_m.quantized_conv2d.default: cmsis_nn_conv_buffer_size, + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: partial( + cmsis_nn_conv_buffer_size, nhwc_logical=True + ), exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default: cmsis_nn_depthwise_conv_buffer_size, + exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default: partial( + cmsis_nn_depthwise_conv_buffer_size, nhwc_logical=True + ), exir_ops.edge.cortex_m.quantized_batch_matmul.default: cmsis_nn_batch_matmul_buffer_size, exir_ops.edge.cortex_m.quantized_transpose_conv2d.default: cmsis_nn_transpose_conv_buffer_size, + exir_ops.edge.cortex_m.quantized_transpose_conv2d_nhwc.default: partial( + cmsis_nn_transpose_conv_buffer_size, nhwc_logical=True + ), exir_ops.edge.cortex_m.quantized_avg_pool2d.default: cmsis_nn_avgpool_buffer_size, + exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default: partial( + cmsis_nn_avgpool_buffer_size, nhwc_logical=True + ), } diff --git a/backends/cortex_m/quantizer/pattern_checkers.py b/backends/cortex_m/quantizer/pattern_checkers.py index ef2a91e6c2c..d1b7f8b8255 100644 --- a/backends/cortex_m/quantizer/pattern_checkers.py +++ b/backends/cortex_m/quantizer/pattern_checkers.py @@ -131,6 +131,18 @@ def check_quantization_config( return is_int8 and is_ch_axis_0 +class CortexMExplicitConv2DCheck(CortexMConv2DCheck): + @classmethod + def check_pattern(cls, pattern): + return all(get_first_fake_tensor(node).dim() == 4 for node in pattern) + + +class CortexMExplicitConv1DCheck(CortexMConv2DCheck): + @classmethod + def check_pattern(cls, pattern): + return all(get_first_fake_tensor(node).dim() == 3 for node in pattern) + + class CortexMLinearCheck(PatternCheck): @classmethod def check_quantization_config( @@ -219,6 +231,8 @@ def check_quantization_config( class CortexMConvTranspose2DCheck(PatternCheck): + require_channels_last = True + @classmethod def _check_node(cls, node: Node) -> bool: if node is None: @@ -228,8 +242,7 @@ def _check_node(cls, node: Node) -> bool: if tensor is None: return False # Reject if no tensor found - # REJECT if using NCHW format (we need channels_last/NHWC) - if not is_channels_last(tensor): + if cls.require_channels_last and not is_channels_last(tensor): return False # Reject NCHW # For aten.conv_transpose2d.input: @@ -288,6 +301,10 @@ def check_quantization_config( return is_int8 and is_ch_axis_1 +class CortexMExplicitConvTranspose2DCheck(CortexMConvTranspose2DCheck): + require_channels_last = False + + class CortexMAvgPool2DCheck(PatternCheck): @classmethod def check_pattern(cls, pattern): diff --git a/backends/cortex_m/quantizer/quantizer.py b/backends/cortex_m/quantizer/quantizer.py index d3f49114144..c1f0778b336 100644 --- a/backends/cortex_m/quantizer/quantizer.py +++ b/backends/cortex_m/quantizer/quantizer.py @@ -25,8 +25,10 @@ ) from executorch.backends.cortex_m.quantizer.quantizer_support import ( __name__ as cortex_m_quantizer_support_module, + CONV1D_OP_PATTERNS, CONV_OP_PATTERNS, CONV_TRANSPOSE_OP_PATTERNS, + CORTEX_M_EXPLICIT_LAYOUT_QUANTIZER_SUPPORT_DICT, CORTEX_M_QUANTIZER_SUPPORT_DICT, ) from executorch.backends.cortex_m.quantizer_reporter import QuantizerReporter @@ -45,8 +47,11 @@ def mark_node_as_annotated( class CortexMQuantizer(ComposableQuantizer): - - def __init__(self, per_tensor_config: Optional[QuantizationConfig] = None) -> None: + def __init__( + self, + per_tensor_config: Optional[QuantizationConfig] = None, + use_explicit_layout: bool = False, + ) -> None: """Cortex-M PT2E quantizer. Args: @@ -57,20 +62,33 @@ def __init__(self, per_tensor_config: Optional[QuantizationConfig] = None) -> No ``INT8_PER_TENSOR_CONFIG``; pass ``INT16_PER_TENSOR_CONFIG`` to quantize the ops that support it (e.g. ``quantized_div``) with int16 activations. + use_explicit_layout: Select the support checks for the experimental + explicit-layout pipeline. Legacy mode continues to require + channels-last convolution inputs. """ per_tensor_config = per_tensor_config or INT8_PER_TENSOR_CONFIG - conv_targets: set[OpOverload] = set() - for key in CONV_OP_PATTERNS.keys() | CONV_TRANSPOSE_OP_PATTERNS.keys(): - conv_targets.update(key) - + support_dict = CORTEX_M_QUANTIZER_SUPPORT_DICT support_dict_name = ( cortex_m_quantizer_support_module + ".CORTEX_M_QUANTIZER_SUPPORT_DICT" ) + conv_patterns = CONV_OP_PATTERNS.keys() | CONV_TRANSPOSE_OP_PATTERNS.keys() + if use_explicit_layout: + support_dict = CORTEX_M_EXPLICIT_LAYOUT_QUANTIZER_SUPPORT_DICT + support_dict_name = ( + cortex_m_quantizer_support_module + + ".CORTEX_M_EXPLICIT_LAYOUT_QUANTIZER_SUPPORT_DICT" + ) + conv_patterns |= CONV1D_OP_PATTERNS.keys() + + conv_targets: set[OpOverload] = set() + for key in conv_patterns: + conv_targets.update(key) + pattern_matcher = PatternMatcher( cast( dict[tuple[OpOverload, ...], Optional[type[PatternCheck]]], - CORTEX_M_QUANTIZER_SUPPORT_DICT, + support_dict, ), support_dict_name=support_dict_name, ) diff --git a/backends/cortex_m/quantizer/quantizer_support.py b/backends/cortex_m/quantizer/quantizer_support.py index d0cce702cff..b00909c3b5f 100644 --- a/backends/cortex_m/quantizer/quantizer_support.py +++ b/backends/cortex_m/quantizer/quantizer_support.py @@ -12,6 +12,9 @@ CortexMConv2DCheck, CortexMConvTranspose2DCheck, CortexMDivCheck, + CortexMExplicitConv1DCheck, + CortexMExplicitConv2DCheck, + CortexMExplicitConvTranspose2DCheck, CortexMLinearCheck, CortexMMaxPool2DCheck, CortexMSoftmaxCheck, @@ -102,6 +105,42 @@ (torch.ops.aten.conv2d.default, torch.ops.aten.clamp_.default): CortexMConv2DCheck, } +CONV1D_OP_PATTERNS = { + (torch.ops.aten.conv1d.default,): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.relu.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.relu_.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.hardtanh.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.hardtanh_.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.hardsigmoid.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.hardsigmoid_.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.clamp.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.clamp_.default, + ): CortexMExplicitConv1DCheck, +} + CONV_TRANSPOSE_OP_PATTERNS = { (torch.ops.aten.conv_transpose2d.input,): CortexMConvTranspose2DCheck, ( @@ -209,3 +248,13 @@ | BMM_OP_PATTERNS | ACTIVATION_OP_PATTERNS ) + +CORTEX_M_EXPLICIT_LAYOUT_QUANTIZER_SUPPORT_DICT = ( + CORTEX_M_QUANTIZER_SUPPORT_DICT + | CONV1D_OP_PATTERNS + | {pattern: CortexMExplicitConv2DCheck for pattern in CONV_OP_PATTERNS} + | { + pattern: CortexMExplicitConvTranspose2DCheck + for pattern in CONV_TRANSPOSE_OP_PATTERNS + } +) diff --git a/backends/cortex_m/test/build_test_runner.sh b/backends/cortex_m/test/build_test_runner.sh index a38c6d53256..39a760ee39f 100755 --- a/backends/cortex_m/test/build_test_runner.sh +++ b/backends/cortex_m/test/build_test_runner.sh @@ -52,6 +52,7 @@ ops_list=( aten::full.out aten::ge.Tensor_out aten::unsqueeze_copy.out + aten::squeeze_copy.dim_out aten::select_copy.int_out aten::amax.out cortex_m::quantize_per_tensor.out @@ -67,10 +68,15 @@ ops_list=( cortex_m::transpose.out cortex_m::pad.out cortex_m::quantized_conv2d.out + cortex_m::quantized_conv2d_nhwc.out cortex_m::quantized_depthwise_conv2d.out + cortex_m::quantized_depthwise_conv2d_nhwc.out cortex_m::quantized_transpose_conv2d.out + cortex_m::quantized_transpose_conv2d_nhwc.out cortex_m::quantized_avg_pool2d.out + cortex_m::quantized_avg_pool2d_nhwc.out cortex_m::quantized_max_pool2d.out + cortex_m::quantized_max_pool2d_nhwc.out cortex_m::quantized_batch_matmul.out ) diff --git a/backends/cortex_m/test/misc/test_portable_int8.py b/backends/cortex_m/test/misc/test_portable_int8.py index 6efeec9e5b0..ce0aeeb1de6 100644 --- a/backends/cortex_m/test/misc/test_portable_int8.py +++ b/backends/cortex_m/test/misc/test_portable_int8.py @@ -283,6 +283,12 @@ def _quantize_and_export( (torch.randn(2, 3, 4, 5), torch.randn(2, 3, 4, 5)), None, ), + "roll": OpCase( + torch.ops.aten.roll.default, + _build_module(lambda x, y: torch.ops.aten.roll.default(x, [1], [2])), + (torch.randn(2, 3, 4, 5), torch.randn(2, 3, 4, 5)), + None, + ), "index_select": OpCase( torch.ops.aten.index_select.default, _build_module( @@ -708,6 +714,7 @@ def _quantize_and_export( "where_default": "MLETORCH-1865: Properly support flaky scalar comparison ops.", "while_loop": "MLETORCH-1866: Support higher-order operators", "cond": "MLETORCH-1866: Support higher-order operators", + "roll": "MLETORCH-2563: Support roll operator", } @@ -716,7 +723,10 @@ def _quantize_and_export( OP_CASES, xfails=xfails, strict=False, - skips={"while_loop": "Has been observed to hang randomly."}, + skips={ + "while_loop": "Has been observed to hang randomly.", + "dropout": "Not training, so it folds away and no node survives to carry int8.", + }, ) def test_shared_qspec_portable_int8_ops(op_case: OpCase) -> None: tester = CortexMTester(op_case.module, op_case.example_inputs) diff --git a/backends/cortex_m/test/models/test_ds_cnn.py b/backends/cortex_m/test/models/test_ds_cnn.py index 206af19a61e..ba5b2da6b78 100644 --- a/backends/cortex_m/test/models/test_ds_cnn.py +++ b/backends/cortex_m/test/models/test_ds_cnn.py @@ -15,7 +15,6 @@ "executorch_exir_dialects_edge__ops_aten_linear_default": 1, "executorch_exir_dialects_edge__ops_aten_relu_default": 9, "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 2, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_channel_default": 18, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 17, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 15, @@ -23,14 +22,13 @@ ops_after_transforms: dict[str, int] = { "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 3, "executorch_exir_dialects_edge__ops_cortex_m_pad_default": 1, - "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 3, "executorch_exir_dialects_edge__ops_cortex_m_quantized_avg_pool2d_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 4, "executorch_exir_dialects_edge__ops_cortex_m_quantized_depthwise_conv2d_default": 5, "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 2, } test_cases = { @@ -43,11 +41,21 @@ } +ops_absent_after_transforms: list[str] = [ + "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default", +] + + @parametrize("test_case", test_cases) def test_dialect_ds_cnn(test_case): inputs = test_case.get_example_inputs() tester = CortexMTester(test_case.model, inputs) - tester.test_dialect(ops_before_transforms, ops_after_transforms, qtol=1) + tester.test_dialect( + ops_before_transforms, + ops_after_transforms, + qtol=1, + ops_absent_after_transforms=ops_absent_after_transforms, + ) @parametrize("test_case", test_cases) diff --git a/backends/cortex_m/test/models/test_mobilenet_v2.py b/backends/cortex_m/test/models/test_mobilenet_v2.py index 67f0937a006..7a3bc9ecfe3 100644 --- a/backends/cortex_m/test/models/test_mobilenet_v2.py +++ b/backends/cortex_m/test/models/test_mobilenet_v2.py @@ -20,7 +20,6 @@ "executorch_exir_dialects_edge__ops_aten_hardtanh_default": 35, "executorch_exir_dialects_edge__ops_aten_linear_default": 1, "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_channel_default": 104, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 79, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 67, @@ -28,21 +27,17 @@ ops_after_transforms: dict[str, int] = { "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, - "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 2, + "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 2, "executorch_exir_dialects_edge__ops_cortex_m_quantized_add_default": 10, "executorch_exir_dialects_edge__ops_cortex_m_quantized_avg_pool2d_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 35, "executorch_exir_dialects_edge__ops_cortex_m_quantized_depthwise_conv2d_default": 17, "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, } # Use larger sample set for calibration to get better quantization -calibration_samples = [ - (torch.randn(1, 3, 224, 224).to(memory_format=torch.channels_last),) - for _ in range(100) -] +calibration_samples = [(torch.randn(1, 3, 224, 224),) for _ in range(100)] test_cases = { "mobilenet_v2": McuTestCase( @@ -54,6 +49,11 @@ } +ops_absent_after_transforms: list[str] = [ + "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default", +] + + @parametrize("test_case", test_cases) def test_dialect_mv2(test_case): inputs = test_case.get_example_inputs() @@ -63,6 +63,7 @@ def test_dialect_mv2(test_case): ops_after_transforms, qtol=10, calibration_samples=calibration_samples, + ops_absent_after_transforms=ops_absent_after_transforms, ) # assert that top 1 output matches diff --git a/backends/cortex_m/test/models/test_mobilenet_v3.py b/backends/cortex_m/test/models/test_mobilenet_v3.py index 08633d54dd6..3a6f0a5004f 100644 --- a/backends/cortex_m/test/models/test_mobilenet_v3.py +++ b/backends/cortex_m/test/models/test_mobilenet_v3.py @@ -39,10 +39,7 @@ } # Use bigger sample set for calibration. -calibration_samples = [ - (torch.randn(1, 3, 232, 232).to(memory_format=torch.channels_last),) - for i in (range(100)) -] +calibration_samples = [(torch.randn(1, 3, 232, 232),) for _ in range(100)] test_cases = { "mobilenet_v3_small": McuTestCase( diff --git a/backends/cortex_m/test/ops/nhwc_test_utils.py b/backends/cortex_m/test/ops/nhwc_test_utils.py new file mode 100644 index 00000000000..6cc02d85602 --- /dev/null +++ b/backends/cortex_m/test/ops/nhwc_test_utils.py @@ -0,0 +1,42 @@ +# 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. + +import math + +import torch + +from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager +from executorch.backends.cortex_m.passes.scratch_buffer_sizes import ( + required_cmsis_nn_buffer_sizes, +) +from executorch.backends.cortex_m.test.tester import CortexMTester +from executorch.backends.test.harness.stages import RunPasses, StageType + + +def int8_values(shape): + values = torch.arange(math.prod(shape), dtype=torch.int32) + return (values.remainder(7) - 3).to(torch.int8).reshape(shape) + + +def run_on_fvp(module, x, target, target_config, scratch_count, atol=0): + sizing_inputs = (x,) + tuple( + torch.empty(0, dtype=torch.uint8) for _ in range(scratch_count) + ) + tester = CortexMTester(module, sizing_inputs, target_config=target_config) + tester.export().to_edge() + program = tester.get_artifact(StageType.TO_EDGE).exported_program() + [node] = [n for n in program.graph.nodes if n.target == target] + scratch_sizes = required_cmsis_nn_buffer_sizes(node, target_config.backend) or [] + assert len(scratch_sizes) == scratch_count + + inputs = (x,) + tuple( + torch.empty(size, dtype=torch.uint8) for size in scratch_sizes + ) + tester = CortexMTester(module, inputs, target_config=target_config) + tester.export().to_edge() + tester.run_passes(RunPasses(CortexMPassManager, pass_list=[])) + tester.to_executorch().serialize() + tester.run_method_and_compare_outputs(inputs=inputs, atol=atol) diff --git a/backends/cortex_m/test/ops/test_nhwc_conv.py b/backends/cortex_m/test/ops/test_nhwc_conv.py new file mode 100644 index 00000000000..da7602447c2 --- /dev/null +++ b/backends/cortex_m/test/ops/test_nhwc_conv.py @@ -0,0 +1,145 @@ +# 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. + +import torch + +from executorch.backends.cortex_m.test.ops.nhwc_test_utils import ( + int8_values, + run_on_fvp, +) +from executorch.exir.dialects._ops import ops as exir_ops + +# Direct-op coverage is temporary until CortexMTester uses explicit layout by default. + + +class Conv2dNhwc(torch.nn.Module): + def __init__(self, grouped=False): + super().__init__() + in_channels = 4 if grouped else 3 + self.register_buffer( + "weight", int8_values((4, 2, 3, 2 if grouped else in_channels)) + ) + self.register_buffer("bias", torch.arange(4, dtype=torch.int32) - 2) + self.register_buffer( + "multipliers", torch.full((4,), 1 << 30, dtype=torch.int32) + ) + self.register_buffer("shifts", torch.full((4,), -1, dtype=torch.int32)) + + def forward(self, x, scratch): + return torch.ops.cortex_m.quantized_conv2d_nhwc.default( + x, + self.weight, + self.bias, + [2, 1], + [1, 0], + [1, 1], + 0, + 0, + self.multipliers, + self.shifts, + -128, + 127, + scratch, + ) + + +class DepthwiseConv2dNhwc(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", int8_values((1, 3, 2, 4))) + self.register_buffer("bias", torch.arange(4, dtype=torch.int32) - 2) + self.register_buffer( + "multipliers", torch.full((4,), 1 << 30, dtype=torch.int32) + ) + self.register_buffer("shifts", torch.full((4,), -1, dtype=torch.int32)) + + def forward(self, x, scratch): + return torch.ops.cortex_m.quantized_depthwise_conv2d_nhwc.default( + x, + self.weight, + self.bias, + [2, 1], + [1, 0], + [1, 1], + 1, + 0, + 0, + self.multipliers, + self.shifts, + -128, + 127, + scratch, + ) + + +class TransposeConv2dNhwc(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", int8_values((4, 2, 4, 2))) + self.register_buffer("bias", torch.arange(4, dtype=torch.int32) - 2) + self.register_buffer( + "multipliers", torch.full((4,), 1 << 30, dtype=torch.int32) + ) + self.register_buffer("shifts", torch.full((4,), -1, dtype=torch.int32)) + + def forward(self, x, scratch, output_scratch): + return torch.ops.cortex_m.quantized_transpose_conv2d_nhwc.default( + x, + self.weight, + self.bias, + [1, 1], + [0, 0], + [0, 0], + [1, 1], + 0, + 0, + self.multipliers, + self.shifts, + -128, + 127, + scratch, + output_scratch, + ) + + +def test_conv2d_nhwc_runs_on_fvp(cortex_m_target): + run_on_fvp( + Conv2dNhwc(), + int8_values((1, 7, 10, 3)), + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + cortex_m_target, + 1, + ) + + +def test_grouped_conv2d_nhwc_runs_on_fvp(cortex_m_target): + run_on_fvp( + Conv2dNhwc(grouped=True), + int8_values((1, 7, 10, 4)), + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + cortex_m_target, + 1, + ) + + +def test_depthwise_conv2d_nhwc_runs_on_fvp(cortex_m_target): + run_on_fvp( + DepthwiseConv2dNhwc(), + int8_values((1, 7, 10, 4)), + exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default, + cortex_m_target, + 1, + ) + + +def test_transpose_conv2d_nhwc_runs_on_fvp(cortex_m_target): + run_on_fvp( + TransposeConv2dNhwc(), + int8_values((1, 5, 6, 2)), + exir_ops.edge.cortex_m.quantized_transpose_conv2d_nhwc.default, + cortex_m_target, + 2, + ) diff --git a/backends/cortex_m/test/ops/test_nhwc_pool.py b/backends/cortex_m/test/ops/test_nhwc_pool.py new file mode 100644 index 00000000000..4377df62f9f --- /dev/null +++ b/backends/cortex_m/test/ops/test_nhwc_pool.py @@ -0,0 +1,67 @@ +# 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. + +import torch + +from executorch.backends.cortex_m.test.ops.nhwc_test_utils import ( + int8_values, + run_on_fvp, +) +from executorch.exir.dialects._ops import ops as exir_ops + +# Direct-op coverage is temporary until CortexMTester uses explicit layout by default. + + +class AvgPool2dNhwc(torch.nn.Module): + def forward(self, x, scratch): + return torch.ops.cortex_m.quantized_avg_pool2d_nhwc.default( + x, + [2, 3], + [2, 1], + [0, 1], + False, + 0, + 1 << 30, + 1, + scratch, + ) + + +class MaxPool2dNhwc(torch.nn.Module): + def forward(self, x): + return torch.ops.cortex_m.quantized_max_pool2d_nhwc.default( + x, + [2, 3], + [2, 1], + [0, 1], + [1, 1], + False, + 0, + 0, + -128, + 127, + ) + + +def test_avg_pool2d_nhwc_runs_on_fvp(cortex_m_target): + run_on_fvp( + AvgPool2dNhwc(), + int8_values((1, 7, 9, 3)), + exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default, + cortex_m_target, + 1, + atol=1, + ) + + +def test_max_pool2d_nhwc_runs_on_fvp(cortex_m_target): + run_on_fvp( + MaxPool2dNhwc(), + int8_values((1, 7, 9, 3)), + exir_ops.edge.cortex_m.quantized_max_pool2d_nhwc.default, + cortex_m_target, + 0, + ) diff --git a/backends/cortex_m/test/ops/test_transpose.py b/backends/cortex_m/test/ops/test_transpose.py index 2e5f5112bd9..26a026358df 100644 --- a/backends/cortex_m/test/ops/test_transpose.py +++ b/backends/cortex_m/test/ops/test_transpose.py @@ -24,6 +24,12 @@ "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, } +RANK5_OPS_AFTER_PASSES = { + "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, +} + class CortexMPermute(torch.nn.Module): ops_before_transforms = OPS_BEFORE_PASSES @@ -98,6 +104,19 @@ def test_dialect_transpose(test_case, cortex_m_target): ) +def test_dialect_rank5_permute_stays_portable(cortex_m_target): + tester = CortexMTester( + CortexMPermute((0, 2, 1, 4, 3)), + (ramp_tensor(-1.0, 1.0, (1, 2, 3, 4, 5)),), + target_config=cortex_m_target, + ) + tester.test_dialect( + OPS_BEFORE_PASSES, + RANK5_OPS_AFTER_PASSES, + qtol=1, + ) + + @parametrize("test_case", test_cases) def test_implementation_transpose(test_case, cortex_m_target): tester = CortexMTester( diff --git a/backends/cortex_m/test/targets.bzl b/backends/cortex_m/test/targets.bzl index 2ea3a5b3b99..06b2daed033 100644 --- a/backends/cortex_m/test/targets.bzl +++ b/backends/cortex_m/test/targets.bzl @@ -5,6 +5,8 @@ # LICENSE file in the root directory of this source tree. load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") +load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbcode_macros//build_defs:python_pytest.bzl", "python_pytest") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbsource//tools/build_defs:platform_defs.bzl", "CXX") @@ -36,6 +38,21 @@ def define_common_targets(is_fbcode = False): define_operator_test_target(op) if is_fbcode: + python_library( + name = "tester", + srcs = ["tester.py"], + deps = [ + "//caffe2:torch", + "//executorch/backends/arm/test:arm_tester", + "//executorch/backends/arm/test:common", + "//executorch/backends/cortex_m:edge_compile_config", + "//executorch/backends/cortex_m:target_config", + "//executorch/backends/cortex_m/passes:cortex_passes", + "//executorch/backends/cortex_m/quantizer:quantizer", + "//executorch/backends/test/harness:tester", + ], + ) + python_unittest( name = "test_activation_lut", srcs = [ @@ -50,6 +67,24 @@ def define_common_targets(is_fbcode = False): ], ) + python_pytest( + name = "test_explicit_layout_pipeline", + srcs = ["test_explicit_layout_pipeline.py"], + compile = "with-source", + typing = False, + deps = [ + "//caffe2:torch", + "//executorch/backends/cortex_m:target_config", + "//executorch/backends/cortex_m/ops:ops", + "//executorch/backends/cortex_m/passes:cortex_passes", + "//executorch/backends/cortex_m/quantizer:quantizer", + "//executorch/backends/test/harness:tester", + "//executorch/exir/dialects:lib", + ":tester", + "fbsource//third-party/pypi/pytest:pytest", + ], + ) + python_unittest( name = "test_replace_quant_nodes", srcs = [ diff --git a/backends/cortex_m/test/test_explicit_layout.py b/backends/cortex_m/test/test_explicit_layout.py new file mode 100644 index 00000000000..1e9ae4488d8 --- /dev/null +++ b/backends/cortex_m/test/test_explicit_layout.py @@ -0,0 +1,180 @@ +# 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. + +from types import SimpleNamespace + +import pytest +import torch +from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig +from executorch.backends.cortex_m.test.tester import CortexMSerialize +from executorch.exir.dialects._ops import ops as exir_ops + + +_LEGACY_SPATIAL_OPS = { + exir_ops.edge.cortex_m.quantized_conv2d.default, + exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default, + exir_ops.edge.cortex_m.quantized_transpose_conv2d.default, + exir_ops.edge.cortex_m.quantized_avg_pool2d.default, + exir_ops.edge.cortex_m.quantized_max_pool2d.default, +} + +# Temporary dual-mode acceptance coverage; see "Explicit-layout migration" in +# the Cortex-M README. + + +class Conv2d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 4, 3, padding=1) + + def forward(self, x): + return self.conv(x) + + +class Conv1d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv1d(2, 4, 3, padding=1) + + def forward(self, x): + return self.conv(x) + + +def _compile(module, inputs, *, explicit_layout: bool, quantize: bool = True): + from executorch.backends.arm.scripts.aot_arm_compiler import _to_edge_cortex_m + + exported_program = torch.export.export(module, inputs, strict=True) + return _to_edge_cortex_m( + exported_program, + SimpleNamespace( + cortex_m_explicit_layout=explicit_layout, + quantize=quantize, + strict_export=True, + ), + exported_program.module(), + inputs, + None, + CortexMTargetConfig(cpu=CortexM.M55), + ) + + +def _count(exported_program, target) -> int: + return sum(node.target == target for node in exported_program.graph.nodes) + + +@pytest.mark.parametrize( + "explicit_layout,expected,unexpected", + [ + ( + False, + exir_ops.edge.cortex_m.quantized_conv2d.default, + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + ), + ( + True, + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + exir_ops.edge.cortex_m.quantized_conv2d.default, + ), + ], +) +def test_aot_layout_mode_selects_operator_family(explicit_layout, expected, unexpected): + _, edge, _ = _compile( + Conv2d().eval(), + (torch.randn(1, 3, 8, 8),), + explicit_layout=explicit_layout, + ) + program = edge.exported_program() + + assert _count(program, expected) == 1 + assert _count(program, unexpected) == 0 + + +def test_aot_explicit_layout_requires_quantization(): + with pytest.raises(RuntimeError, match="requires --quantize"): + _compile( + Conv2d().eval(), + (torch.randn(1, 3, 8, 8),), + explicit_layout=True, + quantize=False, + ) + + +def _assert_explicit_copy_ceiling(module, inputs, ceiling): + _, edge, _ = _compile(module, inputs, explicit_layout=True) + program = edge.exported_program() + copies = [ + node + for node in program.graph.nodes + if node.target + in { + exir_ops.edge.cortex_m.transpose.default, + exir_ops.edge.aten.view_copy.default, + } + ] + + assert len(copies) <= ceiling + assert all( + node.args[0].meta["val"].dtype == torch.int8 + for node in copies + if node.target == exir_ops.edge.cortex_m.transpose.default + ) + assert not any(node.target in _LEGACY_SPATIAL_OPS for node in program.graph.nodes) + assert not any( + getattr(node.target, "namespace", None) == "channels_last" + for node in program.graph.nodes + ) + + +def test_mobilenet_v2_explicit_copy_ceiling(): + torchvision = pytest.importorskip("torchvision") + _assert_explicit_copy_ceiling( + torchvision.models.mobilenet_v2(weights=None).eval(), + (torch.randn(1, 3, 224, 224),), + ceiling=2, + ) + + +def test_resnet8_explicit_copy_ceiling(): + from executorch.examples.models.mlperf_tiny.resnet8 import ResNet8 + + _assert_explicit_copy_ceiling( + ResNet8().eval(), + (torch.rand(1, 3, 32, 32) * 2 - 1,), + ceiling=2, + ) + + +def test_silero_explicit_copy_ceiling(): + from executorch.examples.models.silero_vad.export_silero_vad import ( + CONTEXT_SIZE, + HIDDEN_DIM, + SileroVAD16k, + WINDOW_SIZE, + ) + + _assert_explicit_copy_ceiling( + SileroVAD16k().eval(), + ( + torch.randn(1, CONTEXT_SIZE + WINDOW_SIZE), + torch.zeros(2, 1, HIDDEN_DIM), + ), + ceiling=12, + ) + + +def test_explicit_conv1d_runs_on_fvp(): + inputs = (torch.linspace(-5, 5, steps=16).reshape(1, 2, 8),) + model_quant, edge, runtime_inputs = _compile( + Conv1d().eval(), inputs, explicit_layout=True + ) + program = edge.exported_program() + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + + serialized = CortexMSerialize(CortexMTargetConfig(cpu=CortexM.M55)) + serialized.run(edge.to_executorch()) + [actual] = serialized.run_artifact(runtime_inputs) + expected = model_quant(*runtime_inputs) + torch.testing.assert_close(actual, expected, atol=0.05, rtol=1e-3) diff --git a/backends/cortex_m/test/test_explicit_layout_pipeline.py b/backends/cortex_m/test/test_explicit_layout_pipeline.py new file mode 100644 index 00000000000..7c93233e768 --- /dev/null +++ b/backends/cortex_m/test/test_explicit_layout_pipeline.py @@ -0,0 +1,158 @@ +# 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. + +from functools import partial + +import pytest +import torch +from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager +from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer +from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig +from executorch.backends.cortex_m.test.tester import CortexMTester +from executorch.backends.test.harness.stages import Quantize, RunPasses, StageType +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import Node + +# Temporary opt-in coverage. Move these invariants into the standard Cortex-M +# tests when this pass manager becomes the default. + + +class Conv2d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 4, 3, padding=1) + + def forward(self, x): + return self.conv(x) + + +class Conv1d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv1d(2, 4, 3, padding=1) + + def forward(self, x): + return self.conv(x) + + +class ConvPadConv(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 4, 3, padding=1) + self.conv2 = torch.nn.Conv2d(4, 5, 3, padding=1) + + def forward(self, x): + return self.conv2(torch.nn.functional.pad(self.conv1(x), (1, 1, 1, 1))) + + +class UnsupportedAvgPool(torch.nn.Module): + def forward(self, x): + return torch.nn.functional.avg_pool2d( + x, kernel_size=2, stride=2, divisor_override=3 + ) + + +def _count(exported_program, target) -> int: + return sum(node.target == target for node in exported_program.graph.nodes) + + +def _run_explicit_layout_pass_manager(tester: CortexMTester) -> CortexMTester: + target_config = CortexMTargetConfig(cpu=CortexM.M55) + tester.run_passes( + RunPasses( + partial( + CortexMPassManager, + target_config=target_config, + use_explicit_layout=True, + ), # type: ignore[arg-type] + CortexMPassManager.explicit_layout_pass_list, # type: ignore[arg-type] + ) + ) + return tester + + +def _run_explicit_layout_passes(tester: CortexMTester) -> CortexMTester: + tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=True))) + tester.export().to_edge() + return _run_explicit_layout_pass_manager(tester) + + +def test_layout_pipelines_select_distinct_spatial_operators(): + legacy_input = torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last) + legacy = CortexMTester( + Conv2d().eval().to(memory_format=torch.channels_last), + (legacy_input,), + ) + legacy.quantize().export().to_edge().run_passes() + legacy_program = legacy.get_artifact(StageType.RUN_PASSES).exported_program() + + explicit = _run_explicit_layout_passes( + CortexMTester(Conv2d().eval(), (torch.randn(1, 3, 8, 8),)) + ) + explicit_program = explicit.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(legacy_program, exir_ops.edge.cortex_m.quantized_conv2d.default) == 1 + assert ( + _count( + legacy_program, + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + ) + == 0 + ) + assert ( + _count(explicit_program, exir_ops.edge.cortex_m.quantized_conv2d.default) == 0 + ) + assert ( + _count( + explicit_program, + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + ) + == 1 + ) + assert _count(explicit_program, exir_ops.edge.cortex_m.transpose.default) == 2 + + +def test_conv1d_is_quantized_before_layout_conversion(): + tester = CortexMTester(Conv1d().eval(), (torch.randn(1, 2, 8),)) + tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=True))) + quantized = tester.get_artifact(StageType.QUANTIZE) + [conv1d] = [ + node + for node in quantized.graph.nodes + if node.target == torch.ops.aten.conv1d.default + ] + + weight = conv1d.args[1] + assert isinstance(weight, Node) + assert ( + weight.target == torch.ops.quantized_decomposed.dequantize_per_channel.default + ) + + tester.export().to_edge() + _run_explicit_layout_pass_manager(tester) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + assert _count(program, exir_ops.edge.aten.convolution.default) == 0 + + +def test_explicit_layout_reuses_pad(): + tester = _run_explicit_layout_passes( + CortexMTester(ConvPadConv().eval(), (torch.randn(1, 3, 8, 8),)) + ) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.pad.default) == 1 + + +def test_explicit_layout_rejects_unsupported_spatial_operator(): + tester = CortexMTester(UnsupportedAvgPool(), (torch.randn(1, 3, 8, 8),)) + + with pytest.raises(Exception) as caught: + _run_explicit_layout_passes(tester) + + assert caught.value.__cause__ is not None + assert "NHWC-eligible" in str(caught.value.__cause__) diff --git a/backends/cortex_m/test/tester.py b/backends/cortex_m/test/tester.py index a1b5245b80b..b644db4e6c1 100644 --- a/backends/cortex_m/test/tester.py +++ b/backends/cortex_m/test/tester.py @@ -132,6 +132,7 @@ def test_dialect( qtol=0, atol=1e-03, calibration_samples=None, + ops_absent_after_transforms=None, ): """ Test the python dialect op implementation. @@ -142,13 +143,14 @@ def test_dialect( ) else: quantization_stage = None - self.quantize(quantization_stage) self.export() self.to_edge() self.check_count(ops_before_transforms) self.run_passes() self.check_count(ops_after_transforms) + if ops_absent_after_transforms: + self.check_not(ops_absent_after_transforms) self.run_method_and_compare_outputs( inputs=self.example_inputs, qtol=qtol, atol=atol ) diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 3fe33473472..4602978fad4 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -109,6 +109,22 @@ fbcode_target( ], ) +fbcode_target( + _kind = runtime.python_library, + name = "merge_ptes", + srcs = [ + "merge_ptes.py", + ], + visibility = ["PUBLIC"], + deps = [ + ":cuda_backend", + "//caffe2:torch", + "//executorch/exir/_serialize:lib", + "//executorch/exir:schema", + "//executorch/extension/flat_tensor/serialize:serialize", + ], +) + fbcode_target( _kind = runtime.python_library, name = "cuda_partitioner", diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index c3e5d17809e..6ed8f83678f 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -14,7 +14,7 @@ # ~~~ # It should also be cmake-lint clean. # -cmake_minimum_required(VERSION 3.29) +cmake_minimum_required(VERSION 3.26) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 854ebf8f952..4db99b6b040 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -519,6 +519,14 @@ def _setup_cuda_environment_for_fatbin() -> bool: except Exception: return False + @classmethod + def _should_include_ptx(cls, compile_specs: List[CompileSpec]) -> bool: + include_ptx = True + for spec in compile_specs: + if spec.key == "cuda_include_ptx": + include_ptx = _on_off_compile_spec_value(spec) + return include_ptx and cls._setup_cuda_environment_for_fatbin() + @classmethod def save_data_externally(cls) -> bool: """ @@ -539,7 +547,29 @@ def _preprocess_with_weight_collector( result = super().preprocess(edge_program, compile_specs) if capture.artifact is None: raise RuntimeError("CUDA AOTI did not return a structured Weights output") - collector.add_preprocess_result(result, capture.artifact, cls.get_device_name()) + target_sm = None + ptx_compute = 0 + if torch.version.hip is None: + from torch._inductor.codegen.cuda.compile_utils import ( + _nvcc_arch_as_compile_option, + ) + + compiled_arch = _nvcc_arch_as_compile_option() + target_arch = compiled_arch.removesuffix("a") + if not target_arch.isdigit(): + raise RuntimeError(f"Unsupported CUDA architecture: {compiled_arch}") + target_sm = int(target_arch) + # Architecture-accelerated PTX targets such as compute_90a and + # compute_120a are not forward-compatible fallback images. + if cls._should_include_ptx(compile_specs) and compiled_arch.isdigit(): + ptx_compute = int(compiled_arch) + collector.add_preprocess_result( + result, + capture.artifact, + cls.get_device_name(), + target_sm=target_sm, + ptx_compute=ptx_compute, + ) return result @classmethod @@ -669,7 +699,10 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]: return {} return { "at::_ops::_weight_int4pack_mm::call": None, + # Also under the shim name Inductor derives for it. + "aoti_torch_cuda__weight_int4pack_mm": None, "at::_ops::sort_stable::call": None, + "aoti_torch_cuda_sort_stable": None, "aoti_torch_cuda_randint_low_out": None, "executorch_cuda::int4_plain_mm": None, "aoti_torch_cuda_int4_plain_mm": None, @@ -759,7 +792,7 @@ def get_aoti_compile_options( # Configure CUDA environment variables based on detected version - emit_multi_arch_kernel = CudaBackend._setup_cuda_environment_for_fatbin() + emit_multi_arch_kernel = cls._should_include_ptx(compile_specs) # Base options for all platforms options: Dict[str, typing.Any] = { diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index ac47f7c7e2c..08806c2bc68 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -22,7 +22,7 @@ from executorch.exir.tensor import scalar_type_enum -CUDA_WEIGHT_CACHE_MAGIC = b"ETCUDAFQN3" +CUDA_AOTI_METADATA_MAGIC = b"ETCUDAFQN0" AOTI_DEVICE_TYPE_CPU = 0 AOTI_DEVICE_TYPE_CUDA = 1 @@ -46,6 +46,24 @@ class CudaWeightArtifact: storages: Dict[str, FileBackedData] +@dataclass(frozen=True) +class CudaAotiVariant: + """One AOTI shared library and its CUDA runtime-selection metadata.""" + + target_sm: int + ptx_compute: int + so_blob_key: str + fallback_only: bool = False + + +@dataclass(frozen=True) +class CudaAotiMetadata: + """CUDA AOTI native and fallback variants sharing one weight manifest.""" + + variants: List[CudaAotiVariant] + entries: List[CudaWeightEntry] + + @dataclass class _CudaWeightCapture: collector: "CudaWeightCollector" @@ -114,18 +132,75 @@ def _is_aoti_library_local_fqn(fqn: str) -> bool: return fqn.startswith("_tensor_constant") -def encode_cuda_weight_metadata( - so_blob_key: str, entries: List[CudaWeightEntry] +def _validate_cuda_aoti_variant( + variant: CudaAotiVariant, has_fallback: bool, regular_sms: set[int] +) -> None: + if variant.target_sm < 0: + raise ValueError(f"Invalid CUDA target SM: {variant.target_sm}") + if variant.ptx_compute < 0 or variant.ptx_compute > variant.target_sm: + raise ValueError( + f"Invalid PTX compute target {variant.ptx_compute} for sm{variant.target_sm}" + ) + if not variant.so_blob_key: + raise ValueError("CUDA AOTI variant is missing its shared-object key") + if variant.fallback_only: + if variant.ptx_compute == 0: + raise ValueError("CUDA fallback variant must contain PTX") + return + if variant.target_sm in regular_sms: + raise ValueError(f"Duplicate CUDA target SM: {variant.target_sm}") + regular_sms.add(variant.target_sm) + if has_fallback and variant.ptx_compute != 0: + raise ValueError("Regular CUDA variants cannot advertise PTX fallback") + + +def _validate_cuda_aoti_variants( + variants: List[CudaAotiVariant], has_fallback: bool +) -> None: + untargeted = [variant for variant in variants if variant.target_sm == 0] + if untargeted: + if len(variants) != 1 or untargeted[0].ptx_compute or has_fallback: + raise ValueError( + "Untargeted CUDA AOTI metadata requires one non-fallback variant" + ) + if not untargeted[0].so_blob_key: + raise ValueError("CUDA AOTI variant is missing its shared-object key") + return + if sum(variant.fallback_only for variant in variants) > 1: + raise ValueError("CUDA AOTI metadata supports only one fallback variant") + if len(variants) > 1 and any( + variant.ptx_compute and not variant.fallback_only for variant in variants + ): + raise ValueError( + "Multi-variant CUDA AOTI metadata requires an explicit PTX fallback" + ) + regular_sms: set[int] = set() + for variant in variants: + _validate_cuda_aoti_variant(variant, has_fallback, regular_sms) + + +def encode_cuda_aoti_metadata( + variants: List[CudaAotiVariant], entries: List[CudaWeightEntry] ) -> bytes: - """Encode the per-method FQN-to-tensor metadata consumed by CUDA runtime.""" - output = bytearray(CUDA_WEIGHT_CACHE_MAGIC) + """Encode CUDA AOTI variants followed by one shared weight manifest.""" + if not variants: + raise ValueError("CUDA AOTI metadata requires at least one variant") + + has_fallback = any(variant.fallback_only for variant in variants) + _validate_cuda_aoti_variants(variants, has_fallback) + output = bytearray(CUDA_AOTI_METADATA_MAGIC) def write_string(value: str) -> None: encoded = value.encode("utf-8") output.extend(struct.pack(" None: return bytes(output) +class _MetadataReader: + def __init__(self, data: bytes) -> None: + self._data = memoryview(data) + self._offset = 0 + + def read(self, size: int) -> memoryview: + end = self._offset + size + if size < 0 or end > len(self._data): + raise ValueError("Truncated CUDA AOTI metadata") + value = self._data[self._offset : end] + self._offset = end + return value + + def unpack(self, format: str) -> Tuple[Any, ...]: + size = struct.calcsize(format) + return struct.unpack(format, self.read(size)) + + def read_string(self) -> str: + (size,) = self.unpack(" None: + if self._offset != len(self._data): + raise ValueError("CUDA AOTI metadata contains trailing bytes") + + +def _decode_weight_entries(reader: _MetadataReader) -> List[CudaWeightEntry]: + (num_entries,) = reader.unpack(" 1 << 20: + raise ValueError(f"CUDA AOTI metadata has too many weights: {num_entries}") + + entries = [] + for _ in range(num_entries): + fqn = reader.read_string() + storage_key = reader.read_string() + storage_nbytes, dtype, device_type, storage_offset, ndim = reader.unpack( + " 64: + raise ValueError("CUDA AOTI metadata contains an invalid weight entry") + sizes = reader.unpack(f"<{ndim}q") if ndim else () + strides = reader.unpack(f"<{ndim}q") if ndim else () + if storage_offset < 0 or any(value < 0 for value in sizes + strides): + raise ValueError("CUDA AOTI metadata contains invalid tensor metadata") + entries.append( + CudaWeightEntry( + fqn=fqn, + storage_key=storage_key, + storage_nbytes=storage_nbytes, + dtype=dtype, + device_type=device_type, + storage_offset=storage_offset, + sizes=tuple(sizes), + strides=tuple(strides), + ) + ) + return entries + + +def decode_cuda_aoti_metadata(data: bytes) -> CudaAotiMetadata: + """Decode CUDA AOTI variant and shared-weight metadata.""" + reader = _MetadataReader(data) + magic = bytes(reader.read(len(CUDA_AOTI_METADATA_MAGIC))) + if magic != CUDA_AOTI_METADATA_MAGIC: + raise ValueError("Unrecognized CUDA AOTI metadata") + + variants = [] + (num_variants,) = reader.unpack(" 256: + raise ValueError( + f"CUDA AOTI metadata has invalid variant count: {num_variants}" + ) + for _ in range(num_variants): + target_sm, ptx_compute, flags = reader.unpack(" value`` store for all methods.""" @@ -311,6 +482,8 @@ def add_preprocess_result( result: PreprocessResult, artifact: CudaWeightArtifact, device_name: str, + target_sm: Optional[int] = None, + ptx_compute: int = 0, ) -> None: if result.data_store_output is None: raise RuntimeError("CUDA AOTI preprocess returned no named data") @@ -344,8 +517,9 @@ def add_preprocess_result( self._add_weight(entry, data, external_tag) serialized_entries.append(entry) - result.processed_bytes = encode_cuda_weight_metadata( - so_blob_key, serialized_entries + result.processed_bytes = encode_cuda_aoti_metadata( + [CudaAotiVariant(target_sm or 0, ptx_compute, so_blob_key)], + serialized_entries, ) self._results.append(result) diff --git a/backends/cuda/merge_ptes.py b/backends/cuda/merge_ptes.py new file mode 100644 index 00000000000..2a7212e608a --- /dev/null +++ b/backends/cuda/merge_ptes.py @@ -0,0 +1,675 @@ +# 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. + +import argparse +import copy +import hashlib +import os +import shutil +import tempfile +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Dict, Iterable, List, Optional, Sequence, Tuple + +import torch + +from executorch.backends.cuda.cuda_weight_collector import ( + CudaAotiMetadata, + CudaAotiVariant, + CudaWeightEntry, + decode_cuda_aoti_metadata, + encode_cuda_aoti_metadata, +) +from executorch.exir._serialize._cord import Cord +from executorch.exir._serialize._named_data_store import ( + NamedDataStore, + NamedDataStoreOutput, +) +from executorch.exir._serialize._program import ( + deserialize_pte_binary, + PTEFile, + serialize_pte_binary, +) +from executorch.exir.schema import ( + BackendDelegateDataReference, + BackendDelegateInlineData, + DataLocation, + Program, +) +from executorch.extension.flat_tensor.serialize.serialize import ( + _deserialize_to_flat_tensor, + FlatTensorHeader, +) + + +CUDA_BACKEND_ID = "CudaBackend" + + +@dataclass(frozen=True) +class CudaPteInput: + """One native CUDA export and its optional external tensor data.""" + + pte_path: Path + ptd_path: Optional[Path] + + +@dataclass(frozen=True) +class CudaPteProvenance: + delegate: Tuple[str, int] + kind: str + target_sm: int + ptx_compute: int + source_pte: Path + + +@dataclass(frozen=True) +class CudaPteMergeResult: + pte: Cord + provenance: Tuple[CudaPteProvenance, ...] + + +@dataclass(frozen=True) +class _PtdBlob: + offset: int + size: int + + +class _PtdIndex: + def __init__(self, path: Path) -> None: + self.path = path + with path.open("rb") as source: + prefix = source.read(8 + FlatTensorHeader.EXPECTED_LENGTH) + header = FlatTensorHeader.from_bytes(prefix[8:]) + if not header.is_valid(): + raise ValueError(f"Invalid PTD header in {path}") + source.seek(0) + flatbuffer_size = header.flatbuffer_offset + header.flatbuffer_size + flat_tensor = _deserialize_to_flat_tensor(source.read(flatbuffer_size)) + + file_size = path.stat().st_size + self._blobs: Dict[str, _PtdBlob] = {} + for named_data in flat_tensor.named_data: + if named_data.key in self._blobs: + raise ValueError(f"PTD contains duplicate key {named_data.key!r}") + if named_data.segment_index >= len(flat_tensor.segments): + raise ValueError( + f"PTD key {named_data.key!r} has an invalid segment index" + ) + segment = flat_tensor.segments[named_data.segment_index] + offset = header.segment_base_offset + segment.offset + if offset + segment.size > file_size: + raise ValueError(f"PTD key {named_data.key!r} extends past end of file") + self._blobs[named_data.key] = _PtdBlob(offset, segment.size) + self._digests: Dict[str, bytes] = {} + + def keys(self) -> set[str]: + return set(self._blobs) + + def size(self, key: str) -> int: + try: + return self._blobs[key].size + except KeyError as error: + raise ValueError(f"PTD {self.path} does not contain key {key!r}") from error + + def sha256(self, key: str) -> bytes: + digest = self._digests.get(key) + if digest is not None: + return digest + try: + blob = self._blobs[key] + except KeyError as error: + raise ValueError(f"PTD {self.path} does not contain key {key!r}") from error + + hasher = hashlib.sha256() + remaining = blob.size + with self.path.open("rb") as source: + source.seek(blob.offset) + while remaining: + chunk = source.read(min(8 * 1024 * 1024, remaining)) + if not chunk: + raise ValueError(f"PTD {self.path} is truncated at key {key!r}") + hasher.update(chunk) + remaining -= len(chunk) + digest = hasher.digest() + self._digests[key] = digest + return digest + + +@dataclass +class _CudaDelegate: + identity: Tuple[str, int] + metadata: CudaAotiMetadata + + +@dataclass +class _Artifact: + source: CudaPteInput + pte: PTEFile + pte_named_data: Dict[str, bytes] + ptd: Optional[_PtdIndex] + delegates: List[_CudaDelegate] + + def blob_size(self, key: str) -> int: + data = self.pte_named_data.get(key) + if data is not None: + return len(data) + if self.ptd is not None: + return self.ptd.size(key) + raise ValueError(f"{self.source.pte_path} does not contain named data {key!r}") + + def blob_sha256(self, key: str) -> bytes: + data = self.pte_named_data.get(key) + if data is not None: + return hashlib.sha256(data).digest() + if self.ptd is not None: + return self.ptd.sha256(key) + raise ValueError(f"{self.source.pte_path} does not contain named data {key!r}") + + +def _named_data_bytes(output: Optional[NamedDataStoreOutput]) -> Dict[str, bytes]: + if output is None: + return {} + return { + key: bytes(output.buffers[entry.buffer_index]) + for key, entry in output.pte_data.items() + } + + +def _delegate_payload(program: Program, delegate) -> bytes: + if delegate.processed.location != DataLocation.INLINE: + raise ValueError("PTE deserialization did not restore delegate data inline") + try: + return bytes(program.backend_delegate_data[delegate.processed.index].data) + except IndexError as error: + raise ValueError("CUDA delegate references invalid processed data") from error + + +def _load_artifact(source: CudaPteInput) -> _Artifact: + pte = deserialize_pte_binary(source.pte_path.read_bytes()) + delegates = [] + for plan in pte.program.execution_plan: + for delegate_index, delegate in enumerate(plan.delegates): + if delegate.id != CUDA_BACKEND_ID: + continue + metadata = decode_cuda_aoti_metadata( + _delegate_payload(pte.program, delegate) + ) + if metadata.variants[0].target_sm == 0: + raise ValueError( + f"Untargeted CUDA metadata in {source.pte_path} cannot be merged" + ) + delegates.append(_CudaDelegate((plan.name, delegate_index), metadata)) + if not delegates: + raise ValueError(f"{source.pte_path} contains no CUDA delegates") + ptd = _PtdIndex(source.ptd_path) if source.ptd_path is not None else None + return _Artifact(source, pte, _named_data_bytes(pte.named_data), ptd, delegates) + + +def _normalized_program(program: Program) -> Program: + normalized = copy.deepcopy(program) + payloads = [] + for plan in normalized.execution_plan: + for delegate in plan.delegates: + payload = _delegate_payload(normalized, delegate) + if delegate.id == CUDA_BACKEND_ID: + payload = b"CUDA_AOTI_VARIANTS" + delegate.compile_specs = [ + spec + for spec in delegate.compile_specs + if spec.key != "cuda_include_ptx" + ] + delegate.processed = BackendDelegateDataReference( + location=DataLocation.INLINE, index=len(payloads) + ) + payloads.append(BackendDelegateInlineData(data=payload)) + normalized.backend_delegate_data = payloads + return normalized + + +def _entry_map(entries: Iterable[CudaWeightEntry]) -> Dict[str, CudaWeightEntry]: + result = {} + for entry in entries: + if entry.fqn in result: + raise ValueError(f"Duplicate CUDA weight FQN {entry.fqn!r}") + result[entry.fqn] = entry + return result + + +def _entry_without_storage_key(entry: CudaWeightEntry) -> Tuple[object, ...]: + return ( + entry.fqn, + entry.storage_nbytes, + entry.dtype, + entry.device_type, + entry.storage_offset, + entry.sizes, + entry.strides, + ) + + +def _validate_shared_weights( + reference: _Artifact, + reference_metadata: CudaAotiMetadata, + candidate: _Artifact, + candidate_metadata: CudaAotiMetadata, + identity: Tuple[str, int], +) -> None: + reference_entries = _entry_map(reference_metadata.entries) + candidate_entries = _entry_map(candidate_metadata.entries) + if reference_entries.keys() != candidate_entries.keys(): + raise ValueError(f"CUDA weights differ for delegate {identity}: FQN mismatch") + + for fqn, reference_entry in reference_entries.items(): + candidate_entry = candidate_entries[fqn] + if _entry_without_storage_key(reference_entry) != _entry_without_storage_key( + candidate_entry + ): + raise ValueError( + f"CUDA weight metadata differs for delegate {identity}, FQN {fqn!r}" + ) + if ( + reference.blob_size(reference_entry.storage_key) + != reference_entry.storage_nbytes + ): + raise ValueError( + f"CUDA weight {fqn!r} has an invalid size in {reference.source.pte_path}" + ) + if ( + candidate.blob_size(candidate_entry.storage_key) + != candidate_entry.storage_nbytes + ): + raise ValueError( + f"CUDA weight {fqn!r} has an invalid size in {candidate.source.pte_path}" + ) + if reference.blob_sha256(reference_entry.storage_key) != candidate.blob_sha256( + candidate_entry.storage_key + ): + raise ValueError( + f"CUDA weight content differs for delegate {identity}, FQN {fqn!r}" + ) + + +def _cuda_so_keys(artifact: _Artifact) -> set[str]: + return { + variant.so_blob_key + for delegate in artifact.delegates + for variant in delegate.metadata.variants + } + + +def _cuda_weight_keys(artifact: _Artifact) -> set[str]: + return { + entry.storage_key + for delegate in artifact.delegates + for entry in delegate.metadata.entries + } + + +def _validate_programs(reference: _Artifact, candidate: _Artifact) -> None: + if _normalized_program(reference.pte.program) != _normalized_program( + candidate.pte.program + ): + raise ValueError( + f"ExecuTorch programs differ between {reference.source.pte_path} and " + f"{candidate.source.pte_path}" + ) + if reference.pte.mutable_data != candidate.pte.mutable_data: + raise ValueError( + f"Mutable program data differs between {reference.source.pte_path} and " + f"{candidate.source.pte_path}" + ) + + reference_non_cuda = { + key: value + for key, value in reference.pte_named_data.items() + if key not in _cuda_so_keys(reference) + and key not in _cuda_weight_keys(reference) + } + candidate_non_cuda = { + key: value + for key, value in candidate.pte_named_data.items() + if key not in _cuda_so_keys(candidate) + and key not in _cuda_weight_keys(candidate) + } + if reference_non_cuda != candidate_non_cuda: + raise ValueError( + f"Non-CUDA named data differs between {reference.source.pte_path} and " + f"{candidate.source.pte_path}" + ) + + reference_external = ( + reference.ptd.keys() - _cuda_weight_keys(reference) + if reference.ptd is not None + else set() + ) + candidate_external = ( + candidate.ptd.keys() - _cuda_weight_keys(candidate) + if candidate.ptd is not None + else set() + ) + if reference_external != candidate_external: + raise ValueError( + f"Non-CUDA external data differs between {reference.source.pte_path} and " + f"{candidate.source.pte_path}" + ) + for key in reference_external: + if reference.blob_size(key) != candidate.blob_size( + key + ) or reference.blob_sha256(key) != candidate.blob_sha256(key): + raise ValueError( + f"External data {key!r} differs between " + f"{reference.source.pte_path} and {candidate.source.pte_path}" + ) + + +def _compact_delegate_data(program: Program) -> None: + payloads = [] + for plan in program.execution_plan: + for delegate in plan.delegates: + payload = _delegate_payload(program, delegate) + delegate.processed = BackendDelegateDataReference( + location=DataLocation.INLINE, index=len(payloads) + ) + payloads.append(BackendDelegateInlineData(data=payload)) + program.backend_delegate_data = payloads + + +def _merge_delegate_variants( + regular_artifacts: Sequence[_Artifact], + regular_delegates: Sequence[Dict[Tuple[str, int], CudaAotiMetadata]], + fallback_artifact: Optional[_Artifact], + fallback_delegates: Optional[Dict[Tuple[str, int], CudaAotiMetadata]], + reference_metadata: CudaAotiMetadata, + identity: Tuple[str, int], + merged_store: NamedDataStore, + provenance: List[CudaPteProvenance], +) -> List[CudaAotiVariant]: + variants = [] + target_sms = set() + reference = regular_artifacts[0] + for artifact, delegates in zip(regular_artifacts, regular_delegates): + metadata = delegates[identity] + _validate_shared_weights( + reference, reference_metadata, artifact, metadata, identity + ) + for variant in metadata.variants: + if variant.target_sm in target_sms: + raise ValueError(f"Duplicate CUDA target sm{variant.target_sm}") + target_sms.add(variant.target_sm) + try: + so_data = artifact.pte_named_data[variant.so_blob_key] + except KeyError as error: + raise ValueError( + f"{artifact.source.pte_path} does not contain CUDA SO " + f"{variant.so_blob_key!r}" + ) from error + merged_store.add_named_data(variant.so_blob_key, so_data) + variants.append(replace(variant, ptx_compute=0, fallback_only=False)) + provenance.append( + CudaPteProvenance( + delegate=identity, + kind="cubin", + target_sm=variant.target_sm, + ptx_compute=0, + source_pte=artifact.source.pte_path, + ) + ) + + variants.sort(key=lambda variant: variant.target_sm) + if fallback_artifact is not None: + assert fallback_delegates is not None + metadata = fallback_delegates[identity] + _validate_shared_weights( + reference, reference_metadata, fallback_artifact, metadata, identity + ) + fallback_variants = [ + variant for variant in metadata.variants if variant.ptx_compute + ] + if len(fallback_variants) != 1: + raise ValueError( + f"Fallback PTE {fallback_artifact.source.pte_path} must contain " + f"exactly one PTX-capable variant for delegate {identity}" + ) + fallback = replace(fallback_variants[0], fallback_only=True) + try: + so_data = fallback_artifact.pte_named_data[fallback.so_blob_key] + except KeyError as error: + raise ValueError( + f"{fallback_artifact.source.pte_path} does not contain CUDA SO " + f"{fallback.so_blob_key!r}" + ) from error + merged_store.add_named_data(fallback.so_blob_key, so_data) + variants.append(fallback) + provenance.append( + CudaPteProvenance( + delegate=identity, + kind="ptx-fallback", + target_sm=fallback.target_sm, + ptx_compute=fallback.ptx_compute, + source_pte=fallback_artifact.source.pte_path, + ) + ) + return variants + + +def _load_merge_artifacts( + inputs: Sequence[CudaPteInput], fallback: Optional[CudaPteInput] +) -> Tuple[List[_Artifact], Optional[_Artifact], List[Tuple[str, int]]]: + regular_artifacts = [_load_artifact(source) for source in inputs] + fallback_artifact = _load_artifact(fallback) if fallback is not None else None + artifacts = [*regular_artifacts] + if fallback_artifact is not None: + artifacts.append(fallback_artifact) + + reference = regular_artifacts[0] + reference_identities = [delegate.identity for delegate in reference.delegates] + for candidate in artifacts[1:]: + _validate_programs(reference, candidate) + candidate_identities = [delegate.identity for delegate in candidate.delegates] + if candidate_identities != reference_identities: + raise ValueError( + f"CUDA delegate layout differs between {reference.source.pte_path} " + f"and {candidate.source.pte_path}" + ) + return regular_artifacts, fallback_artifact, reference_identities + + +def _prepare_merged_output(reference: _Artifact) -> Tuple[Program, NamedDataStore]: + merged_program = copy.deepcopy(reference.pte.program) + for plan in merged_program.execution_plan: + for delegate in plan.delegates: + if delegate.id == CUDA_BACKEND_ID: + delegate.compile_specs = [ + spec + for spec in delegate.compile_specs + if spec.key != "cuda_include_ptx" + ] + + merged_store = NamedDataStore() + if reference.pte.named_data is not None: + merged_store.merge_named_data_store(reference.pte.named_data) + return merged_program, merged_store + + +def merge_cuda_pte_files_with_provenance( + inputs: Sequence[CudaPteInput], fallback: Optional[CudaPteInput] = None +) -> CudaPteMergeResult: + """Merge exact-SM CUDA exports and an optional PTX-only fallback source.""" + if torch.version.hip is not None: + raise RuntimeError( + "CUDA PTE merging supports only NVIDIA CUDA and is not supported on ROCm" + ) + if not inputs: + raise ValueError("At least one regular CUDA PTE input is required") + if len(inputs) + int(fallback is not None) < 2: + raise ValueError("At least two CUDA PTE inputs are required") + regular_artifacts, fallback_artifact, reference_identities = _load_merge_artifacts( + inputs, fallback + ) + reference = regular_artifacts[0] + merged_program, merged_store = _prepare_merged_output(reference) + + regular_delegates = [ + {delegate.identity: delegate.metadata for delegate in artifact.delegates} + for artifact in regular_artifacts + ] + fallback_delegates = ( + { + delegate.identity: delegate.metadata + for delegate in fallback_artifact.delegates + } + if fallback_artifact is not None + else None + ) + expected_variants = None + provenance: List[CudaPteProvenance] = [] + for identity_index, identity in enumerate(reference_identities): + reference_metadata = reference.delegates[identity_index].metadata + variants = _merge_delegate_variants( + regular_artifacts, + regular_delegates, + fallback_artifact, + fallback_delegates, + reference_metadata, + identity, + merged_store, + provenance, + ) + current_variants = tuple( + (variant.target_sm, variant.ptx_compute, variant.fallback_only) + for variant in variants + ) + if expected_variants is None: + expected_variants = current_variants + elif current_variants != expected_variants: + raise ValueError( + f"CUDA target variants differ across delegates at {identity}" + ) + merged_payload = encode_cuda_aoti_metadata(variants, reference_metadata.entries) + plan_name, delegate_index = identity + plan = next( + plan for plan in merged_program.execution_plan if plan.name == plan_name + ) + delegate = plan.delegates[delegate_index] + delegate.processed = BackendDelegateDataReference( + location=DataLocation.INLINE, + index=len(merged_program.backend_delegate_data), + ) + merged_program.backend_delegate_data.append( + BackendDelegateInlineData(data=merged_payload) + ) + + _compact_delegate_data(merged_program) + return CudaPteMergeResult( + pte=serialize_pte_binary( + PTEFile( + program=merged_program, + mutable_data=reference.pte.mutable_data, + named_data=merged_store.get_named_data_store_output(), + ), + extract_delegate_segments=True, + ), + provenance=tuple(provenance), + ) + + +def merge_cuda_pte_files( + inputs: Sequence[CudaPteInput], fallback: Optional[CudaPteInput] = None +) -> Cord: + return merge_cuda_pte_files_with_provenance(inputs, fallback).pte + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=("Merge exact-SM CUDA PTEs with an optional explicit PTX fallback") + ) + parser.add_argument( + "--input-pte", + action="append", + required=True, + type=Path, + help=( + "Regular CUDA PTE contributing exact-SM native cubins; the first " + "input supplies common data" + ), + ) + parser.add_argument( + "--input-ptd", + action="append", + type=Path, + default=[], + help="PTD paired by position with --input-pte", + ) + parser.add_argument( + "--fallback-pte", + type=Path, + help="Single CUDA PTE contributing only the PTX runtime fallback", + ) + parser.add_argument( + "--fallback-ptd", + type=Path, + help="PTD paired with --fallback-pte", + ) + parser.add_argument("--output-pte", required=True, type=Path) + parser.add_argument("--output-ptd", type=Path) + return parser.parse_args() + + +def main() -> None: + """Command-line entry point for CUDA PTE merging.""" + args = _parse_args() + if args.input_ptd and len(args.input_ptd) != len(args.input_pte): + raise ValueError("--input-ptd must be provided once per --input-pte") + if args.input_ptd and args.output_ptd is None: + raise ValueError("--output-ptd is required when --input-ptd is provided") + if args.fallback_ptd is not None and args.fallback_pte is None: + raise ValueError("--fallback-ptd requires --fallback-pte") + sources = [ + CudaPteInput( + pte_path=pte_path, + ptd_path=args.input_ptd[index] if args.input_ptd else None, + ) + for index, pte_path in enumerate(args.input_pte) + ] + fallback = ( + CudaPteInput( + pte_path=args.fallback_pte, + ptd_path=args.fallback_ptd, + ) + if args.fallback_pte is not None + else None + ) + result = merge_cuda_pte_files_with_provenance(sources, fallback) + args.output_pte.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + dir=args.output_pte.parent, prefix=f".{args.output_pte.name}.", delete=False + ) as temporary: + temporary_path = Path(temporary.name) + result.pte.write_to_file(temporary) + os.replace(temporary_path, args.output_pte) + + if args.output_ptd is not None: + if not args.input_ptd: + raise ValueError("--output-ptd requires --input-ptd") + args.output_ptd.parent.mkdir(parents=True, exist_ok=True) + if args.input_ptd[0].resolve() != args.output_ptd.resolve(): + shutil.copyfile(args.input_ptd[0], args.output_ptd) + + print("Merged CUDA code provenance:") + print("delegate\tkind\ttarget\tsource PTE") + for entry in result.provenance: + delegate = f"{entry.delegate[0]}[{entry.delegate[1]}]" + if entry.kind == "cubin": + target = f"sm{entry.target_sm}" + source = str(entry.source_pte) + else: + target = f"compute_{entry.ptx_compute} (source sm{entry.target_sm})" + source = f"{entry.source_pte} [fallback]" + print(f"{delegate}\t{entry.kind}\t{target}\t{source}") + + +if __name__ == "__main__": + main() diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index 8bb03cc1c1e..c0d61c61100 100644 Binary files a/backends/cuda/runtime/aoti_cuda_shims.lib and b/backends/cuda/runtime/aoti_cuda_shims.lib differ diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib.md b/backends/cuda/runtime/aoti_cuda_shims.lib.md new file mode 100644 index 00000000000..af2fb65ca13 --- /dev/null +++ b/backends/cuda/runtime/aoti_cuda_shims.lib.md @@ -0,0 +1,36 @@ +# aoti_cuda_shims.lib + +Import library for `aoti_cuda_shims.dll`. Lowering a CUDA model for a Windows target +links the generated wrapper against this file, so it has to advertise every shim name +the wrapper can reference. It is checked in because the build that produces the DLL +does not run on the machines that lower for Windows. + +Regenerate it whenever a shim is added, which in practice means whenever the PyTorch +pin moves and the generated wrapper starts calling something new: + +``` +nm --defined-only aoti_cuda_shims.lib \ + | sed -n 's/.* T _\?\(aoti_torch_[a-z_0-9]*\)$/\1/p' \ + | sort -u > exports.txt +# add the new names to exports.txt, then +{ echo 'LIBRARY aoti_cuda_shims.dll'; echo 'EXPORTS'; sed 's/^/ /' exports.txt; } \ + > aoti_cuda_shims.def +x86_64-w64-mingw32-dlltool -d aoti_cuda_shims.def -l aoti_cuda_shims.lib \ + --dllname aoti_cuda_shims.dll +# zero the archive metadata so the file is reproducible +mkdir extract && cd extract && x86_64-w64-mingw32-ar x ../aoti_cuda_shims.lib \ + && rm -f ../aoti_cuda_shims.lib \ + && x86_64-w64-mingw32-ar rcsD ../aoti_cuda_shims.lib $(ls | sort) +``` + +The archive has to be built at a fresh path. Updating it in place keeps the reverse +member order the generator produced, so the same export list would not give the same +bytes twice. + +Then check the result is a superset of what it replaced and that no member carries a +timestamp, since this file ships in the wheel and a stamped one makes two builds of +the same source differ. + +A name only resolves if something in the DLL defines it. The DLL is built from the +CUDA shims plus the SlimTensor common shims, so a shim added to the ETensor common +shims will link and then fail to load. diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 0f63f662c83..3e0c3175818 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -322,7 +322,46 @@ class ET_EXPERIMENTAL CudaBackend final CudaWeightCache::parse( processed->data(), processed->size(), fqn_weights), "Malformed CUDA FQN weight metadata"); - so_blob_key = fqn_weights.so_blob_key; + size_t variant_index = 0; + bool uses_ptx_fallback = false; + uint32_t current_sm = 0; + if (!(fqn_weights.variants.size() == 1 && + fqn_weights.variants[0].target_sm == 0)) { +#if defined(EXECUTORCH_USE_HIP) + ET_LOG( + Error, + "Multi-SM CUDA AOTI metadata is not supported by the ROCm runtime"); + return Error::NotSupported; +#else + int device_index = 0; + cudaDeviceProp device_properties{}; + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaGetDevice(&device_index)); + ET_CUDA_CHECK_OR_RETURN_ERROR( + cudaGetDeviceProperties(&device_properties, device_index)); + current_sm = static_cast( + device_properties.major * 10 + device_properties.minor); + ET_CHECK_OK_OR_RETURN_ERROR( + CudaWeightCache::select_variant( + fqn_weights, current_sm, variant_index, uses_ptx_fallback), + "Failed to select a CUDA AOTI variant for sm%u", + current_sm); +#endif + } + const auto& variant = fqn_weights.variants[variant_index]; + so_blob_key = variant.so_blob_key; + if (variant.target_sm == 0) { + ET_LOG(Info, "Selected untargeted CUDA AOTI variant"); + } else if (uses_ptx_fallback) { + ET_LOG( + Info, + "Selected sm%u CUDA AOTI PTX fallback (compute_%u) for sm%u", + variant.target_sm, + variant.ptx_compute, + current_sm); + } else { + ET_LOG( + Info, "Selected native sm%u CUDA AOTI variant", variant.target_sm); + } } else { ET_CHECK_OK_OR_RETURN_ERROR( executorch::backends::aoti::resolve_blob_keys( diff --git a/backends/cuda/runtime/cuda_weight_cache.cpp b/backends/cuda/runtime/cuda_weight_cache.cpp index 09b4bd76312..5a2e465f243 100644 --- a/backends/cuda/runtime/cuda_weight_cache.cpp +++ b/backends/cuda/runtime/cuda_weight_cache.cpp @@ -152,9 +152,51 @@ Error CudaWeightCache::parse( } MetadataReader reader(data, size); - if (!reader.skip(kFormatMagicSize) || - !reader.read_string(metadata.so_blob_key) || - metadata.so_blob_key.empty()) { + if (!reader.skip(kFormatMagicSize)) { + return Error::InvalidProgram; + } + + metadata.variants.clear(); + uint32_t num_variants = 0; + constexpr uint32_t kMaxVariants = 256; + if (!reader.read_u32(num_variants) || num_variants == 0 || + num_variants > kMaxVariants) { + return Error::InvalidProgram; + } + metadata.variants.reserve(num_variants); + std::unordered_set target_sms; + bool found_fallback = false; + bool regular_has_ptx = false; + for (uint32_t index = 0; index < num_variants; ++index) { + Variant variant; + uint32_t flags = 0; + if (!reader.read_u32(variant.target_sm) || + !reader.read_u32(variant.ptx_compute) || + variant.ptx_compute > variant.target_sm || !reader.read_u32(flags) || + (flags & ~1U) != 0 || !reader.read_string(variant.so_blob_key) || + variant.so_blob_key.empty()) { + return Error::InvalidProgram; + } + variant.fallback_only = (flags & 1U) != 0; + if (variant.target_sm == 0) { + if (num_variants != 1 || variant.ptx_compute != 0 || + variant.fallback_only) { + return Error::InvalidProgram; + } + } else if (variant.fallback_only) { + if (variant.ptx_compute == 0 || found_fallback) { + return Error::InvalidProgram; + } + found_fallback = true; + } else { + if (!target_sms.emplace(variant.target_sm).second) { + return Error::InvalidProgram; + } + regular_has_ptx |= variant.ptx_compute != 0; + } + metadata.variants.push_back(std::move(variant)); + } + if (num_variants > 1 && regular_has_ptx) { return Error::InvalidProgram; } @@ -202,6 +244,45 @@ Error CudaWeightCache::parse( return reader.empty() ? Error::Ok : Error::InvalidProgram; } +Error CudaWeightCache::select_variant( + const Metadata& metadata, + uint32_t current_sm, + size_t& variant_index, + bool& uses_ptx_fallback) { + ET_CHECK_OR_RETURN_ERROR( + !metadata.variants.empty(), InvalidProgram, "CUDA AOTI has no variants"); + + if (metadata.variants.size() == 1 && metadata.variants[0].target_sm == 0) { + variant_index = 0; + uses_ptx_fallback = false; + return Error::Ok; + } + + for (size_t index = 0; index < metadata.variants.size(); ++index) { + if (!metadata.variants[index].fallback_only && + metadata.variants[index].target_sm == current_sm) { + variant_index = index; + uses_ptx_fallback = false; + return Error::Ok; + } + } + + for (size_t index = 0; index < metadata.variants.size(); ++index) { + const Variant& variant = metadata.variants[index]; + if (variant.ptx_compute != 0 && variant.ptx_compute <= current_sm && + (variant.fallback_only || metadata.variants.size() == 1)) { + variant_index = index; + uses_ptx_fallback = true; + return Error::Ok; + } + } + ET_CHECK_OR_RETURN_ERROR( + false, + NotSupported, + "CUDA AOTI has no native or PTX variant compatible with sm%u", + current_sm); +} + Error CudaWeightCache::validate_view(const Entry& entry) { uint64_t item_size = 0; switch (static_cast(entry.dtype)) { diff --git a/backends/cuda/runtime/cuda_weight_cache.h b/backends/cuda/runtime/cuda_weight_cache.h index eb58ebc8fb1..ee77f95bcd7 100644 --- a/backends/cuda/runtime/cuda_weight_cache.h +++ b/backends/cuda/runtime/cuda_weight_cache.h @@ -24,9 +24,16 @@ namespace executorch::backends::cuda { class CudaWeightCache final { public: - static constexpr char kFormatMagic[] = "ETCUDAFQN3"; + static constexpr char kFormatMagic[] = "ETCUDAFQN0"; static constexpr size_t kFormatMagicSize = sizeof(kFormatMagic) - 1; + struct Variant { + uint32_t target_sm{0}; + uint32_t ptx_compute{0}; + std::string so_blob_key; + bool fallback_only{false}; + }; + struct Entry { std::string fqn; std::string storage_key; @@ -39,7 +46,7 @@ class CudaWeightCache final { }; struct Metadata { - std::string so_blob_key; + std::vector variants; std::vector entries; }; @@ -48,6 +55,12 @@ class CudaWeightCache final { static runtime::Error parse(const void* data, size_t size, Metadata& metadata); + static runtime::Error select_variant( + const Metadata& metadata, + uint32_t current_sm, + size_t& variant_index, + bool& uses_ptx_fallback); + runtime::Error load( CudaDelegateHandle* handle, const runtime::NamedDataMap* named_data_map, diff --git a/backends/cuda/runtime/shims/memory.cpp b/backends/cuda/runtime/shims/memory.cpp index 8a81916ab6c..bfd43c25090 100644 --- a/backends/cuda/runtime/shims/memory.cpp +++ b/backends/cuda/runtime/shims/memory.cpp @@ -6,7 +6,9 @@ * LICENSE file in the root directory of this source tree. */ +#include #include +#include #include #include @@ -31,10 +33,63 @@ const PalInitializer kPalInitializer{}; } // namespace namespace c10 = executorch::backends::aoti::slim::c10; + using c10::Device; using c10::DeviceIndex; using c10::DeviceType; using c10::ScalarType; + +namespace { + +// Reads the one element as T. Dispatch on the dtype the tensor actually holds: +// item() copies sizeof(T) bytes and does not check, so asking for the wrong +// width reads past a one-element allocation. +// +// The dtype in each entry point's name is the type the caller wants back, not +// the type the tensor holds. Generated code reads a boolean branch selector +// through the int64 entry point, for instance. So convert, and refuse only when +// the value does not fit, which is what the reference implementation does. +template +AOTITorchError narrow_to(From value, To* ret_value) { + // Anything at all converts to a boolean, so there is nothing to check there. + if (!std::is_same_v && ::c10::overflows(value)) { + ET_CHECK_OR_RETURN_ERROR( + false, InvalidArgument, "reading a single element: value does not fit"); + } + *ret_value = static_cast(value); + return Error::Ok; +} + +template +AOTITorchError read_one_element(const SlimTensor* tensor, T* ret_value) { + switch (tensor->dtype()) { + case ScalarType::Byte: + return narrow_to(tensor->item(), ret_value); + case ScalarType::Char: + return narrow_to(tensor->item(), ret_value); + case ScalarType::Short: + return narrow_to(tensor->item(), ret_value); + case ScalarType::Int: + return narrow_to(tensor->item(), ret_value); + case ScalarType::Long: + return narrow_to(tensor->item(), ret_value); + case ScalarType::Float: + return narrow_to(tensor->item(), ret_value); + case ScalarType::BFloat16: + return narrow_to( + static_cast(tensor->item()), ret_value); + case ScalarType::Bool: + return narrow_to(tensor->item(), ret_value); + default: + ET_CHECK_OR_RETURN_ERROR( + false, + InvalidArgument, + "reading a single element: dtype %d is not supported", + static_cast(tensor->dtype())); + } +} + +} // namespace using executorch::backends::aoti::slim::empty_strided; using executorch::backends::aoti::slim::from_blob; using executorch::backends::aoti::slim::IntArrayRef; @@ -195,6 +250,31 @@ AOTITorchError aoti_torch_empty_strided( return Error::Ok; } +AOTITorchError aoti_torch_empty_strided_pinned( + int64_t ndim, + const int64_t* sizes_ptr, + const int64_t* strides_ptr, + int32_t dtype, + int32_t device_type, + int32_t device_index, + SlimTensor** ret_new_tensor) { + ET_CHECK_OR_RETURN_ERROR( + static_cast(device_type) == DeviceType::CPU, + InvalidArgument, + "aoti_torch_empty_strided_pinned: pinned memory is host memory, so the " + "device type must be CPU, got %d", + device_type); + + return aoti_torch_empty_strided( + ndim, + sizes_ptr, + strides_ptr, + dtype, + device_type, + device_index, + ret_new_tensor); +} + AOTITorchError aoti_torch_delete_tensor_object(SlimTensor* tensor) { ET_CHECK_OR_RETURN_ERROR( tensor != nullptr, @@ -322,34 +402,36 @@ aoti_torch_copy_(SlimTensor* self, SlimTensor* src, int32_t non_blocking) { return Error::Ok; } -AOTITorchError aoti_torch_item_bool(SlimTensor* tensor, bool* ret_value) { - ET_CHECK_OR_RETURN_ERROR( - tensor != nullptr, - InvalidArgument, - "aoti_torch_item_bool: tensor is null"); - - ET_CHECK_OR_RETURN_ERROR( - ret_value != nullptr, - InvalidArgument, - "aoti_torch_item_bool: ret_value is null"); - - ET_CHECK_OR_RETURN_ERROR( - tensor->numel() == 1, - InvalidArgument, - "aoti_torch_item_bool: tensor must have exactly 1 element, got %zu", - tensor->numel()); - - ET_CHECK_OR_RETURN_ERROR( - tensor->dtype() == ScalarType::Bool, - InvalidArgument, - "aoti_torch_item_bool: tensor dtype must be Bool"); +#define ET_CUDA_DEFINE_ITEM_SHIM(SUFFIX, CTYPE) \ + AOTITorchError aoti_torch_item_##SUFFIX( \ + SlimTensor* tensor, CTYPE* ret_value) { \ + ET_CHECK_OR_RETURN_ERROR( \ + tensor != nullptr, \ + InvalidArgument, \ + "aoti_torch_item_" #SUFFIX ": tensor is null"); \ + ET_CHECK_OR_RETURN_ERROR( \ + ret_value != nullptr, \ + InvalidArgument, \ + "aoti_torch_item_" #SUFFIX ": ret_value is null"); \ + ET_CHECK_OR_RETURN_ERROR( \ + tensor->numel() == 1, \ + InvalidArgument, \ + "aoti_torch_item_" #SUFFIX \ + ": tensor must have exactly 1 element, got %zu", \ + tensor->numel()); \ + return read_one_element(tensor, ret_value); \ + } - // SlimTensor::item() handles both CPU and CUDA tensors. - // For CUDA tensors, it copies the value to CPU automatically. - *ret_value = tensor->item(); +ET_CUDA_DEFINE_ITEM_SHIM(uint8, uint8_t) +ET_CUDA_DEFINE_ITEM_SHIM(int8, int8_t) +ET_CUDA_DEFINE_ITEM_SHIM(int16, int16_t) +ET_CUDA_DEFINE_ITEM_SHIM(int32, int32_t) +ET_CUDA_DEFINE_ITEM_SHIM(int64, int64_t) +ET_CUDA_DEFINE_ITEM_SHIM(float32, float) +ET_CUDA_DEFINE_ITEM_SHIM(bfloat16, c10::BFloat16) +ET_CUDA_DEFINE_ITEM_SHIM(bool, bool) - return Error::Ok; -} +#undef ET_CUDA_DEFINE_ITEM_SHIM AOTITorchError aoti_torch_assign_tensors_out( SlimTensor* src, diff --git a/backends/cuda/runtime/shims/memory.h b/backends/cuda/runtime/shims/memory.h index ca464a9acf5..96c47c03aed 100644 --- a/backends/cuda/runtime/shims/memory.h +++ b/backends/cuda/runtime/shims/memory.h @@ -95,6 +95,31 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_empty_strided( int32_t device_index, SlimTensor** ret_new_tensor); +/** + * Allocates ordinary host memory where pinned host memory was asked for. + * + * There is no pinned allocator here. Pinning only lets a copy to the device + * overlap other work, so ordinary memory is correct and slower. + * + * @param ndim Number of dimensions + * @param sizes_ptr Pointer to the sizes, ndim of them + * @param strides_ptr Pointer to the strides, ndim of them, or null for + * contiguous + * @param dtype Element type, as a scalar type value + * @param device_type Must be CPU, since pinned memory is host memory + * @param device_index Device index, unused for CPU + * @param ret_new_tensor Receives the new tensor + * @return Error::Ok on success, Error::InvalidArgument if the device is not CPU + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_empty_strided_pinned( + int64_t ndim, + const int64_t* sizes_ptr, + const int64_t* strides_ptr, + int32_t dtype, + int32_t device_type, + int32_t device_index, + SlimTensor** ret_new_tensor); + /** * Deletes a tensor object and frees associated resources. * @@ -170,18 +195,49 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch__reinterpret_tensor( AOTI_SHIM_EXPORT AOTITorchError aoti_torch_copy_(SlimTensor* self, SlimTensor* src, int32_t non_blocking); +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_bool(SlimTensor* tensor, bool* ret_value); + /** - * Extracts a boolean scalar value from a single-element tensor. + * Extracts a scalar value from a single-element tensor. * - * The tensor must contain exactly one element and have Bool dtype. - * For CUDA tensors, this will synchronize to copy the value to CPU. + * The type in the name is the type returned, not the tensor's dtype: generated + * code reads a boolean branch selector through the int64 entry point. The value + * is converted, and rejected only when it does not fit. The tensor must contain + * exactly one element. For CUDA tensors, this will synchronize to copy the + * value to CPU. * - * @param tensor Single-element boolean tensor (must not be null) - * @param ret_value Output parameter for the extracted boolean value + * @param tensor Single-element tensor (must not be null) + * @param ret_value Output parameter for the extracted value * @return AOTITorchError error code (Error::Ok on success) */ AOTI_SHIM_EXPORT AOTITorchError -aoti_torch_item_bool(SlimTensor* tensor, bool* ret_value); +aoti_torch_item_uint8(SlimTensor* tensor, uint8_t* ret_value); + +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_int8(SlimTensor* tensor, int8_t* ret_value); + +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_int16(SlimTensor* tensor, int16_t* ret_value); + +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_int32(SlimTensor* tensor, int32_t* ret_value); + +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_int64(SlimTensor* tensor, int64_t* ret_value); + +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_float32(SlimTensor* tensor, float* ret_value); + +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_bfloat16(SlimTensor* tensor, c10::BFloat16* ret_value); /** * Moves a tensor into a new handle and assigns it to the output parameter. diff --git a/backends/cuda/runtime/shims/tests/test_aoti_torch_item_bool.cpp b/backends/cuda/runtime/shims/tests/test_aoti_torch_item_bool.cpp index fa8b5bb9245..cc4b68d372c 100644 --- a/backends/cuda/runtime/shims/tests/test_aoti_torch_item_bool.cpp +++ b/backends/cuda/runtime/shims/tests/test_aoti_torch_item_bool.cpp @@ -155,49 +155,62 @@ TEST_F(AOTITorchItemBoolSlimTest, NullReturnValue) { EXPECT_EQ(aoti_torch_delete_tensor_object(tensor), Error::Ok); } -TEST_F(AOTITorchItemBoolSlimTest, MultiElementTensor) { - std::vector sizes = {2, 3}; +TEST_F(AOTITorchItemBoolSlimTest, ConvertsFromLong) { + // Generated code reads a value through whichever entry point matches the type + // it wants back, not the type the tensor holds. + std::vector sizes = {1}; Tensor* tensor = createTestTensor( sizes, - static_cast(slim_c10::ScalarType::Bool), + static_cast(slim_c10::ScalarType::Long), static_cast(slim_c10::DeviceType::CPU), 0); ASSERT_NE(tensor, nullptr); - EXPECT_GT(tensor->numel(), 1); + *static_cast(tensor->data_ptr()) = 1; bool result = false; - AOTITorchError error = aoti_torch_item_bool(tensor, &result); + EXPECT_EQ(aoti_torch_item_bool(tensor, &result), Error::Ok); + EXPECT_TRUE(result); - EXPECT_EQ(error, Error::InvalidArgument); + EXPECT_EQ(aoti_torch_delete_tensor_object(tensor), Error::Ok); +} + +TEST_F(AOTITorchItemBoolSlimTest, BoolReadAsInt64) { + Tensor* tensor = createScalarBoolTensor( + true, static_cast(slim_c10::DeviceType::CPU), 0); + ASSERT_NE(tensor, nullptr); + + int64_t result = -1; + EXPECT_EQ(aoti_torch_item_int64(tensor, &result), Error::Ok); + EXPECT_EQ(result, 1); EXPECT_EQ(aoti_torch_delete_tensor_object(tensor), Error::Ok); } -TEST_F(AOTITorchItemBoolSlimTest, WrongDtype_Float) { +TEST_F(AOTITorchItemBoolSlimTest, RejectsValueThatDoesNotFit) { std::vector sizes = {1}; Tensor* tensor = createTestTensor( sizes, - static_cast(slim_c10::ScalarType::Float), + static_cast(slim_c10::ScalarType::Long), static_cast(slim_c10::DeviceType::CPU), 0); ASSERT_NE(tensor, nullptr); + *static_cast(tensor->data_ptr()) = 300; - bool result = false; - AOTITorchError error = aoti_torch_item_bool(tensor, &result); - - EXPECT_EQ(error, Error::InvalidArgument); + uint8_t result = 0; + EXPECT_EQ(aoti_torch_item_uint8(tensor, &result), Error::InvalidArgument); EXPECT_EQ(aoti_torch_delete_tensor_object(tensor), Error::Ok); } -TEST_F(AOTITorchItemBoolSlimTest, WrongDtype_Long) { - std::vector sizes = {1}; +TEST_F(AOTITorchItemBoolSlimTest, MultiElementTensor) { + std::vector sizes = {2, 3}; Tensor* tensor = createTestTensor( sizes, - static_cast(slim_c10::ScalarType::Long), + static_cast(slim_c10::ScalarType::Bool), static_cast(slim_c10::DeviceType::CPU), 0); ASSERT_NE(tensor, nullptr); + EXPECT_GT(tensor->numel(), 1); bool result = false; AOTITorchError error = aoti_torch_item_bool(tensor, &result); @@ -270,24 +283,3 @@ TEST_F(AOTITorchItemBoolSlimTest, MultiElementTensor_CUDA) { EXPECT_EQ(aoti_torch_delete_tensor_object(tensor), Error::Ok); } - -TEST_F(AOTITorchItemBoolSlimTest, WrongDtype_Float_CUDA) { - if (!isCudaAvailable()) { - GTEST_SKIP() << "CUDA not available"; - } - - std::vector sizes = {1}; - Tensor* tensor = createTestTensor( - sizes, - static_cast(slim_c10::ScalarType::Float), - static_cast(slim_c10::DeviceType::CUDA), - 0); - ASSERT_NE(tensor, nullptr); - - bool result = false; - AOTITorchError error = aoti_torch_item_bool(tensor, &result); - - EXPECT_EQ(error, Error::InvalidArgument); - - EXPECT_EQ(aoti_torch_delete_tensor_object(tensor), Error::Ok); -} diff --git a/backends/cuda/runtime/test/test_cuda_weight_cache.cpp b/backends/cuda/runtime/test/test_cuda_weight_cache.cpp index d16e7143058..daa0a41bf0f 100644 --- a/backends/cuda/runtime/test/test_cuda_weight_cache.cpp +++ b/backends/cuda/runtime/test/test_cuda_weight_cache.cpp @@ -43,6 +43,10 @@ std::vector serialized_metadata( cuda::CudaWeightCache::kFormatMagic, cuda::CudaWeightCache::kFormatMagic + cuda::CudaWeightCache::kFormatMagicSize); + append_u32(output, 1); // variants + append_u32(output, 0); // untargeted (ROCm) + append_u32(output, 0); // no PTX + append_u32(output, 0); // regular append_string(output, "so-key"); append_u32(output, 1); // entries append_string(output, "model.weight"); @@ -59,9 +63,60 @@ std::vector serialized_metadata( return output; } +std::vector serialized_multi_arch_metadata() { + std::vector output( + cuda::CudaWeightCache::kFormatMagic, + cuda::CudaWeightCache::kFormatMagic + + cuda::CudaWeightCache::kFormatMagicSize); + append_u32(output, 3); // variants + append_u32(output, 80); + append_u32(output, 0); + append_u32(output, 0); // regular + append_string(output, "sm80-so"); + append_u32(output, 90); + append_u32(output, 0); + append_u32(output, 0); // regular + append_string(output, "sm90-so"); + append_u32(output, 120); + append_u32(output, 0); + append_u32(output, 0); // regular + append_string(output, "sm120-so"); + append_u32(output, 1); // entries + append_string(output, "model.weight"); + append_string(output, "storage-key"); + append_u64(output, 24); + append_u32(output, 6); + append_u32(output, 1); + append_u64(output, 0); + append_u32(output, 2); + append_u64(output, 2); + append_u64(output, 3); + append_u64(output, 3); + append_u64(output, 1); + return output; +} + +std::vector serialized_fallback_metadata() { + std::vector output( + cuda::CudaWeightCache::kFormatMagic, + cuda::CudaWeightCache::kFormatMagic + + cuda::CudaWeightCache::kFormatMagicSize); + append_u32(output, 2); // variants + append_u32(output, 80); + append_u32(output, 0); + append_u32(output, 0); // regular + append_string(output, "sm80-so"); + append_u32(output, 80); + append_u32(output, 80); + append_u32(output, 1); // fallback only + append_string(output, "fallback-so"); + append_u32(output, 0); // entries + return output; +} + } // namespace -TEST(CudaWeightCacheTest, LegacyPayloadIsNotMisdetected) { +TEST(CudaWeightCacheTest, RawAotiPayloadIsNotMisdetected) { const std::string legacy = "so-key\nweights-key"; EXPECT_FALSE( cuda::CudaWeightCache::is_serialized(legacy.data(), legacy.size())); @@ -73,7 +128,9 @@ TEST(CudaWeightCacheTest, ParsesSerializedMetadata) { ASSERT_EQ( cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), Error::Ok); - ASSERT_EQ(metadata.so_blob_key, "so-key"); + ASSERT_EQ(metadata.variants.size(), 1u); + EXPECT_EQ(metadata.variants[0].target_sm, 0u); + EXPECT_EQ(metadata.variants[0].so_blob_key, "so-key"); ASSERT_EQ(metadata.entries.size(), 1u); const auto& entry = metadata.entries[0]; EXPECT_EQ(entry.fqn, "model.weight"); @@ -85,6 +142,126 @@ TEST(CudaWeightCacheTest, ParsesSerializedMetadata) { EXPECT_EQ(entry.strides, (std::vector{3, 1})); } +TEST(CudaWeightCacheTest, ParsesAndSelectsMultiArchMetadata) { + const std::vector bytes = serialized_multi_arch_metadata(); + cuda::CudaWeightCache::Metadata metadata; + ASSERT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::Ok); + ASSERT_EQ(metadata.variants.size(), 3u); + EXPECT_EQ(metadata.variants[0].target_sm, 80u); + EXPECT_EQ(metadata.variants[2].so_blob_key, "sm120-so"); + + size_t variant_index = 0; + bool uses_ptx_fallback = false; + EXPECT_EQ( + cuda::CudaWeightCache::select_variant( + metadata, 120, variant_index, uses_ptx_fallback), + Error::Ok); + EXPECT_EQ(variant_index, 2u); + EXPECT_FALSE(uses_ptx_fallback); + + EXPECT_EQ( + cuda::CudaWeightCache::select_variant( + metadata, 100, variant_index, uses_ptx_fallback), + Error::NotSupported); +} + +TEST(CudaWeightCacheTest, SelectsPtxFromPortableSingleVariant) { + cuda::CudaWeightCache::Metadata metadata; + metadata.variants = {{80, 80, "sm80-so"}}; + size_t variant_index = 0; + bool uses_ptx_fallback = false; + ASSERT_EQ( + cuda::CudaWeightCache::select_variant( + metadata, 100, variant_index, uses_ptx_fallback), + Error::Ok); + EXPECT_EQ(variant_index, 0u); + EXPECT_TRUE(uses_ptx_fallback); +} + +TEST(CudaWeightCacheTest, FallbackOnlyVariantNeverWinsNativeMatch) { + const std::vector bytes = serialized_fallback_metadata(); + cuda::CudaWeightCache::Metadata metadata; + ASSERT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::Ok); + ASSERT_EQ(metadata.variants.size(), 2u); + EXPECT_FALSE(metadata.variants[0].fallback_only); + EXPECT_TRUE(metadata.variants[1].fallback_only); + + size_t variant_index = 0; + bool uses_ptx_fallback = false; + ASSERT_EQ( + cuda::CudaWeightCache::select_variant( + metadata, 80, variant_index, uses_ptx_fallback), + Error::Ok); + EXPECT_EQ(variant_index, 0u); + EXPECT_FALSE(uses_ptx_fallback); + + ASSERT_EQ( + cuda::CudaWeightCache::select_variant( + metadata, 90, variant_index, uses_ptx_fallback), + Error::Ok); + EXPECT_EQ(variant_index, 1u); + EXPECT_TRUE(uses_ptx_fallback); +} + +TEST(CudaWeightCacheTest, RejectsWhenNoVariantIsCompatible) { + cuda::CudaWeightCache::Metadata metadata; + metadata.variants = {{120, 0, "sm120-so"}}; + size_t variant_index = 0; + bool uses_ptx_fallback = false; + EXPECT_EQ( + cuda::CudaWeightCache::select_variant( + metadata, 90, variant_index, uses_ptx_fallback), + Error::NotSupported); +} + +TEST(CudaWeightCacheTest, RejectsDuplicateMultiArchTarget) { + std::vector bytes( + cuda::CudaWeightCache::kFormatMagic, + cuda::CudaWeightCache::kFormatMagic + + cuda::CudaWeightCache::kFormatMagicSize); + append_u32(bytes, 2); + append_u32(bytes, 80); + append_u32(bytes, 0); + append_u32(bytes, 0); + append_string(bytes, "first-so"); + append_u32(bytes, 80); + append_u32(bytes, 0); + append_u32(bytes, 0); + append_string(bytes, "second-so"); + append_u32(bytes, 0); + + cuda::CudaWeightCache::Metadata metadata; + EXPECT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::InvalidProgram); +} + +TEST(CudaWeightCacheTest, RejectsImplicitPtxFallbackInMultiVariantMetadata) { + std::vector bytes( + cuda::CudaWeightCache::kFormatMagic, + cuda::CudaWeightCache::kFormatMagic + + cuda::CudaWeightCache::kFormatMagicSize); + append_u32(bytes, 2); + append_u32(bytes, 80); + append_u32(bytes, 80); + append_u32(bytes, 0); + append_string(bytes, "sm80-so"); + append_u32(bytes, 120); + append_u32(bytes, 0); + append_u32(bytes, 0); + append_string(bytes, "sm120-so"); + append_u32(bytes, 0); + + cuda::CudaWeightCache::Metadata metadata; + EXPECT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::InvalidProgram); +} + TEST(CudaWeightCacheTest, RejectsTruncationAndTrailingData) { std::vector bytes = serialized_metadata(); cuda::CudaWeightCache::Metadata metadata; diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index eda5e46de41..7b45e677da2 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -133,6 +133,33 @@ def test_target_smem_context_only_patches_exact_triton_limit(self): self.assertIs(triton_compiler.max_shared_mem, local_max_shared_mem) + def test_cuda_include_ptx_compile_spec(self): + with mock.patch.object( + CudaBackend, "_setup_cuda_environment_for_fatbin", return_value=True + ): + options = CudaBackend.get_aoti_compile_options( + [CompileSpec(key="cuda_include_ptx", value=b"ON")] + ) + + self.assertTrue(options["aot_inductor.emit_multi_arch_kernel"]) + + def test_cuda_include_ptx_off_disables_multi_arch_kernel(self): + with mock.patch.object( + CudaBackend, "_setup_cuda_environment_for_fatbin" + ) as setup_fatbin: + options = CudaBackend.get_aoti_compile_options( + [CompileSpec(key="cuda_include_ptx", value=b"OFF")] + ) + + setup_fatbin.assert_not_called() + self.assertFalse(options["aot_inductor.emit_multi_arch_kernel"]) + + def test_invalid_cuda_include_ptx_compile_spec(self): + with self.assertRaisesRegex(ValueError, "Invalid cuda_include_ptx"): + CudaBackend.get_aoti_compile_options( + [CompileSpec(key="cuda_include_ptx", value=b"MAYBE")] + ) + class TestCudaExport(unittest.TestCase): """Test CUDA export functionality for various operations using to_edge_transform_and_lower.""" diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index 153828b2cb0..9a1fafaab68 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -21,9 +21,10 @@ from executorch.backends.cuda.cuda_weight_collector import ( AOTI_DEVICE_TYPE_CPU, AOTI_DEVICE_TYPE_CUDA, - CUDA_WEIGHT_CACHE_MAGIC, + CUDA_AOTI_METADATA_MAGIC, + CudaAotiVariant, CudaWeightCollector, - encode_cuda_weight_metadata, + encode_cuda_aoti_metadata, ) from executorch.exir._serialize._cord import FileBackedData from executorch.exir._serialize._named_data_store import NamedDataStore @@ -95,8 +96,10 @@ def test_weights_are_materialized_as_independent_storages(self) -> None: artifact.storages[artifact.entries[1].storage_key].to_bytes(), ) - metadata = encode_cuda_weight_metadata("so-key", artifact.entries) - self.assertTrue(metadata.startswith(CUDA_WEIGHT_CACHE_MAGIC)) + metadata = encode_cuda_aoti_metadata( + [CudaAotiVariant(0, 0, "so-key")], artifact.entries + ) + self.assertTrue(metadata.startswith(CUDA_AOTI_METADATA_MAGIC)) self.assertIn(b"first", metadata) self.assertIn(b"second", metadata) for storage in artifact.storages.values(): diff --git a/backends/cuda/tests/test_cuda_weight_metadata.py b/backends/cuda/tests/test_cuda_weight_metadata.py new file mode 100644 index 00000000000..d7fb1d93179 --- /dev/null +++ b/backends/cuda/tests/test_cuda_weight_metadata.py @@ -0,0 +1,130 @@ +# 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. + +import unittest + +from executorch.backends.cuda.cuda_weight_collector import ( + AOTI_DEVICE_TYPE_CUDA, + CUDA_AOTI_METADATA_MAGIC, + CudaAotiVariant, + CudaWeightEntry, + decode_cuda_aoti_metadata, + encode_cuda_aoti_metadata, +) + + +class TestCudaWeightMetadata(unittest.TestCase): + @staticmethod + def _entry() -> CudaWeightEntry: + return CudaWeightEntry( + fqn="model.weight", + storage_key="cuda_fqn_weight:cuda:model.weight", + storage_nbytes=24, + dtype=6, + device_type=AOTI_DEVICE_TYPE_CUDA, + storage_offset=0, + sizes=(2, 3), + strides=(3, 1), + ) + + def test_targeted_metadata_has_shared_weights(self) -> None: + entry = self._entry() + encoded = encode_cuda_aoti_metadata( + [ + CudaAotiVariant(80, 0, "sm80-so"), + CudaAotiVariant(120, 0, "sm120-so"), + ], + [entry], + ) + self.assertTrue(encoded.startswith(CUDA_AOTI_METADATA_MAGIC)) + decoded = decode_cuda_aoti_metadata(encoded) + self.assertEqual( + decoded.variants, + [ + CudaAotiVariant(80, 0, "sm80-so"), + CudaAotiVariant(120, 0, "sm120-so"), + ], + ) + self.assertEqual(decoded.entries, [entry]) + + def test_metadata_rejects_duplicate_target(self) -> None: + with self.assertRaisesRegex(ValueError, "Duplicate CUDA target SM"): + encode_cuda_aoti_metadata( + [ + CudaAotiVariant(80, 0, "first"), + CudaAotiVariant(80, 0, "second"), + ], + [self._entry()], + ) + + def test_fallback_metadata_allows_matching_regular_target(self) -> None: + entry = self._entry() + encoded = encode_cuda_aoti_metadata( + [ + CudaAotiVariant(80, 0, "sm80-so"), + CudaAotiVariant(80, 80, "fallback-so", fallback_only=True), + ], + [entry], + ) + self.assertTrue(encoded.startswith(CUDA_AOTI_METADATA_MAGIC)) + decoded = decode_cuda_aoti_metadata(encoded) + self.assertEqual( + decoded.variants, + [ + CudaAotiVariant(80, 0, "sm80-so"), + CudaAotiVariant(80, 80, "fallback-so", fallback_only=True), + ], + ) + + def test_fallback_metadata_rejects_multiple_fallbacks(self) -> None: + with self.assertRaisesRegex(ValueError, "only one fallback"): + encode_cuda_aoti_metadata( + [ + CudaAotiVariant(80, 80, "first", fallback_only=True), + CudaAotiVariant(75, 75, "second", fallback_only=True), + ], + [self._entry()], + ) + + def test_multi_variant_metadata_rejects_implicit_ptx_fallback(self) -> None: + with self.assertRaisesRegex(ValueError, "explicit PTX fallback"): + encode_cuda_aoti_metadata( + [ + CudaAotiVariant(80, 80, "sm80-so"), + CudaAotiVariant(120, 0, "sm120-so"), + ], + [self._entry()], + ) + + def test_untargeted_metadata_for_rocm(self) -> None: + encoded = encode_cuda_aoti_metadata( + [CudaAotiVariant(0, 0, "rocm-so")], [self._entry()] + ) + self.assertTrue(encoded.startswith(CUDA_AOTI_METADATA_MAGIC)) + decoded = decode_cuda_aoti_metadata(encoded) + self.assertEqual(decoded.variants, [CudaAotiVariant(0, 0, "rocm-so")]) + self.assertEqual(decoded.entries, [self._entry()]) + + def test_untargeted_metadata_cannot_mix_with_targeted_variants(self) -> None: + with self.assertRaisesRegex(ValueError, "requires one non-fallback variant"): + encode_cuda_aoti_metadata( + [ + CudaAotiVariant(0, 0, "rocm-so"), + CudaAotiVariant(80, 0, "sm80-so"), + ], + [self._entry()], + ) + + def test_metadata_rejects_trailing_data(self) -> None: + encoded = encode_cuda_aoti_metadata( + [CudaAotiVariant(80, 80, "sm80-so")], [self._entry()] + ) + with self.assertRaisesRegex(ValueError, "trailing bytes"): + decode_cuda_aoti_metadata(encoded + b"\0") + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/cuda/tests/test_merge_ptes.py b/backends/cuda/tests/test_merge_ptes.py new file mode 100644 index 00000000000..7872f066321 --- /dev/null +++ b/backends/cuda/tests/test_merge_ptes.py @@ -0,0 +1,387 @@ +# 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. + +import hashlib +import io +import tempfile +import unittest +from contextlib import redirect_stdout +from pathlib import Path +from unittest.mock import patch + +from executorch.backends.cuda.cuda_weight_collector import ( + AOTI_DEVICE_TYPE_CUDA, + CudaAotiVariant, + CudaWeightEntry, + decode_cuda_aoti_metadata, + encode_cuda_aoti_metadata, +) +from executorch.backends.cuda.merge_ptes import ( + CudaPteInput, + main as merge_main, + merge_cuda_pte_files, +) +from executorch.exir._serialize._named_data_store import NamedDataStore +from executorch.exir._serialize._program import ( + deserialize_pte_binary, + PTEFile, + serialize_pte_binary, +) +from executorch.exir._serialize.data_serializer import DataEntry, DataPayload +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.schema import ( + BackendDelegate, + BackendDelegateDataReference, + BackendDelegateInlineData, + ContainerMetadata, + DataLocation, + ExecutionPlan, + Program, + SubsegmentOffsets, +) +from executorch.extension.flat_tensor.serialize.serialize import FlatTensorSerializer + + +@patch("executorch.backends.cuda.merge_ptes.torch.version.hip", None) +class TestMergeCudaPtes(unittest.TestCase): + def _write_artifact( + self, + directory: Path, + target_sm: int, + so_data: bytes, + weight_data: bytes, + *, + fqn: str = "model.weight", + weight_key: str = "cuda_fqn_weight:cuda:model.weight", + ptx_compute: int | None = None, + ) -> CudaPteInput: + if ptx_compute is None: + ptx_compute = target_sm + so_key = hashlib.sha256(so_data).hexdigest() + "_so_blob" + entry = CudaWeightEntry( + fqn=fqn, + storage_key=weight_key, + storage_nbytes=len(weight_data), + dtype=1, + device_type=AOTI_DEVICE_TYPE_CUDA, + storage_offset=0, + sizes=(len(weight_data),), + strides=(1,), + ) + metadata = encode_cuda_aoti_metadata( + [CudaAotiVariant(target_sm, ptx_compute, so_key)], [entry] + ) + compile_specs = [CompileSpec("method_name", b"forward")] + compile_specs.append( + CompileSpec("cuda_include_ptx", b"ON" if ptx_compute else b"OFF") + ) + delegate = BackendDelegate( + id="CudaBackend", + processed=BackendDelegateDataReference(DataLocation.INLINE, 0), + compile_specs=compile_specs, + ) + program = Program( + version=0, + execution_plan=[ + ExecutionPlan( + name="forward", + container_meta_type=ContainerMetadata("", ""), + values=[], + inputs=[], + outputs=[], + chains=[], + operators=[], + delegates=[delegate], + non_const_buffer_sizes=[], + ) + ], + constant_buffer=[], + backend_delegate_data=[BackendDelegateInlineData(metadata)], + segments=[], + constant_segment=SubsegmentOffsets(0, []), + ) + store = NamedDataStore() + store.add_named_data(so_key, so_data) + pte_path = directory / "model.pte" + with pte_path.open("wb") as output: + serialize_pte_binary( + PTEFile( + program=program, named_data=store.get_named_data_store_output() + ), + extract_delegate_segments=True, + ).write_to_file(output) + + ptd_path = directory / "aoti_cuda_blob.ptd" + serializer = FlatTensorSerializer() + with ptd_path.open("wb") as output: + serializer.serialize( + DataPayload( + buffers=[weight_data], + named_data={weight_key: DataEntry(0, 1, None)}, + ) + ).write_to_file(output) + return CudaPteInput( + pte_path=pte_path, + ptd_path=ptd_path, + ) + + def test_merges_variants_and_keeps_one_weight_manifest(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + sm80 = root / "sm80" + sm120 = root / "sm120" + sm80.mkdir() + sm120.mkdir() + inputs = [ + self._write_artifact(sm80, 80, b"sm80-so", b"shared-weight"), + self._write_artifact(sm120, 120, b"sm120-so", b"shared-weight"), + ] + + merged = deserialize_pte_binary(bytes(merge_cuda_pte_files(inputs))) + delegate = merged.program.execution_plan[0].delegates[0] + payload = merged.program.backend_delegate_data[ + delegate.processed.index + ].data + metadata = decode_cuda_aoti_metadata(payload) + self.assertEqual( + [variant.target_sm for variant in metadata.variants], [80, 120] + ) + self.assertEqual( + [variant.ptx_compute for variant in metadata.variants], [0, 0] + ) + self.assertFalse( + any(variant.fallback_only for variant in metadata.variants) + ) + self.assertEqual(len(metadata.entries), 1) + self.assertEqual(metadata.entries[0].fqn, "model.weight") + self.assertEqual( + set(merged.named_data.pte_data), + { + hashlib.sha256(b"sm80-so").hexdigest() + "_so_blob", + hashlib.sha256(b"sm120-so").hexdigest() + "_so_blob", + }, + ) + self.assertNotIn( + "cuda_include_ptx", + { + spec.key + for spec in merged.program.execution_plan[0] + .delegates[0] + .compile_specs + }, + ) + + def test_uses_ptx_only_from_explicit_fallback(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + sm80 = root / "sm80" + sm120 = root / "sm120" + fallback_dir = root / "fallback" + sm80.mkdir() + sm120.mkdir() + fallback_dir.mkdir() + inputs = [ + self._write_artifact(sm80, 80, b"sm80-so", b"weight"), + self._write_artifact( + sm120, + 120, + b"sm120-so", + b"weight", + ptx_compute=0, + ), + ] + fallback = self._write_artifact(fallback_dir, 75, b"fallback-so", b"weight") + + merged = deserialize_pte_binary( + bytes(merge_cuda_pte_files(inputs, fallback)) + ) + delegate = merged.program.execution_plan[0].delegates[0] + metadata = decode_cuda_aoti_metadata( + merged.program.backend_delegate_data[delegate.processed.index].data + ) + self.assertEqual( + [variant.target_sm for variant in metadata.variants], [80, 120, 75] + ) + self.assertEqual( + [variant.ptx_compute for variant in metadata.variants], [0, 0, 75] + ) + self.assertEqual( + [variant.fallback_only for variant in metadata.variants], + [False, False, True], + ) + + def test_preserves_no_ptx_fallback(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + sm80 = root / "sm80" + sm120 = root / "sm120" + sm80.mkdir() + sm120.mkdir() + inputs = [ + self._write_artifact(sm80, 80, b"sm80-so", b"weight", ptx_compute=0), + self._write_artifact( + sm120, + 120, + b"sm120-so", + b"weight", + ptx_compute=0, + ), + ] + + merged = deserialize_pte_binary(bytes(merge_cuda_pte_files(inputs))) + delegate = merged.program.execution_plan[0].delegates[0] + metadata = decode_cuda_aoti_metadata( + merged.program.backend_delegate_data[delegate.processed.index].data + ) + self.assertEqual( + [variant.ptx_compute for variant in metadata.variants], [0, 0] + ) + + def test_rejects_different_weight_content(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + sm80 = root / "sm80" + sm120 = root / "sm120" + sm80.mkdir() + sm120.mkdir() + inputs = [ + self._write_artifact(sm80, 80, b"sm80-so", b"first-weight"), + self._write_artifact(sm120, 120, b"sm120-so", b"other-weight"), + ] + with self.assertRaisesRegex(ValueError, "weight content differs"): + merge_cuda_pte_files(inputs) + + def test_library_local_weight_keys_are_normalized_to_base(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + sm80 = root / "sm80" + sm120 = root / "sm120" + sm80.mkdir() + sm120.mkdir() + inputs = [ + self._write_artifact( + sm80, + 80, + b"sm80-so", + b"constant", + fqn="_tensor_constant0", + weight_key="cuda_fqn_weight:cuda:sm80-so:_tensor_constant0", + ), + self._write_artifact( + sm120, + 120, + b"sm120-so", + b"constant", + fqn="_tensor_constant0", + weight_key="cuda_fqn_weight:cuda:sm120-so:_tensor_constant0", + ), + ] + + merged = deserialize_pte_binary(bytes(merge_cuda_pte_files(inputs))) + delegate = merged.program.execution_plan[0].delegates[0] + metadata = decode_cuda_aoti_metadata( + merged.program.backend_delegate_data[delegate.processed.index].data + ) + self.assertEqual( + metadata.entries[0].storage_key, + "cuda_fqn_weight:cuda:sm80-so:_tensor_constant0", + ) + + def test_rejects_duplicate_target_sm(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + first = root / "first" + second = root / "second" + first.mkdir() + second.mkdir() + inputs = [ + self._write_artifact(first, 80, b"first-so", b"weight"), + self._write_artifact(second, 80, b"second-so", b"weight"), + ] + with self.assertRaisesRegex(ValueError, "Duplicate CUDA target sm80"): + merge_cuda_pte_files(inputs) + + def test_rejects_fallback_without_ptx(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + regular_dir = root / "regular" + fallback_dir = root / "fallback" + regular_dir.mkdir() + fallback_dir.mkdir() + regular = self._write_artifact(regular_dir, 80, b"sm80-so", b"weight") + fallback = self._write_artifact( + fallback_dir, + 75, + b"fallback-so", + b"weight", + ptx_compute=0, + ) + + with self.assertRaisesRegex(ValueError, "exactly one PTX-capable"): + merge_cuda_pte_files([regular], fallback) + + def test_rejects_rocm(self) -> None: + with patch( + "executorch.backends.cuda.merge_ptes.torch.version.hip", "6.3" + ), self.assertRaisesRegex(RuntimeError, "only NVIDIA CUDA"): + merge_cuda_pte_files([]) + + def test_cli_writes_merged_pte_and_reuses_base_ptd(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + sm80 = root / "sm80" + sm120 = root / "sm120" + fallback_dir = root / "fallback" + output = root / "output" + sm80.mkdir() + sm120.mkdir() + fallback_dir.mkdir() + first = self._write_artifact(sm80, 80, b"sm80-so", b"weight") + second = self._write_artifact(sm120, 120, b"sm120-so", b"weight") + fallback = self._write_artifact(fallback_dir, 75, b"fallback-so", b"weight") + self.assertIsNotNone(first.ptd_path) + output_pte = output / "model.pte" + output_ptd = output / "aoti_cuda_blob.ptd" + + with patch( + "sys.argv", + [ + "merge_ptes", + "--input-pte", + str(first.pte_path), + "--input-pte", + str(second.pte_path), + "--input-ptd", + str(first.ptd_path), + "--input-ptd", + str(second.ptd_path), + "--fallback-pte", + str(fallback.pte_path), + "--fallback-ptd", + str(fallback.ptd_path), + "--output-pte", + str(output_pte), + "--output-ptd", + str(output_ptd), + ], + ), redirect_stdout(io.StringIO()) as stdout: + merge_main() + + self.assertTrue(output_pte.is_file()) + assert first.ptd_path is not None + self.assertEqual(output_ptd.read_bytes(), first.ptd_path.read_bytes()) + report = stdout.getvalue() + self.assertIn(f"cubin\tsm80\t{first.pte_path}", report) + self.assertIn(f"cubin\tsm120\t{second.pte_path}", report) + self.assertIn( + f"ptx-fallback\tcompute_75 (source sm75)\t" + f"{fallback.pte_path} [fallback]", + report, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/cuda/tests/test_missing_weights_blob.py b/backends/cuda/tests/test_missing_weights_blob.py index bbd85c29da4..410184ee238 100644 --- a/backends/cuda/tests/test_missing_weights_blob.py +++ b/backends/cuda/tests/test_missing_weights_blob.py @@ -14,10 +14,9 @@ Two payload shapes reach that code. A current export carries per-name weight metadata and goes through the weight cache, which reports a missing blob itself. -A library built before external weights carries only the two blob keys, newline -separated, and goes through the legacy path. That legacy path is the one the -check was added to, so it is the one this file exercises, by rewriting the -payload of a real export into the older shape in place. +The legacy external-weight format carries two blob keys, newline separated. +That path is the one the check was added to, so this file exercises it by +rewriting the payload of a real export into the older shape in place. """ import os @@ -27,7 +26,7 @@ import torch from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner -from executorch.backends.cuda.cuda_weight_collector import CUDA_WEIGHT_CACHE_MAGIC +from executorch.backends.cuda.cuda_weight_collector import decode_cuda_aoti_metadata from executorch.exir import to_edge_transform_and_lower from executorch.exir._serialize._program import deserialize_pte_binary from torch.export import export @@ -91,10 +90,9 @@ def _rewrite_payload_as_legacy(self, path: str) -> None: ) self.assertEqual(len(payloads), 1, "expected one CUDA delegate") payload = payloads[0] - self.assertTrue( - payload.startswith(CUDA_WEIGHT_CACHE_MAGIC), - "expected the weight metadata payload this rewrite consumes", - ) + metadata = decode_cuda_aoti_metadata(payload) + self.assertEqual(len(metadata.variants), 1, "expected one compiled variant") + so_key = metadata.variants[0].so_blob_key # The payload carries a content hash, so it occurs once. offset = raw.find(payload) @@ -103,14 +101,9 @@ def _rewrite_payload_as_legacy(self, path: str) -> None: raw.find(payload, offset + 1), -1, "payload is not unique in the file" ) - # Derived from the shared library key, so it carries the library's hash - # rather than the blob's and would not resolve even if a sidecar were - # supplied. That is fine here: the point is that an unresolvable key now - # fails the load rather than binding nothing. - marker = b"_so_blob" - end = payload.index(marker) + len(marker) - so_key = payload[:end].rsplit(b"\x00", 1)[-1].decode("utf-8") - weights_key = so_key.replace("_so_blob", "_weights_blob") + # Keep the absent key short enough for a payload with no weight entries. + weights_key = "missing_weights" + self.assertNotIn(weights_key.encode(), raw) # The two keys, then zeros to keep the payload its original length. The # runtime reads the blob key as a C string, so the filler is not part of the @@ -123,6 +116,15 @@ def _rewrite_payload_as_legacy(self, path: str) -> None: blob[offset : offset + len(payload)] = legacy with open(path, "wb") as f: f.write(bytes(blob)) + with open(path, "rb") as f: + rewritten = deserialize_pte_binary(f.read()).program + rewritten_payloads = [ + bytes(rewritten.backend_delegate_data[delegate.processed.index].data) + for plan in rewritten.execution_plan + for delegate in plan.delegates + if delegate.id == "CudaBackend" + ] + self.assertEqual(rewritten_payloads, [legacy]) def test_load_reports_not_found_when_blob_is_absent(self) -> None: from executorch.runtime import Runtime @@ -149,16 +151,6 @@ def test_load_reports_not_found_when_blob_is_absent(self) -> None: self._rewrite_payload_as_legacy(path) - # Without this the test would still pass if the rewrite stopped - # working, by exercising the weight cache path instead, which reports - # the same error number for the same program. - with open(path, "rb") as f: - self.assertNotIn( - CUDA_WEIGHT_CACHE_MAGIC, - f.read(), - "the rewrite left the metadata payload in place", - ) - # The blob is never supplied, so the load must fail. The runtime's # exception carries only the method name and the error number, so the # cause is asserted through the rewrite check above rather than here. @@ -171,9 +163,7 @@ def test_load_reports_not_found_when_blob_is_absent(self) -> None: def test_load_succeeds_when_the_model_has_no_constants(self) -> None: """A model with nothing to bind still loads without its weights blob. - The refusal must not fire on a model that has no constants, and that branch - has no other coverage. A model with no parameters or buffers already emits - the two-key payload this loader handles, so no rewrite is needed here. + Rewrite this export too, so it exercises the legacy constant-count check. """ with tempfile.TemporaryDirectory() as outdir: @@ -194,6 +184,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: with open(path, "wb") as f: lowered.to_executorch().write_to_file(f) + self._rewrite_payload_as_legacy(path) + from executorch.runtime import Runtime method = Runtime.get().load_program(path).load_method("forward") diff --git a/backends/cuda/tests/test_sort_shim.py b/backends/cuda/tests/test_sort_shim.py index fc5f870fc42..da1ee0c3e2c 100644 --- a/backends/cuda/tests/test_sort_shim.py +++ b/backends/cuda/tests/test_sort_shim.py @@ -37,7 +37,9 @@ _CUDA_FALLBACK_KERNELS = frozenset( { "at::_ops::_weight_int4pack_mm::call", + "aoti_torch_cuda__weight_int4pack_mm", "at::_ops::sort_stable::call", + "aoti_torch_cuda_sort_stable", "aoti_torch_cuda_randint_low_out", "executorch_cuda::int4_plain_mm", "aoti_torch_cuda_int4_plain_mm", diff --git a/backends/mediatek/BUCK b/backends/mediatek/BUCK index fa68ea9c2be..c194e3335a9 100644 --- a/backends/mediatek/BUCK +++ b/backends/mediatek/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target(_kind = runtime.python_library, name = "preprocess", srcs = [ diff --git a/backends/mediatek/_passes/BUCK b/backends/mediatek/_passes/BUCK index 931c337cf4e..4f8b78983b7 100644 --- a/backends/mediatek/_passes/BUCK +++ b/backends/mediatek/_passes/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target(_kind = runtime.python_library, name = "passes", srcs = [ diff --git a/backends/mediatek/quantizer/BUCK b/backends/mediatek/quantizer/BUCK index 847bfc547b2..1136a5c3888 100644 --- a/backends/mediatek/quantizer/BUCK +++ b/backends/mediatek/quantizer/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "quantizer", diff --git a/backends/mediatek/scripts/mtk_build.sh b/backends/mediatek/scripts/mtk_build.sh index d42e5f7e10a..6d5b013f9f3 100755 --- a/backends/mediatek/scripts/mtk_build.sh +++ b/backends/mediatek/scripts/mtk_build.sh @@ -35,7 +35,7 @@ cmake -DCMAKE_INSTALL_PREFIX="${build_dir}" \ -B"${build_dir}" # Build the project -cmake --build "${build_dir}" --target install --config Release -j5 +cmake --build "${build_dir}" --target install --config Release -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) # Switch back to the original directory cd - > /dev/null diff --git a/backends/mlx/CMakeLists.txt b/backends/mlx/CMakeLists.txt index 2a3d1546755..b58c82b93cc 100644 --- a/backends/mlx/CMakeLists.txt +++ b/backends/mlx/CMakeLists.txt @@ -167,8 +167,9 @@ endif() # with add_subdirectory would drop its whole project into ExecuTorch's # target/option namespace, which collides with shared deps MLX fetches (e.g. # nlohmann_json) and leaks MLX's MLX_BUILD_* options into our cache. The -# isolated scope runs MLX's FetchContent in its own namespace, so no collision -# and no submodule patching are needed. +# isolated scope runs MLX's FetchContent in its own namespace, so dependency +# collisions do not require patching. The platform packaging patches below are +# still applied before the external configure step. include(ExternalProject) set(_mlx_binary_dir ${CMAKE_CURRENT_BINARY_DIR}/mlx) @@ -185,37 +186,9 @@ message( # mismatch after an MLX bump fails loudly rather than silently no-op'ing. See # each patch file under patches/ for its rationale. set(_mlx_patches - ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_nax_jit_sdk_gate.patch - ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_qmm_splitk_bk_align.patch - ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_gather_mm_rhs_lda.patch ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_metal_sdk_per_platform.patch - ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_swiftpm_metallib_name.patch + ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_metal_remove_addrspace_compat.patch ) -# In a framework build the delegate is static, so MLX cannot find a colocated -# metallib and instead loads one from a SwiftPM resource bundle. SWIFTPM_BUNDLE -# is the bundle name and MLX_SWIFTPM_METALLIB_NAME the per-slice file inside it -# (one per slice, since a slice's metallib does not load on another). PLATFORM -# is set only for the Apple presets; a plain wheel build leaves it empty and -# keeps the colocated path. -set(_mlx_extra_cxx_flags "") -set(_mlx_cxx_flags_arg "") -if(PLATFORM) - if(PLATFORM STREQUAL "OS64") - set(_mlx_metallib_slice "ios") - elseif(PLATFORM STREQUAL "SIMULATORARM64") - set(_mlx_metallib_slice "ios-simulator") - else() - set(_mlx_metallib_slice "macos") - endif() - set(_mlx_extra_cxx_flags - "-DSWIFTPM_BUNDLE=\\\"executorch_backend_mlx_resources\\\" -DMLX_SWIFTPM_METALLIB_NAME=\\\"mlx-${_mlx_metallib_slice}\\\"" - ) - # Only override the sub-build's CXX flags when there is something to add. An - # empty -DCMAKE_CXX_FLAGS= on the command line beats the environment, so - # passing it unconditionally would silently drop a wheel build's CXXFLAGS for - # MLX only. - set(_mlx_cxx_flags_arg "-DCMAKE_CXX_FLAGS=${_mlx_extra_cxx_flags}") -endif() # Prefer the preset's per-slice DEPLOYMENT_TARGET; fall back to the toolchain # value. @@ -265,7 +238,6 @@ ExternalProject_Add( CMAKE_GENERATOR "Unix Makefiles" CMAKE_ARGS "-DCMAKE_BUILD_TYPE=${_mlx_build_type}" -DCMAKE_CXX_STANDARD=${CMAKE_CXX_STANDARD} - ${_mlx_cxx_flags_arg} # Feed the preset's per-slice DEPLOYMENT_TARGET in as # CMAKE_OSX_DEPLOYMENT_TARGET: the ios.toolchain does not reliably # carry it to this sub-build, and the shader-flag patch reads @@ -310,14 +282,12 @@ ExternalProject_Add( # ExternalProject stamps the patch step and BUILD_ALWAYS does not re-run it, so # a reused build directory whose MLX source was reset (patches reverted) would -# recompile an unpatched MLX and silently drop the iOS Metal SDK selection and -# the SwiftPM metallib name. Re-apply the patches on every build; apply.sh is -# idempotent (it reverse-checks each patch, skipping those already applied). -# DEPENDERS configure, not build: mlx_metal_sdk_per_platform.patch edits the -# sub-project's own CMake to pick the Metal SDK from PLATFORM, so it must land -# before the sub-configure runs, or the shaders are built against the macOS SDK -# for every slice. That does mean the sub-configure re-runs on each build, which -# is the price of the ordering. +# recompile an unpatched MLX and silently drop the iOS Metal SDK selection. +# Re-apply the patches on every build; apply.sh is idempotent (it reverse-checks +# each patch, skipping those already applied). DEPENDERS configure, not build: +# mlx_metal_sdk_per_platform.patch edits the sub-project's own CMake to select +# the Metal SDK, so it must land before the sub-configure runs. That means the +# sub-configure re-runs on each build, which is the price of the ordering. ExternalProject_Add_Step( mlx_external reapply_patches COMMAND bash ${CMAKE_CURRENT_SOURCE_DIR}/patches/apply.sh ${MLX_SOURCE_DIR} @@ -367,6 +337,21 @@ set(_mlx_backend__srcs ${CMAKE_CURRENT_SOURCE_DIR}/runtime/MLXBackend.cpp ${CMAKE_CURRENT_SOURCE_DIR}/runtime/mlx_mutable_state.cpp ) +if(EXECUTORCH_MLX_SWIFTPM_RESOURCES) + if(NOT APPLE) + message( + FATAL_ERROR "EXECUTORCH_MLX_SWIFTPM_RESOURCES requires an Apple target" + ) + endif() + enable_language(OBJCXX) + list(APPEND _mlx_backend__srcs + ${CMAKE_CURRENT_SOURCE_DIR}/runtime/SwiftPMMetallibPath.mm + ) + set_source_files_properties( + ${CMAKE_CURRENT_SOURCE_DIR}/runtime/SwiftPMMetallibPath.mm + PROPERTIES COMPILE_FLAGS "-fobjc-arc" + ) +endif() # Build the delegate as a shared library for the wheel so a C++ consumer can # link it, and keep it static everywhere else so no other build changes. @@ -425,6 +410,11 @@ target_include_directories( target_link_libraries( mlxdelegate PRIVATE mlx_schema extension_llm_cache $ ) +if(EXECUTORCH_MLX_SWIFTPM_RESOURCES) + target_compile_definitions( + mlxdelegate PRIVATE EXECUTORCH_MLX_SWIFTPM_RESOURCES + ) +endif() if(EXECUTORCH_BUILD_SHARED) set_target_properties( mlxdelegate PROPERTIES OUTPUT_NAME executorch_backend_mlx @@ -452,6 +442,14 @@ target_compile_options(mlxdelegate PRIVATE ${_common_compile_options}) # Core tensor headers carry pre-existing narrowing conversions that trip Xcode's # -Wshorten-64-to-32 -Werror; suppress it here as XNNPACK and abseil already do. target_compile_options(mlxdelegate PRIVATE -Wno-shorten-64-to-32) +# MLX headers use C++20 defaulted comparison operators while we build C++17. +# They normally arrive via -isystem from the imported mlx target, which +# suppresses the warning. But this backend also installs them for downstream +# consumers, and every -I is searched before any -isystem, so once an install +# has run they resolve out of ${CMAKE_INSTALL_PREFIX}/include instead and +# -Werror turns fatal. That is why an incremental rebuild used to need a +# cmake-out wipe: a clean tree has nothing installed yet. +target_compile_options(mlxdelegate PRIVATE -Wno-c++20-extensions) if(EXECUTORCH_MLX_ENABLE_SANITIZERS) target_link_options(mlxdelegate PRIVATE ${_mlx_sanitizer_link_options}) endif() diff --git a/backends/mlx/examples/llm/CMakeLists.txt b/backends/mlx/examples/llm/CMakeLists.txt index fa041b7129a..b8ef4981bce 100644 --- a/backends/mlx/examples/llm/CMakeLists.txt +++ b/backends/mlx/examples/llm/CMakeLists.txt @@ -35,18 +35,10 @@ executorch_target_link_options_shared_lib(executorch) set(gflags_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../../../third-party/gflags) find_package(gflags REQUIRED) -set(link_libraries - executorch - extension_module - extension_tensor - extension_llm_cache - extension_llm_runner - extension_llm_sampler - gflags -) +set(link_libraries executorch extension_llm_runner gflags) if(NOT TARGET mlxdelegate) - message(FATAL_ERROR "mlx_run_llm_hf requires the MLX backend (mlxdelegate)") + message(FATAL_ERROR "MLX LLM runners require the MLX backend (mlxdelegate)") endif() list(APPEND link_libraries mlxdelegate mlx) executorch_target_link_options_shared_lib(mlxdelegate) @@ -67,7 +59,27 @@ target_include_directories( ) target_link_libraries(mlx_run_llm_hf PUBLIC ${link_libraries}) +add_executable(mlx_run_llm_batched run_llm_batched.cpp) +target_include_directories( + mlx_run_llm_batched PUBLIC ${_common_include_directories} ${_json_include} + ${_flatbuffers_include} +) +target_link_libraries( + mlx_run_llm_batched + PUBLIC executorch + extension_llm_batching_module + extension_llm_runner + gflags + mlxdelegate + mlx + tokenizers::tokenizers +) +if(TARGET optimized_native_cpu_ops_lib) + target_link_libraries(mlx_run_llm_batched PUBLIC optimized_native_cpu_ops_lib) +endif() + # The copy helper is gated on EXECUTORCH_BUILD_MLX, which the installed config # does not set; reaching here means mlxdelegate exists. set(EXECUTORCH_BUILD_MLX ON) executorch_target_copy_mlx_metallib(mlx_run_llm_hf) +executorch_target_copy_mlx_metallib(mlx_run_llm_batched) diff --git a/backends/mlx/examples/llm/README.md b/backends/mlx/examples/llm/README.md index 21c6f72f548..8969a3be3ab 100644 --- a/backends/mlx/examples/llm/README.md +++ b/backends/mlx/examples/llm/README.md @@ -96,7 +96,7 @@ pip install -U "transformers @ git+https://github.com/huggingface/transformers.g | `--use-custom-sdpa` | `False` | Use MLX custom SDPA (`mlx::custom_sdpa`) | | `--use-custom-kv-cache` | `False` | Use MLX custom KV cache (`mlx::kv_cache_update`) | | `--use-offgraph-cache` | `False` | Use the off-graph KV cache (`kvcache::update_and_attend`); replaces the two flags above | -| `--prefill-chunk-size` | `512` | Max tokens per forward step. Bounds the traced `seq_len` dimension and is published as `get_prefill_chunk_size` for the runner. It is also the largest single cache write, so a ring layer is sized `window + chunk - 1`; it may not exceed the sliding window or the context length. Ignored on the optimum-executorch path, which owns its own `seq_len` bound | +| `--prefill-chunk-size` | `512` | Max tokens per forward step. Bounds the traced `seq_len` dimension and is published as `get_max_seq_len` for the runner. It is also the largest single cache write, so a ring layer is sized `window + chunk - 1`; it may not exceed the sliding window or the context length. Ignored on the optimum-executorch path, which owns its own `seq_len` bound | Off-graph exports keep no cache in the `.pte`, so the pybindings `run_llm_hf` cannot run them — use [`mlx_run_llm_hf`](#mlx_run_llm_hf-c) below, which builds @@ -231,7 +231,7 @@ default. | `--temperature` | `0` | Sampling temperature; 0 is greedy argmax, which is what makes two `.pte` files comparable | | `--chat` | `llama3` | Chat template: `llama3`, `gemma`, `gemma4`, or `0` to disable | | `--kv-max-capacity` | `0` | Off-graph: history the cache may hold. Setting it selects the off-graph path | -| `--kv-storage-dtype` | `bf16` | Off-graph: KV storage dtype (`bf16`, `fp16`, `fp32`) | +| `--kv-storage-dtype` | PTE activation dtype | Off-graph: optional KV storage override (`bf16`, `fp16`, `fp32`); defaults to the PTE's `get_activation_dtype`, which is required, so a `.pte` exported before this metadata must be re-exported | | `--kv-initial-capacity` | `-1` | Off-graph: starting pool size; grows by doubling up to capacity | | `--kv-windows` | *(model's own)* | Off-graph: attention pattern override, e.g. `512` | | `--interactive` | `false` | Multi-turn chat on stdin; off-graph only | diff --git a/backends/mlx/examples/llm/dflash/export.py b/backends/mlx/examples/llm/dflash/export.py index 72baca5cc45..fea62c55e89 100644 --- a/backends/mlx/examples/llm/dflash/export.py +++ b/backends/mlx/examples/llm/dflash/export.py @@ -15,7 +15,6 @@ from pathlib import Path import torch - from executorch.backends.mlx.examples.llm.dflash.adapters import get_adapter from executorch.backends.mlx.examples.llm.dflash.model import ( DFlashDraftModel, @@ -151,7 +150,6 @@ def main(): block_dim = Dim("block_len", min=2, max=block_size) import torch.fx.experimental._config as fx_config - from executorch.backends.mlx.examples.llm.dflash.cache import DFlashDraftKVCache class DFlashCachedDraftModel(torch.nn.Module): @@ -195,6 +193,7 @@ def forward(self, tokens, new_ctx, cache_position): from executorch.backends.mlx.examples.llm.export_llm_hf import ( build_hf_exported_program, + model_constant_methods, ) print( @@ -202,7 +201,7 @@ def forward(self, tokens, new_ctx, cache_position): f"and quant {qlinear}/{qembedding} g={qlinear_group_size}/{qembedding_group_size} " f"max_ctx_len {max_ctx_len} prefill_chunk_size {prefill_chunk_size}..." ) - target_exported, prefill_chunk_size = build_hf_exported_program( + target_exported, prefill_chunk_size, vocab_size = build_hf_exported_program( model_id=args.target_model, revision=None, max_ctx_len=max_ctx_len, @@ -226,8 +225,13 @@ def forward(self, tokens, new_ctx, cache_position): from executorch.exir.passes import MemoryPlanningPass constant_methods = { - "get_max_ctx_len": max_ctx_len, - "get_prefill_chunk_size": prefill_chunk_size, + **model_constant_methods( + max_context_len=max_ctx_len, + logits_to_keep="full", + activation_dtype=args.dtype, + vocab_size=vocab_size, + max_seq_len=prefill_chunk_size, + ), "get_max_block_len": block_size, "get_mask_token_id": draft_config.mask_token_id, } diff --git a/backends/mlx/examples/llm/dflash/run.py b/backends/mlx/examples/llm/dflash/run.py index 31d900b5730..660c95f87cd 100644 --- a/backends/mlx/examples/llm/dflash/run.py +++ b/backends/mlx/examples/llm/dflash/run.py @@ -24,7 +24,6 @@ import time import torch - from executorch.backends.mlx.examples.llm.runtime_meta import ( apply_chat_template, chunked_prefill, @@ -112,7 +111,7 @@ def resolve_limits(program, pte_path, n_draft_arg): max_ctx_len, prefill_chunk_size = read_model_limits(program) if max_ctx_len is None: raise ValueError( - f"{pte_path} publishes no get_max_ctx_len; re-export it with " + f"{pte_path} publishes no get_max_context_len; re-export it with " "dflash/export.py." ) if prefill_chunk_size is None: diff --git a/backends/mlx/examples/llm/export_llm_hf.py b/backends/mlx/examples/llm/export_llm_hf.py index 825fcb36365..25f2915a329 100644 --- a/backends/mlx/examples/llm/export_llm_hf.py +++ b/backends/mlx/examples/llm/export_llm_hf.py @@ -42,12 +42,43 @@ from typing import Optional import torch +from executorch.extension.llm.export.model_metadata import ( + model_vocab_size, + write_activation_dtype, + write_logits_to_keep_mode, + write_max_context_len, + write_max_seq_len, + write_vocab_size, +) FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" logging.basicConfig(level=logging.INFO, format=FORMAT) logger = logging.getLogger(__name__) +def model_constant_methods( + *, + max_context_len: int, + logits_to_keep: str, + activation_dtype: str, + vocab_size: int, + max_seq_len: Optional[int] = None, +) -> dict[str, int]: + """Build the full metadata set an MLX HF export publishes. + + ``max_seq_len`` is the largest single forward step this export traces + (serialized as ``get_max_seq_len``); ``max_context_len`` is the KV-cache + capacity (``get_max_context_len``). Composes the shared per-constant writers. + """ + return { + **write_max_context_len(max_context_len), + **write_vocab_size(vocab_size), + **write_activation_dtype(activation_dtype), + **write_logits_to_keep_mode(logits_to_keep), + **write_max_seq_len(max_seq_len), + } + + def resolve_prefill_chunk_size( prefill_chunk_size: Optional[int], max_ctx_len: int, @@ -108,6 +139,7 @@ def _export_with_optimum( dtype=dtype_str, max_seq_len=max_ctx_len, ) + vocab_size = model_vocab_size(exportable.model) from executorch.backends.mlx.llm.quantization import quantize_model_ @@ -137,7 +169,15 @@ def _export_with_optimum( # optimum drives its own torch.export call, so it owns the seq-len bound. constant_methods = dict(exportable.metadata) - constant_methods["get_max_ctx_len"] = max_ctx_len + constant_methods.update( + model_constant_methods( + max_context_len=max_ctx_len, + logits_to_keep="full", + activation_dtype=dtype, + vocab_size=vocab_size, + max_seq_len=max_ctx_len, + ) + ) edge_program = exir.to_edge_transform_and_lower( exported_progs, @@ -172,15 +212,18 @@ def build_hf_exported_program( qembedding_group_size: Optional[int] = None, tap_layers: Optional[list[int]] = None, prefill_chunk_size: Optional[int] = None, + logits_to_keep: str = "full", ): """Build the torch.export program for an HF model with custom MLX components. - Returns ``(exported_program, resolved_prefill_chunk_size)``. The resolved chunk + Returns ``(exported_program, resolved_prefill_chunk_size, vocab_size)``. The resolved chunk is the traced ``seq_len`` upper bound and is what callers should publish as - ``get_prefill_chunk_size``. + ``get_max_seq_len``. """ + from executorch.backends.mlx.llm.exportable import LogitsToKeepMode from transformers import AutoModelForCausalLM + logits_to_keep_mode = int(LogitsToKeepMode.from_value(logits_to_keep)) torch_dtype_map = { "fp32": torch.float32, "fp16": torch.float16, @@ -206,6 +249,7 @@ def build_hf_exported_program( if attn_implementation: load_kwargs["attn_implementation"] = attn_implementation model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs) + vocab_size = model_vocab_size(model) # Check if model uses sliding window attention. Multimodal configs like # Gemma 4 keep transformer attributes under text_config. @@ -240,6 +284,7 @@ def build_hf_exported_program( model=model, max_cache_len=effective_cache_len, tap_layers=tap_layers, + logits_to_keep_mode=logits_to_keep_mode, ) if use_custom_kv_cache: @@ -283,19 +328,30 @@ def build_hf_exported_program( # prefill_chunk_size is the largest single step: it bounds the traced seq_len # and sizes the ring buffer as window + chunk - 1. seq_len_dim = torch.export.Dim("seq_length_dim", max=prefill_chunk_size) + export_kwargs = { + "input_ids": example_input_ids, + "cache_position": example_cache_position, + } dynamic_shapes = { "input_ids": {1: seq_len_dim}, "cache_position": {0: seq_len_dim}, } + if logits_to_keep == "selected": + logits_example_length = min(seq_length, prefill_chunk_size) + export_kwargs["logits_to_keep"] = torch.arange( + logits_example_length, dtype=torch.int64 + ) + dynamic_shapes["logits_to_keep"] = ( + {0: torch.export.Dim("logits_to_keep_dim", min=1, max=prefill_chunk_size)} + if prefill_chunk_size > 1 + else None + ) with torch.no_grad(): exported_program = torch.export.export( exportable, args=(), - kwargs={ - "input_ids": example_input_ids, - "cache_position": example_cache_position, - }, + kwargs=export_kwargs, dynamic_shapes=dynamic_shapes, strict=True, ) @@ -304,7 +360,7 @@ def build_hf_exported_program( for sym, constraint in exported_program.range_constraints.items(): logger.info(f" Range constraint: {sym}: {constraint}") - return exported_program, prefill_chunk_size + return exported_program, prefill_chunk_size, vocab_size def _export_with_custom_components( @@ -322,6 +378,7 @@ def _export_with_custom_components( qembedding_group_size: Optional[int] = None, tap_layers: Optional[list[int]] = None, prefill_chunk_size: Optional[int] = None, + logits_to_keep: str = "full", ) -> None: """Export using direct HF model with custom MLX components.""" import executorch.exir as exir @@ -331,7 +388,7 @@ def _export_with_custom_components( from executorch.exir.capture._config import ExecutorchBackendConfig from executorch.exir.passes import MemoryPlanningPass - exported_program, prefill_chunk_size = build_hf_exported_program( + exported_program, prefill_chunk_size, vocab_size = build_hf_exported_program( model_id=model_id, revision=revision, max_ctx_len=max_ctx_len, @@ -345,6 +402,7 @@ def _export_with_custom_components( qembedding_group_size=qembedding_group_size, tap_layers=tap_layers, prefill_chunk_size=prefill_chunk_size, + logits_to_keep=logits_to_keep, ) logger.info("Delegating to MLX backend...") @@ -353,10 +411,13 @@ def _export_with_custom_components( _skip_dim_order=True, ) - constant_methods = { - "get_max_ctx_len": max_ctx_len, - "get_prefill_chunk_size": prefill_chunk_size, - } + constant_methods = model_constant_methods( + max_context_len=max_ctx_len, + logits_to_keep=logits_to_keep, + activation_dtype=dtype, + vocab_size=vocab_size, + max_seq_len=prefill_chunk_size, + ) edge_program = exir.to_edge_transform_and_lower( {"forward": exported_program}, @@ -389,6 +450,7 @@ def _export_with_offgraph_cache( qlinear_group_size: Optional[int] = None, qembedding_group_size: Optional[int] = None, prefill_chunk_size: Optional[int] = None, + logits_to_keep: str = "full", ) -> None: """Export using the off-graph KV cache op (kvcache::update_and_attend).""" import executorch.exir as exir @@ -422,6 +484,7 @@ def _export_with_offgraph_cache( if revision is not None: load_kwargs["revision"] = revision model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs) + vocab_size = model_vocab_size(model) model.eval() from executorch.backends.mlx.llm.quantization import quantize_model_ @@ -436,7 +499,7 @@ def _export_with_offgraph_cache( and not no_tie_word_embeddings, ) - exportable = OffGraphExportWrapper(model) + exportable = OffGraphExportWrapper(model, logits_to_keep) from executorch.backends.mlx.llm.cache import resolve_hf_cache_layout @@ -451,14 +514,21 @@ def _export_with_offgraph_cache( prefill_chunk_size, max_ctx_len, min(sliding) if sliding else None ) - kv_metadata = { - "get_n_caches": len(layer_types), - "get_kv_heads": torch.tensor(cache_kv_heads, dtype=torch.int32), - "get_head_dims": torch.tensor(cache_head_dims, dtype=torch.int32), - "get_windows": torch.tensor(cache_windows, dtype=torch.int32), - "get_prefill_chunk_size": prefill_chunk_size, - "get_max_ctx_len": max_ctx_len, - } + kv_metadata = model_constant_methods( + max_context_len=max_ctx_len, + logits_to_keep=logits_to_keep, + activation_dtype=dtype, + vocab_size=vocab_size, + max_seq_len=prefill_chunk_size, + ) + kv_metadata.update( + { + "get_n_caches": len(layer_types), + "get_kv_heads": torch.tensor(cache_kv_heads, dtype=torch.int32), + "get_head_dims": torch.tensor(cache_head_dims, dtype=torch.int32), + "get_windows": torch.tensor(cache_windows, dtype=torch.int32), + } + ) logger.info( f"KV cache layout: {len(layer_types)} caches, " f"{sum(1 for w in cache_windows if w)} sliding (window {sliding_window})" @@ -470,19 +540,30 @@ def _export_with_offgraph_cache( example_cache_position = torch.arange(seq_length, dtype=torch.long) seq_len_dim = torch.export.Dim("seq_length_dim", max=prefill_chunk_size) + export_kwargs = { + "input_ids": example_input_ids, + "cache_position": example_cache_position, + } dynamic_shapes = { "input_ids": {1: seq_len_dim}, "cache_position": {0: seq_len_dim}, } + if logits_to_keep == "selected": + logits_example_length = min(seq_length, prefill_chunk_size) + export_kwargs["logits_to_keep"] = torch.arange( + logits_example_length, dtype=torch.int64 + ) + dynamic_shapes["logits_to_keep"] = ( + {0: torch.export.Dim("logits_to_keep_dim", min=1, max=prefill_chunk_size)} + if prefill_chunk_size > 1 + else None + ) with torch.no_grad(): exported_program = torch.export.export( exportable, args=(), - kwargs={ - "input_ids": example_input_ids, - "cache_position": example_cache_position, - }, + kwargs=export_kwargs, dynamic_shapes=dynamic_shapes, strict=True, ) @@ -534,6 +615,7 @@ def export_llama_hf( qembedding_group_size: Optional[int] = None, tap_layers: Optional[list[int]] = None, prefill_chunk_size: Optional[int] = None, + logits_to_keep: str = "full", ) -> None: if use_offgraph_cache: if use_custom_sdpa or use_custom_kv_cache: @@ -554,12 +636,19 @@ def export_llama_hf( qlinear_group_size=qlinear_group_size, qembedding_group_size=qembedding_group_size, prefill_chunk_size=prefill_chunk_size, + logits_to_keep=logits_to_keep, ) - elif use_custom_sdpa or use_custom_kv_cache or tap_layers is not None: + elif ( + use_custom_sdpa + or use_custom_kv_cache + or tap_layers is not None + or logits_to_keep != "full" + ): logger.info( f"Using custom components: sdpa={use_custom_sdpa}, " f"kv_cache={use_custom_kv_cache}, tap_layers={tap_layers}, " - f"prefill_chunk_size={prefill_chunk_size}" + f"prefill_chunk_size={prefill_chunk_size}, " + f"logits_to_keep={logits_to_keep}" ) _export_with_custom_components( model_id=model_id, @@ -576,6 +665,7 @@ def export_llama_hf( qembedding_group_size=qembedding_group_size, tap_layers=tap_layers, prefill_chunk_size=prefill_chunk_size, + logits_to_keep=logits_to_keep, ) else: logger.info("Using optimum-executorch pipeline (no custom components)") @@ -628,6 +718,12 @@ def main(): help="Comma-separated layer indices whose hidden states are concatenated and returned alongside logits. E.g. '1,9,17,25,33'", ) parser.add_argument("--use-offgraph-cache", action="store_true", default=False) + parser.add_argument( + "--logits-to-keep", + choices=("full", "last", "selected"), + default="full", + help="Logits output: full sequence, last token, or runtime-selected positions.", + ) parser.add_argument( "--prefill-chunk-size", type=int, @@ -658,6 +754,7 @@ def main(): qembedding_group_size=args.qembedding_group_size, tap_layers=tap_layers, prefill_chunk_size=args.prefill_chunk_size, + logits_to_keep=args.logits_to_keep, ) diff --git a/backends/mlx/examples/llm/run_llm_batched.cpp b/backends/mlx/examples/llm/run_llm_batched.cpp new file mode 100644 index 00000000000..c7333b3a79c --- /dev/null +++ b/backends/mlx/examples/llm/run_llm_batched.cpp @@ -0,0 +1,426 @@ +/* + * 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. + */ + +// Sample application demonstrating continuous batching and streaming output +// for independent prompts submitted from one thread. +// +// Required flags are --pte and --tokenizer. Each remaining positional argument +// is a prompt, and generated text is streamed to +// _.txt. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +DEFINE_string(pte, "", "Path to the .pte exported with --use-offgraph-cache"); +DEFINE_string(tokenizer, "", "Path to a supported tokenizer file"); +DEFINE_string(out_prefix, "gen", "Output files are _.txt"); +DEFINE_int32(max_session_tokens, 2048, "Maximum tokens retained per session"); +DEFINE_string( + kv_storage_dtype, + "", + "Override KV storage dtype with bf16, fp16, or fp32. Defaults to the PTE " + "activation dtype, or bf16 when metadata is absent."); +DEFINE_int32( + kv_initial_capacity, + -1, + "Initial cache pool capacity; -1 keeps the cache default"); +DEFINE_int32(max_new_tokens, 128, "Maximum generated tokens per prompt"); +DEFINE_int32(flush_every, 8, "Flush each output file every N generated tokens"); +DEFINE_int32( + max_decode_sequences, + 32, + "Maximum decode sequences admitted to one batch"); +DEFINE_double(temperature, 0.0, "Sampling temperature; 0 is greedy"); +DEFINE_double(top_p, 1.0, "Nucleus sampling probability"); +DEFINE_int32(top_k, 0, "Top-k sampling limit; 0 disables it"); +DEFINE_uint64(seed, 42, "Per-generation sampling seed"); +DEFINE_bool(metrics, true, "Print per-generation and engine reports"); +DEFINE_string( + chat, + "llama3", + "Chat template: llama3, gemma, gemma4, or 0 for raw text"); + +namespace batching = ::executorch::extension::llm::batching; +using ::executorch::backends::mlx::examples::llm::resolve_kv_storage_dtype; +using ::executorch::backends::mlx::examples::llm::resolve_stop_tokens; +using ::executorch::backends::mlx::examples::llm::StopTokens; +using ::executorch::backends::mlx::examples::llm::wrap_turn; +using ::executorch::extension::Module; +using ::executorch::extension::llm::TextStream; +using ::executorch::runtime::Error; + +namespace { + +struct Emitter { + Emitter( + const tokenizers::Tokenizer& tokenizer, + batching::Token previous, + const std::string& path, + std::size_t flush_every) + : file(path, std::ios::binary), + flush_every(flush_every), + stream( + tokenizer, + [this](const std::string& piece) { + file.write( + piece.data(), static_cast(piece.size())); + }, + previous) {} + + void append(batching::Token token) { + if (stream.append(token) != Error::Ok || !file) { + throw std::runtime_error("failed to decode or write output"); + } + if (++tokens_since_flush == flush_every) { + file.flush(); + tokens_since_flush = 0; + if (!file) { + throw std::runtime_error("failed to flush output"); + } + } + } + + void finish() { + stream.flush(); + file.flush(); + if (!file) { + throw std::runtime_error("failed to flush output"); + } + } + + std::ofstream file; + const std::size_t flush_every; + std::size_t tokens_since_flush = 0; + TextStream stream; +}; + +struct JobResult { + std::optional session; + batching::GenerationHandle handle; + std::optional reason; + std::optional metrics; + std::string message; + std::string output_path; + + bool failed() const { + return !reason || *reason == batching::FinishReason::Cancelled || + *reason == batching::FinishReason::Failed; + } +}; + +const char* reason_name(const std::optional& reason) { + if (!reason) { + return "never started"; + } + switch (*reason) { + case batching::FinishReason::StopToken: + return "stop token"; + case batching::FinishReason::NewTokenLimit: + return "token limit"; + case batching::FinishReason::Cancelled: + return "cancelled"; + case batching::FinishReason::Failed: + return "failed"; + } + return "unknown"; +} + +void submit_prompt( + batching::Session& session, + const tokenizers::Tokenizer& tokenizer, + const std::string& prompt, + const std::vector& stop_tokens, + JobResult& result) { + try { + std::string wrapped; + if (!wrap_turn(FLAGS_chat, prompt, true, wrapped)) { + result.message = "unknown --chat template: " + FLAGS_chat; + return; + } + auto encoded = tokenizer.encode(wrapped, FLAGS_chat == "0" ? 1 : 0, 0); + if (!encoded.ok() || encoded->empty()) { + result.message = "could not encode prompt"; + return; + } + // Reserve the full generation budget so an admitted job is never shortened. + if (encoded->size() > + static_cast( + FLAGS_max_session_tokens - FLAGS_max_new_tokens)) { + result.message = + "prompt plus --max_new_tokens exceeds --max_session_tokens"; + return; + } + + auto emitter = std::make_shared( + tokenizer, + encoded->back(), + result.output_path, + static_cast(FLAGS_flush_every)); + if (!emitter->file) { + result.message = "could not open output file"; + return; + } + + batching::GenConfig config; + config.max_new_tokens = FLAGS_max_new_tokens; + config.sampling.temperature = static_cast(FLAGS_temperature); + config.sampling.top_p = static_cast(FLAGS_top_p); + config.sampling.top_k = FLAGS_top_k; + config.stop_tokens = stop_tokens; + config.seed = FLAGS_seed; + + result.handle = session.generate_async( + std::move(*encoded), + std::move(config), + [emitter](const batching::GenerationUpdate& update) { + std::size_t count = update.tokens.size(); + if (update.finish_reason == batching::FinishReason::StopToken && + count > 0) { + --count; + } + for (std::size_t i = 0; i < count; ++i) { + emitter->append(update.tokens[i]); + } + if (update.finish_reason) { + emitter->finish(); + } + }); + } catch (const std::exception& error) { + result.message = error.what(); + } +} + +} // namespace + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + const std::vector prompts(argv + 1, argv + argc); + + if (FLAGS_pte.empty() || FLAGS_tokenizer.empty() || prompts.empty()) { + std::cerr << "usage: " << argv[0] + << " --pte model.pte --tokenizer tokenizer-file \"prompt\" [...]" + << std::endl; + return 1; + } + if (FLAGS_max_session_tokens <= 0 || FLAGS_max_new_tokens <= 0 || + FLAGS_max_decode_sequences <= 0 || FLAGS_flush_every <= 0) { + std::cerr + << "session, generation, decode, and flush limits must be positive" + << std::endl; + return 1; + } + if (FLAGS_temperature < 0.0 || FLAGS_top_p <= 0.0 || FLAGS_top_p > 1.0 || + FLAGS_top_k < 0) { + std::cerr << "invalid sampling parameters" << std::endl; + return 1; + } + if (FLAGS_kv_initial_capacity < -1) { + std::cerr << "--kv_initial_capacity must be -1 or non-negative" + << std::endl; + return 1; + } + if (prompts.size() > + static_cast(std::numeric_limits::max())) { + std::cerr << "too many prompts" << std::endl; + return 1; + } + if (FLAGS_max_new_tokens > FLAGS_max_session_tokens) { + std::cerr << "--max_new_tokens exceeds --max_session_tokens" << std::endl; + return 1; + } + + auto tokenizer = + ::executorch::extension::llm::load_tokenizer(FLAGS_tokenizer); + if (!tokenizer) { + std::cerr << "could not load tokenizer: " << FLAGS_tokenizer << std::endl; + return 1; + } + + auto module = std::make_unique(FLAGS_pte); + if (module->load() != Error::Ok) { + std::cerr << "could not load " << FLAGS_pte << std::endl; + return 1; + } + + const auto activation_dtype = + ::executorch::extension::llm::read_activation_dtype(*module); + if (!activation_dtype.ok()) { + std::cerr << "could not read model metadata" << std::endl; + return 1; + } + const int kv_dtype = + resolve_kv_storage_dtype(FLAGS_kv_storage_dtype, *activation_dtype); + if (kv_dtype < 0) { + std::cerr << "--kv_storage_dtype must be bf16, fp16, or fp32" << std::endl; + return 1; + } + const auto max_context_length = + ::executorch::extension::llm::read_max_context_length(*module); + if (!max_context_length.ok()) { + std::cerr << "could not read model metadata" << std::endl; + return 1; + } + if (FLAGS_max_session_tokens > *max_context_length) { + std::cerr << "--max_session_tokens " << FLAGS_max_session_tokens + << " exceeds the model context limit " << *max_context_length + << std::endl; + return 1; + } + + StopTokens resolved_stop_tokens; + if (!resolve_stop_tokens( + *tokenizer, *module, FLAGS_chat, resolved_stop_tokens)) { + std::cerr << "could not resolve stop tokens for --chat=" << FLAGS_chat + << std::endl; + return 1; + } + const std::vector stop_tokens( + resolved_stop_tokens.ids.begin(), resolved_stop_tokens.ids.end()); + + auto executor = batching::ModuleExecutor::create( + std::move(module), + static_cast(prompts.size()), + FLAGS_max_session_tokens, + kv_dtype, + FLAGS_kv_initial_capacity); + if (!executor.ok()) { + std::cerr << "could not create executor: " + << ::executorch::runtime::to_string(executor.error()) + << std::endl; + return 1; + } + + const std::size_t width = (*executor)->preferred_batch_tokens(); + const std::size_t decode_slots = + static_cast(FLAGS_max_decode_sequences); + if (width == 0) { + std::cerr << "the model's forward token input has no usable width (its " + "traced seq_len dimension is 0); re-export it with a dynamic " + "token dimension" + << std::endl; + return 1; + } + if (decode_slots >= width) { + std::cerr << "--max_decode_sequences " << decode_slots + << " leaves no room for prefill in a " << width + << "-token forward" << std::endl; + return 1; + } + if (decode_slots > width / 4) { + std::cerr << "warning: --max_decode_sequences " << decode_slots + << " leaves only " << width - decode_slots + << " prefill tokens of a " << width << "-token forward" + << std::endl; + } + + auto scheduler = batching::DecodeFirstScheduler::create( + width, decode_slots, width - decode_slots); + if (!scheduler) { + std::cerr << "the scheduler refused those limits" << std::endl; + return 1; + } + + batching::Runner runner(**executor, std::move(scheduler)); + std::vector results(prompts.size()); + // Fail before opening sessions if any output path cannot be created. + for (std::size_t i = 0; i < results.size(); ++i) { + results[i].output_path = + FLAGS_out_prefix + "_" + std::to_string(i) + ".txt"; + std::ofstream output( + results[i].output_path, std::ios::binary | std::ios::trunc); + if (!output) { + runner.shutdown(); + std::cerr << "could not create " << results[i].output_path << std::endl; + return 1; + } + } + + std::vector>> session_futures; + session_futures.reserve(prompts.size()); + for (std::size_t i = 0; i < prompts.size(); ++i) { + session_futures.push_back(runner.open_session_async()); + } + for (std::size_t i = 0; i < prompts.size(); ++i) { + results[i].session = session_futures[i].get(); + } + for (std::size_t i = 0; i < prompts.size(); ++i) { + if (!results[i].session) { + results[i].message = "could not open session"; + continue; + } + submit_prompt( + *results[i].session, *tokenizer, prompts[i], stop_tokens, results[i]); + } + for (JobResult& result : results) { + if (!result.handle.valid()) { + continue; + } + result.handle.wait(); + result.metrics = result.handle.metrics(); + result.reason = result.handle.finish_reason(); + result.message = result.handle.error_message(); + } + + runner.shutdown(); + const batching::EngineMetrics engine = runner.metrics(); + + if (FLAGS_metrics) { + std::cout << "\n"; + for (std::size_t i = 0; i < results.size(); ++i) { + if (results[i].metrics) { + std::cout << "[" << i << "] " + << batching::format_report(*results[i].metrics); + } + } + std::cout << "\n" << batching::format_report(engine); + } + + std::size_t failures = 0; + for (const JobResult& result : results) { + failures += result.failed() ? 1 : 0; + } + if (failures > 0) { + std::cout << "\nfailures:\n"; + for (std::size_t i = 0; i < results.size(); ++i) { + if (!results[i].failed()) { + continue; + } + std::cout << " [" << i << "] " << results[i].output_path << ": " + << reason_name(results[i].reason); + if (!results[i].message.empty()) { + std::cout << ": " << results[i].message; + } + std::cout << "\n"; + } + } + + std::cout << prompts.size() - failures << "/" << prompts.size() + << " generations completed" << std::endl; + return failures == 0 ? 0 : 1; +} diff --git a/backends/mlx/examples/llm/run_llm_hf.cpp b/backends/mlx/examples/llm/run_llm_hf.cpp index 37ae7bc8921..4c6a55c875a 100644 --- a/backends/mlx/examples/llm/run_llm_hf.cpp +++ b/backends/mlx/examples/llm/run_llm_hf.cpp @@ -31,7 +31,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -42,14 +44,14 @@ #include #include -#include +#include + #include #include #include #include #include #include -#include #include DEFINE_string(pte, "", "Model .pte file."); @@ -79,8 +81,9 @@ DEFINE_int32( "only choose policy."); DEFINE_string( kv_storage_dtype, - "bf16", - "Off-graph: KV storage dtype, bf16|fp16|fp32."); + "", + "Off-graph: override KV storage dtype with bf16|fp16|fp32. Defaults to " + "the PTE activation dtype, or bf16 when metadata is absent."); DEFINE_int32( kv_initial_capacity, -1, @@ -100,11 +103,20 @@ DEFINE_bool( false, "Run once before measuring, to absorb JIT and pool growth."); +using ::executorch::backends::mlx::examples::llm::resolve_kv_storage_dtype; +using ::executorch::backends::mlx::examples::llm::resolve_stop_tokens; +using ::executorch::backends::mlx::examples::llm::StopTokens; +using ::executorch::backends::mlx::examples::llm::wrap_turn; using ::executorch::extension::make_tensor_ptr; using ::executorch::extension::Module; -using ::executorch::extension::TensorPtr; +using ::executorch::extension::llm::check_vocab_size; +using ::executorch::extension::llm::LogitsToKeepMode; +using ::executorch::extension::llm::read_activation_dtype; +using ::executorch::extension::llm::read_logits_to_keep_mode; +using ::executorch::extension::llm::read_max_seq_len; +using ::executorch::extension::llm::read_vocab_size; +using ::executorch::extension::llm::TextStream; using ::executorch::runtime::Error; -using ::executorch::runtime::EValue; namespace cache = ::executorch::extension::llm::cache; @@ -142,20 +154,6 @@ bool parse_int_list( return true; } -int storage_dtype(const std::string& name) { - using S = ::executorch::runtime::etensor::ScalarType; - if (name == "bf16") { - return static_cast(S::BFloat16); - } - if (name == "fp16") { - return static_cast(S::Half); - } - if (name == "fp32") { - return static_cast(S::Float); - } - return -1; -} - // Constant methods the export publishes (get_n_caches and friends). They carry // no delegate, so reading them only needs the program loaded -- which is what // lets the cache be built before forward's backend init consumes its key. @@ -167,6 +165,53 @@ std::optional const_int(Module& module, const char* name) { return r->at(0).toInt(); } +// The sampler (sample_from_logits) fatally aborts on any other dtype, so an +// unsupported logits type must be rejected at startup rather than at inference. +bool is_supported_logits_type(::executorch::aten::ScalarType type) { + using ScalarType = ::executorch::aten::ScalarType; + return type == ScalarType::Float || type == ScalarType::Half || + type == ScalarType::BFloat16 || type == ScalarType::UInt16; +} + +bool validate_forward_abi( + Module& module, + LogitsToKeepMode logits_to_keep_mode, + std::int64_t& vocab_size) { + const auto meta = module.method_meta("forward"); + if (!meta.ok()) { + std::cerr << "Forward metadata is unavailable" << std::endl; + return false; + } + // The runner feeds tokens + positions, plus a selector in Selected mode; a + // mismatch means the published logits mode disagrees with the traced graph. + const std::size_t expected_inputs = + logits_to_keep_mode == LogitsToKeepMode::Selected ? 3 : 2; + if (meta->num_inputs() != expected_inputs) { + std::cerr << "Forward must take " << expected_inputs + << " inputs for its logits-to-keep mode, got " + << meta->num_inputs() << std::endl; + return false; + } + // The logits output's last dim is the observed vocab width, cross-checked + // against the published get_vocab_size by the caller; its dtype must be one + // the sampler supports. + if (meta->num_outputs() == 0) { + std::cerr << "Forward publishes no logits output" << std::endl; + return false; + } + const auto logits = meta->output_tensor_meta(0); + if (!logits.ok() || logits->sizes().size() < 2 || + logits->sizes()[logits->sizes().size() - 1] <= 0 || + !is_supported_logits_type(logits->scalar_type())) { + std::cerr << "Forward logits must have a sampler-supported dtype and shape " + "[..., vocab]" + << std::endl; + return false; + } + vocab_size = logits->sizes()[logits->sizes().size() - 1]; + return true; +} + std::optional> const_ints(Module& module, const char* name) { const auto r = module.execute(name); if (!r.ok() || r->empty() || !r->at(0).isTensor()) { @@ -181,21 +226,21 @@ std::optional> const_ints(Module& module, const char* name) { } // Fill in the cache geometry the export published: get_n_caches, then one -// entry per cache in get_kv_heads / get_head_dims / get_windows (0 = flat), -// plus get_prefill_chunk_size, which the export validates against the sliding -// window and which becomes max_write -- the largest step the cache may see. +// entry per cache in get_kv_heads / get_head_dims / get_windows (0 = flat). // Capacity and dtype stay with the flags. False means this is not an off-graph // model. -bool read_kv_layout(Module& module, cache::CacheConfig& cfg) { +bool read_kv_layout( + Module& module, + int prefill_chunk, + cache::CacheConfig& cfg) { const auto n_caches = const_int(module, "get_n_caches"); const auto kv_heads = const_ints(module, "get_kv_heads"); const auto head_dims = const_ints(module, "get_head_dims"); const auto windows = const_ints(module, "get_windows"); - const auto chunk = const_int(module, "get_prefill_chunk_size"); - if (!n_caches || !kv_heads || !head_dims || !windows || !chunk) { + if (!n_caches || !kv_heads || !head_dims || !windows) { return false; } - cfg.max_write = static_cast(*chunk); + cfg.max_write = prefill_chunk; const size_t n = static_cast(*n_caches); if (kv_heads->size() != n || head_dims->size() != n || windows->size() != n) { return false; @@ -292,33 +337,6 @@ void print_cache_summary(const cache::CacheConfig& cfg) { std::cout << std::endl; } -// One user turn wrapped in the model's instruct template. Returns false for an -// unknown template name. The leading BOS belongs to the first turn only, so a -// continuing conversation passes with_bos=false. -bool wrap_turn( - const std::string& chat, - const std::string& prompt, - bool with_bos, - std::string& out) { - if (chat == "0") { - out = prompt; - } else if (chat == "llama3") { - out = std::string(with_bos ? "<|begin_of_text|>" : "") + - "<|start_header_id|>user<|end_header_id|>\n\n" + prompt + - "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"; - } else if (chat == "gemma") { - out = std::string(with_bos ? "" : "") + "user\n" + - prompt + "\nmodel\n"; - } else if (chat == "gemma4") { - // Gemma 4 renamed the turn markers; its own is also the eos. - out = std::string(with_bos ? "" : "") + "<|turn>user\n" + prompt + - "\n<|turn>model\n"; - } else { - return false; - } - return true; -} - } // namespace int main(int argc, char** argv) { @@ -341,6 +359,12 @@ int main(int argc, char** argv) { "[--kv-max-capacity N for off-graph models]\n"; return 1; } + if (warmup && kv_capacity <= 0) { + std::cerr << "--warmup requires an off-graph cache selected with " + "--kv_max_capacity" + << std::endl; + return 1; + } try { // The shared loader sniffs the format, so --tokenizer takes any of the @@ -351,19 +375,10 @@ int main(int argc, char** argv) { return 1; } - // Off-graph models (update_and_attend) need a cache bound via cache_key; - // in-graph models (mlx::kv_cache_update) don't -- omit --kv-max-capacity - // for those. session/options are outer-scoped: session must outlive the - // Module (it keeps the cache in the registry) and mlx_opts must outlive - // load_method() (the map holds a view into it). - std::optional session; + // Outer-scoped because mlx_opts must outlive load_method(): the map holds + // a view into it. ::executorch::runtime::BackendOptions<1> mlx_opts; ::executorch::runtime::LoadBackendOptionsMap options_map; - const bool off_graph = kv_capacity > 0; - // Tokens per prefill step, from the .pte. 0 means one step: an - // in-graph model publishes no chunk and has no ring to bound. - int prefill_chunk = 0; - // Load the program but not forward: the cache must exist before forward's // backend init reads its key, and the layout it needs is published by // constant methods in the same file. @@ -373,337 +388,413 @@ int main(int argc, char** argv) { std::cerr << "Failed to load " << pte << std::endl; return 1; } - - if (off_graph) { - cache::CacheConfig cfg{}; - cfg.capacity = kv_capacity; - cfg.kv_dtype = storage_dtype(kv_dtype); - if (cfg.kv_dtype < 0) { - std::cerr << "Invalid --kv-storage-dtype: " << kv_dtype - << " (bf16|fp16|fp32)" << std::endl; - return 1; - } - if (!read_kv_layout(module, cfg)) { - std::cerr << "No KV cache layout in " << pte - << "; re-export with --use-offgraph-cache" << std::endl; - return 1; - } - if (!kv_windows.empty() && !apply_window_override(kv_windows, cfg)) { - std::cerr << "Invalid --kv-windows: " << kv_windows << std::endl; - return 1; - } - if (!cache::valid(cfg)) { - std::cerr << "Invalid cache config" << std::endl; - return 1; - } - if (initial_capacity >= 0) { - cfg.initial_capacity = initial_capacity; - } - auto built = cache::CacheBuilderRegistry::global().build( - ::executorch::backends::mlx::kMLXBackendId, "seq", cfg); - if (!built.ok()) { - std::cerr << "Failed to build cache: " - << static_cast(built.error()) << std::endl; - return 1; - } - prefill_chunk = cfg.max_write ? *cfg.max_write : 0; - session.emplace(cache::make_unique_key(), built.get()); - - print_cache_summary(cfg); - if (mlx_opts.set_option( - ::executorch::backends::mlx::kCacheKeyKey, - session->key().c_str()) != Error::Ok || - options_map.set_options( - ::executorch::backends::mlx::kMLXBackendId, mlx_opts.view()) != - Error::Ok) { - std::cerr << "Failed to set cache_key option" << std::endl; - return 1; - } + const auto logits_to_keep_mode_result = read_logits_to_keep_mode(module); + if (!logits_to_keep_mode_result.ok()) { + std::cerr << "Invalid model metadata in " << pte << std::endl; + return 1; } - - if (module.load_method( - "forward", - /*planned_memory=*/nullptr, - /*event_tracer=*/nullptr, - off_graph ? &options_map : nullptr) != Error::Ok) { - std::cerr << "Failed to load forward" << std::endl; + const LogitsToKeepMode logits_to_keep_mode = *logits_to_keep_mode_result; + std::int64_t output_vocab_size = 0; + if (!validate_forward_abi(module, logits_to_keep_mode, output_vocab_size)) { return 1; } - // Timings reported at the end, in the shared runner's format. - ::executorch::extension::llm::Stats stats; - stats.model_load_start_ms = load_start_ms; - stats.model_load_end_ms = ::executorch::extension::llm::time_in_ms(); - - // Weights-only baseline, so the deltas below isolate the cache. - const double mem_at_load = ::mlx::core::get_active_memory() / 1048576.0; - std::cout << "[mem] after load : " << mem_at_load << " MiB" << std::endl; - - // Encode. HFTokenizer maps special-token markers in the string to their - // ids, so the template's <|...|> tokens encode correctly; it already - // carries <|begin_of_text|>, so pass bos=0 to avoid a doubled BOS. - std::string enc_input; - if (!wrap_turn(chat, prompt, /*with_bos=*/true, enc_input)) { - std::cerr << "Unknown --chat template: " << chat - << " (expected llama3, gemma, gemma4, or 0)" << std::endl; + const auto published_vocab_size = read_vocab_size(module); + if (!published_vocab_size.ok()) { + std::cerr << "Invalid get_vocab_size in " << pte << std::endl; return 1; } - // The template carries its own BOS, so only a raw prompt asks for one. - const int8_t bos = chat == "0" ? 1 : 0; - auto enc = tokenizer->encode(enc_input, bos, /*eos=*/0); - if (!enc.ok()) { - std::cerr << "Encode failed" << std::endl; + const auto vocab_size_result = + check_vocab_size(*published_vocab_size, output_vocab_size); + if (!vocab_size_result.ok()) { + std::cerr << "Invalid get_vocab_size for the forward output in " << pte + << std::endl; return 1; } - std::vector tokens = std::move(*enc); - const int prompt_len = static_cast(tokens.size()); - - // End-of-text from the model's metadata when it publishes any, else the - // tokenizer's. The turn-end token is ours: it depends on --chat, which the - // .pte knows nothing about. - std::unordered_set stop_ids = - ::executorch::extension::llm::get_eos_ids(tokenizer.get(), &module); - std::optional turn_end_id; - if (chat != "0") { - const char* turn_end = chat == "llama3" ? "<|eot_id|>" - : chat == "gemma4" ? "" - : ""; - if (auto eot = tokenizer->piece_to_id(turn_end); eot.ok()) { - turn_end_id = static_cast(*eot); - stop_ids.insert(*eot); - } + const std::int32_t vocab_size = *vocab_size_result; + const auto max_seq_len = read_max_seq_len(module); + if (!max_seq_len.ok()) { + std::cerr << "Invalid or missing get_max_seq_len in " << pte << std::endl; + return 1; } - auto is_stop = [&](int64_t t) { - for (uint64_t s : stop_ids) { - if (t == static_cast(s)) { - return true; - } - } - return false; - }; - - // One Sampler for the whole run, as the shared runner does: constructing - // one per token would reseed its RNG from the wall clock every time. Built - // on first use because the vocab size comes from the logits -- this export - // publishes no get_vocab_size. - std::optional<::executorch::extension::llm::Sampler> sampler; - - auto step = [&](const std::vector& ids, - const std::vector& pos) { - auto in = - make_tensor_ptr({1, (int)ids.size()}, std::vector(ids)); - auto cp = make_tensor_ptr({(int)pos.size()}, std::vector(pos)); - auto out = module.execute("forward", {in, cp}); - if (!out.ok()) { - throw std::runtime_error("execute failed"); - } - const auto& logits = out->at(0).toTensor(); - if (!sampler) { - sampler.emplace( - static_cast(logits.size(logits.dim() - 1)), temperature); - } - stats.on_sampling_begin(); - const int32_t tok = - ::executorch::extension::llm::sample_from_logits(logits, *sampler); - stats.on_sampling_end(); - return static_cast(tok); - }; - - // Prefill in chunks, so a ring layer holds window + chunk - 1 slots rather - // than growing with the prompt. Only the last chunk's token is kept; the - // earlier ones exist to place their K/V in the cache. - auto prefill = [&](const std::vector& ids, - const std::vector& pos) { - const size_t step_size = - prefill_chunk > 0 ? static_cast(prefill_chunk) : ids.size(); - int64_t next = 0; - for (size_t off = 0; off < ids.size(); off += step_size) { - const size_t n = std::min(step_size, ids.size() - off); - next = step( - {ids.begin() + off, ids.begin() + off + n}, - {pos.begin() + off, pos.begin() + off + n}); - } - return next; + const int prefill_chunk = static_cast(*max_seq_len); + StopTokens stop_tokens; + if (!resolve_stop_tokens(*tokenizer, module, chat, stop_tokens)) { + std::cerr << "Could not resolve stop tokens for --chat=" << chat + << std::endl; + return 1; + } + auto write_text = [](const std::string& text) { + std::cout << text << std::flush; }; - // Multi-turn: history stays in the cache, so each turn only prefills its - // own tokens at the running position. /reset and /undo drive the cache's - // control face directly -- off-graph only, since an in-graph cache gives - // the runner no handle to its state. - if (interactive) { - if (!off_graph) { - std::cerr << "--interactive requires --kv-max-capacity\n"; + // Everything past load_method is identical for both model kinds; only + // setup differs. ctl is null for an in-graph model, which owns its cache + // inside the graph and exposes no control face. + auto run = + [&](cache::SequenceControl* ctl, + const ::executorch::runtime::LoadBackendOptionsMap* load_opts, + int run_prefill_chunk) -> int { + if (module.load_method( + "forward", + /*planned_memory=*/nullptr, + /*event_tracer=*/nullptr, + load_opts) != Error::Ok) { + std::cerr << "Failed to load forward" << std::endl; return 1; } - auto* control = session->control(); - std::cout << "Multi-turn chat. /reset clears, /undo drops the last turn, " - "/undo N drops N tokens, /quit exits.\n"; - int64_t position = 0; - int64_t turn_start = 0; // position this turn began at, for /undo - std::string line; - while (std::cout << "\n> " && std::getline(std::cin, line)) { - if (line == "/quit") { - break; + // Timings reported at the end, in the shared runner's format. + ::executorch::extension::llm::Stats stats; + stats.model_load_start_ms = load_start_ms; + stats.model_load_end_ms = ::executorch::extension::llm::time_in_ms(); + + // Weights-only baseline, so the deltas below isolate the cache. + const double mem_at_load = ::mlx::core::get_active_memory() / 1048576.0; + std::cout << "[mem] after load : " << mem_at_load << " MiB" + << std::endl; + + auto is_stop = [&](int64_t token) { + return stop_tokens.ids.count(static_cast(token)) != 0; + }; + + // One Sampler for the whole run, as the shared runner does: constructing + // one per token would reseed its RNG from the wall clock every time. + // Built on first use because the vocab size comes from the logits -- this + // export publishes no get_vocab_size. + std::optional<::executorch::extension::llm::Sampler> sampler; + + auto step = [&](const std::vector& ids, + const std::vector& pos) { + auto in = + make_tensor_ptr({1, (int)ids.size()}, std::vector(ids)); + auto cp = make_tensor_ptr({(int)pos.size()}, std::vector(pos)); + auto out = [&]() -> ::executorch::runtime::Result< + std::vector<::executorch::runtime::EValue>> { + if (logits_to_keep_mode == LogitsToKeepMode::Selected) { + auto selector = make_tensor_ptr( + {1}, + std::vector{static_cast(ids.size() - 1)}); + return module.execute("forward", {in, cp, selector}); + } + return module.execute("forward", {in, cp}); + }(); + if (!out.ok()) { + throw std::runtime_error("execute failed"); + } + if (out->empty() || !out->at(0).isTensor()) { + throw std::runtime_error("forward returned no logits"); + } + const auto& logits = out->at(0).toTensor(); + const int64_t actual_vocab_size = logits.dim() == 0 + ? 0 + : static_cast(logits.size(logits.dim() - 1)); + const int64_t expected_rows = + logits_to_keep_mode == LogitsToKeepMode::Full + ? static_cast(ids.size()) + : 1; + if (logits.dim() != 3 || logits.size(0) != 1 || + logits.size(1) != expected_rows || + actual_vocab_size != vocab_size) { + throw std::runtime_error("forward returned an invalid logits shape"); } - if (line == "/reset") { - control->clear(); - position = turn_start = 0; - std::cout << "[cleared]\n"; - continue; + if (!sampler) { + sampler.emplace(vocab_size, temperature); } - if (line == "/undo" || line.rfind("/undo ", 0) == 0) { - // Bare /undo drops the last turn; /undo N drops N tokens. - int64_t target = turn_start; - if (line.size() > 6) { - try { - const int64_t n = std::stoll(line.substr(6)); - target = n >= position ? 0 : position - n; - } catch (const std::exception&) { - std::cout << "[usage: /undo [n_tokens]]\n"; - continue; + stats.on_sampling_begin(); + const int32_t tok = + ::executorch::extension::llm::sample_from_logits(logits, *sampler); + stats.on_sampling_end(); + return static_cast(tok); + }; + + // Prefill in chunks, so a ring layer holds window + chunk - 1 slots + // rather than growing with the prompt. Only the last chunk's token is + // kept; the earlier ones exist to place their K/V in the cache. + auto prefill = [&](const std::vector& ids, + const std::vector& pos) { + const size_t step_size = static_cast(run_prefill_chunk); + int64_t next = 0; + for (size_t off = 0; off < ids.size(); off += step_size) { + const size_t n = std::min(step_size, ids.size() - off); + next = step( + {ids.begin() + off, ids.begin() + off + n}, + {pos.begin() + off, pos.begin() + off + n}); + } + return next; + }; + + // Multi-turn: history stays in the cache, so each turn only prefills its + // own tokens at the running position. /reset and /undo drive the cache's + // control face directly -- off-graph only, since an in-graph cache gives + // the runner no handle to its state. + if (interactive) { + if (ctl == nullptr) { + std::cerr << "--interactive requires --kv-max-capacity\n"; + return 1; + } + std::cout + << "Multi-turn chat. /reset clears, /undo drops the last turn, " + "/undo N drops N tokens, /quit exits.\n"; + int64_t position = 0; + int64_t turn_start = 0; // position this turn began at, for /undo + std::string line; + while (std::cout << "\n> " && std::getline(std::cin, line)) { + if (line == "/quit") { + break; + } + if (line == "/reset") { + ctl->clear(); + position = turn_start = 0; + std::cout << "[cleared]\n"; + continue; + } + if (line == "/undo" || line.rfind("/undo ", 0) == 0) { + // Bare /undo drops the last turn; /undo N drops N tokens. + int64_t target = turn_start; + if (line.size() > 6) { + try { + const int64_t n = std::stoll(line.substr(6)); + target = n >= position ? 0 : position - n; + } catch (const std::exception&) { + std::cout << "[usage: /undo [n_tokens]]\n"; + continue; + } } + if (ctl->rewind(static_cast(target))) { + position = target; + turn_start = std::min(turn_start, position); + std::cout << "[rewound to " << position << "]\n"; + } else { + // A sliding-window layer has physically dropped those cells. + std::cout << "[cannot rewind to " << target << "]\n"; + } + continue; } - if (control->rewind(static_cast(target))) { - position = target; - turn_start = std::min(turn_start, position); - std::cout << "[rewound to " << position << "]\n"; - } else { - // A sliding-window layer has physically dropped those cells. - std::cout << "[cannot rewind to " << target << "]\n"; + if (line.empty()) { + continue; } - continue; - } - if (line.empty()) { - continue; + + std::string turn; + wrap_turn(chat, line, /*with_bos=*/position == 0, turn); + auto te = tokenizer->encode(turn, /*bos=*/chat == "0" ? 1 : 0, 0); + if (!te.ok() || te->empty()) { + std::cerr << "Encode failed or produced no tokens\n"; + continue; + } + const int n = static_cast(te->size()); + // Admit the turn if its prompt plus one token fits; reserving the + // whole max_new budget up front would report "full" with most of the + // cache still free. Generation is then clamped to the room that + // remains. + if (!ctl->can_extend(n + 1)) { + std::cout << "[cache full: " << position << "/" << ctl->capacity() + << ", turn " << n << " tokens" + << (ctl->can_extend(1) ? "" : ", length at capacity") + << ", use /reset]\n"; + continue; + } + const int budget = std::min( + max_new, ctl->capacity() - static_cast(position) - n); + + turn_start = position; + std::vector tin(te->begin(), te->end()), tpos; + for (int i = 0; i < n; ++i) { + tpos.push_back(position + i); + } + int64_t next = prefill(tin, tpos); + position += n; + + TextStream text_stream(*tokenizer, write_text, te->back()); + for (int i = 0; i < budget && !is_stop(next); ++i) { + if (text_stream.append(static_cast(next)) != Error::Ok) { + text_stream.flush(); + std::cerr << "Failed to decode generated token" << std::endl; + return 1; + } + next = step({next}, {position}); + ++position; + } + text_stream.flush(); + // The turn-end token stops generation, so it is neither printed nor + // fed back -- but the next turn opens without closing this one, and + // an unterminated assistant turn compounds over a session. Commit it, + // at the cost of one extra step per turn. + if (stop_tokens.turn_end_id && + static_cast(next) == *stop_tokens.turn_end_id && + ctl->can_extend(1)) { + step({next}, {position}); + ++position; + } + std::cout << "\n[" << position << "/" << ctl->capacity() << " tokens" + << (budget < max_new ? ", generation capped by capacity" + : "") + << "]\n"; } + return 0; + } - std::string turn; - wrap_turn(chat, line, /*with_bos=*/position == 0, turn); - auto te = tokenizer->encode(turn, /*bos=*/chat == "0" ? 1 : 0, 0); - if (!te.ok()) { - std::cerr << "Encode failed\n"; - continue; + std::string enc_input; + if (!wrap_turn(chat, prompt, /*with_bos=*/true, enc_input)) { + std::cerr << "Unknown --chat template: " << chat + << " (expected llama3, gemma, gemma4, or 0)" << std::endl; + return 1; + } + const int8_t bos = chat == "0" ? 1 : 0; + auto enc = tokenizer->encode(enc_input, bos, /*eos=*/0); + if (!enc.ok() || enc->empty()) { + std::cerr << "Encode failed or produced no tokens" << std::endl; + return 1; + } + std::vector tokens = std::move(*enc); + const int prompt_len = static_cast(tokens.size()); + std::vector ids(tokens.begin(), tokens.end()), prefill_pos; + for (int i = 0; i < prompt_len; ++i) { + prefill_pos.push_back(i); + } + // Sequence length against the configured ceiling, with what MLX actually + // holds for it. Pools start at initial_capacity and grow by doubling, so + // the bytes lag the token count in steps; bf16 storage (kv_dtype 15) + // halves them vs fp32 (6). + auto print_footprint = [&](const char* when, int len) { + if (ctl == nullptr) { + return; } - const int n = static_cast(te->size()); - // Admit the turn if its prompt plus one token fits; reserving the whole - // max_new budget up front would report "full" with most of the cache - // still free. Generation is then clamped to the room that remains. - if (!control->can_extend(n + 1)) { - std::cout << "[cache full: " << position << "/" << control->capacity() - << ", turn " << n << " tokens" - << (control->can_extend(1) ? "" : ", length at capacity") - << ", use /reset]\n"; - continue; + const int cap = ctl->capacity(); + const double pct = cap > 0 ? 100.0 * len / cap : 0.0; + std::cout << "[cache] " << when << ": " << len << " / " << cap + << " tokens (" << pct << "%)" << std::endl; + const double mem = ::mlx::core::get_active_memory() / 1048576.0; + std::cout << "[mem] " << when << ": " << mem << " MiB (+" + << (mem - mem_at_load) << " MiB since load)" << std::endl; + }; + + // One optional warmup run to absorb JIT and pool growth, then one + // measured run, as the shared LLM runners do. Repeats belong in a harness + // that restarts the process: clear() rewinds the sequence but leaves the + // pools at their grown size, so an in-process repeat cannot see + // reallocation. + for (int iter = 0; iter < (warmup ? 2 : 1); ++iter) { + const bool measured = !warmup || iter == 1; + if (iter > 0 && ctl != nullptr) { + ctl->clear(); + } + stats.inference_start_ms = ::executorch::extension::llm::time_in_ms(); + int64_t next = prefill(ids, prefill_pos); + stats.prompt_eval_end_ms = ::executorch::extension::llm::time_in_ms(); + // prefill returns the first generated token, so TTFT ends with prefill + stats.first_token_ms = stats.prompt_eval_end_ms; + if (measured) { + std::cout << "\n"; + print_footprint("after prefill", prompt_len); + std::cout << "\n"; // blank line before the streamed generation } - const int budget = std::min( - max_new, control->capacity() - static_cast(position) - n); - turn_start = position; - std::vector tin(te->begin(), te->end()), tpos; - for (int i = 0; i < n; ++i) { - tpos.push_back(position + i); + TextStream::Sink sink; + if (measured) { + sink = write_text; } - int64_t next = prefill(tin, tpos); - position += n; - - uint64_t prev = te->back(); - for (int i = 0; i < budget && !is_stop(next); ++i) { - if (auto piece = tokenizer->decode(prev, static_cast(next)); - piece.ok()) { - std::cout << *piece << std::flush; + TextStream text_stream(*tokenizer, std::move(sink), tokens.back()); + int generated = 0; + for (int i = 0; i < max_new; ++i) { + if (is_stop(next)) { + break; + } + if (text_stream.append(static_cast(next)) != Error::Ok) { + text_stream.flush(); + std::cerr << "Failed to decode generated token" << std::endl; + return 1; } - prev = static_cast(next); - next = step({next}, {position}); - ++position; + ++generated; + next = step({next}, {prompt_len + i}); } - // The turn-end token stops generation, so it is neither printed nor - // fed back -- but the next turn opens without closing this one, and an - // unterminated assistant turn compounds over a session. Commit it, at - // the cost of one extra step per turn. - if (turn_end_id && next == *turn_end_id && control->can_extend(1)) { - step({next}, {position}); - ++position; + text_stream.flush(); + stats.inference_end_ms = ::executorch::extension::llm::time_in_ms(); + if (measured) { + std::cout << "\n\n"; // close the generation line + blank separator + // trailing space aligns the colon with the "after prefill" line above + print_footprint("after decode ", prompt_len + generated); + stats.num_prompt_tokens = prompt_len; + stats.num_generated_tokens = generated; } - std::cout << "\n[" << position << "/" << control->capacity() - << " tokens" - << (budget < max_new ? ", generation capped by capacity" : "") - << "]\n"; } + std::cout << std::endl; + ::executorch::extension::llm::print_report(stats); return 0; + }; + + // An in-graph model (mlx::kv_cache_update) binds no cache: nothing to + // build, no key to hand the delegate, and so no registry entry to guard. + if (kv_capacity <= 0) { + return run( + /*ctl=*/nullptr, + /*load_opts=*/nullptr, + /*run_prefill_chunk=*/prefill_chunk); } - std::vector ids(tokens.begin(), tokens.end()), prefill_pos; - for (int i = 0; i < prompt_len; ++i) { - prefill_pos.push_back(i); + const auto activation_dtype = read_activation_dtype(module); + if (!activation_dtype.ok()) { + std::cerr << "Invalid get_activation_dtype in " << pte << std::endl; + return 1; + } + cache::CacheConfig cfg{}; + cfg.capacity = kv_capacity; + cfg.kv_dtype = resolve_kv_storage_dtype(kv_dtype, *activation_dtype); + if (cfg.kv_dtype < 0) { + std::cerr << "Invalid --kv-storage-dtype override: " << kv_dtype + << " (bf16|fp16|fp32)" << std::endl; + return 1; + } + if (!read_kv_layout(module, prefill_chunk, cfg)) { + std::cerr << "No KV cache layout in " << pte + << "; re-export with --use-offgraph-cache" << std::endl; + return 1; + } + if (!kv_windows.empty() && !apply_window_override(kv_windows, cfg)) { + std::cerr << "Invalid --kv-windows: " << kv_windows << std::endl; + return 1; + } + if (!cache::valid(cfg)) { + std::cerr << "Invalid cache config" << std::endl; + return 1; + } + if (initial_capacity >= 0) { + cfg.initial_capacity = initial_capacity; } - auto ms = [](auto a, auto b) { - return std::chrono::duration(b - a).count(); - }; - // Sequence length against the configured ceiling, with what MLX actually - // holds for it. Pools start at initial_capacity and grow by doubling, so - // the bytes lag the token count in steps; bf16 storage (kv_dtype 15) halves - // them vs fp32 (6). - auto print_footprint = [&](const char* when, int len) { - if (!session) { - return; - } - const int cap = session->control()->capacity(); - const double pct = cap > 0 ? 100.0 * len / cap : 0.0; - std::cout << "[cache] " << when << ": " << len << " / " << cap - << " tokens (" << pct << "%)" << std::endl; - const double mem = ::mlx::core::get_active_memory() / 1048576.0; - std::cout << "[mem] " << when << ": " << mem << " MiB (+" - << (mem - mem_at_load) << " MiB since load)" << std::endl; - }; - // One optional warmup run to absorb JIT and pool growth, then one measured - // run, as the shared LLM runners do. Repeats belong in a harness that - // restarts the process: clear() rewinds the sequence but leaves the pools - // at their grown size, so an in-process repeat cannot see reallocation. - for (int iter = 0; iter < (warmup ? 2 : 1); ++iter) { - const bool measured = !warmup || iter == 1; - if (iter > 0 && off_graph) { - session->control()->clear(); - } - stats.inference_start_ms = ::executorch::extension::llm::time_in_ms(); - int64_t next = prefill(ids, prefill_pos); - stats.prompt_eval_end_ms = ::executorch::extension::llm::time_in_ms(); - // prefill returns the first generated token, so TTFT ends with prefill - stats.first_token_ms = stats.prompt_eval_end_ms; - if (measured) { - std::cout << "\n"; - print_footprint("after prefill", prompt_len); - std::cout << "\n"; // blank line before the streamed generation - } + const char* const cache_kind = cache::kind::kSingle; + auto built = cache::CacheFactory::global().build( + ::executorch::backends::mlx::kMLXBackendId, cache_kind, cfg); + if (!built.ok()) { + std::cerr << "Failed to build cache: " << static_cast(built.error()) + << std::endl; + return 1; + } + const std::shared_ptr kv = built.get(); + + // Published for the delegate to find by key, and erased when this scope + // exits. That is after run() returns, so the entry is still there for the + // load_method() inside it. + const cache::InstallGuard guard{kv}; + + print_cache_summary(cfg); + if (guard.set_option(mlx_opts) != Error::Ok || + options_map.set_options( + ::executorch::backends::mlx::kMLXBackendId, mlx_opts.view()) != + Error::Ok) { + std::cerr << "Failed to set cache_key option" << std::endl; + return 1; + } - uint64_t prev = tokens.back(); - int generated = 0; - for (int i = 0; i < max_new; ++i) { - if (is_stop(next)) { - break; - } - if (measured) { - if (auto piece = tokenizer->decode(prev, static_cast(next)); - piece.ok()) { - ::executorch::extension::llm::safe_printf(piece->c_str()); - fflush(stdout); - } - } - prev = static_cast(next); - ++generated; - next = step({next}, {prompt_len + i}); - } - stats.inference_end_ms = ::executorch::extension::llm::time_in_ms(); - if (measured) { - std::cout << "\n\n"; // close the generation line + blank separator - // trailing space aligns the colon with the "after prefill" line above - print_footprint("after decode ", prompt_len + generated); - stats.num_prompt_tokens = prompt_len; - stats.num_generated_tokens = generated; - } + // Checked here so a null ctl inside run() can only mean "in-graph model". + // A cache kind that offers BatchControl instead would otherwise be run as + // if it had no cache at all, with a key published and options set. + auto* ctl = kv->as(); + if (ctl == nullptr) { + std::cerr << "Cache kind '" << cache_kind + << "' offers no single-sequence control face" << std::endl; + return 1; } - std::cout << std::endl; - ::executorch::extension::llm::print_report(stats); - return 0; + + return run(ctl, &options_map, *cfg.max_write); } catch (const std::exception& e) { std::cerr << "Error: " << e.what() << std::endl; return 1; diff --git a/backends/mlx/examples/llm/run_llm_hf.py b/backends/mlx/examples/llm/run_llm_hf.py index e73da846623..79081884c2c 100644 --- a/backends/mlx/examples/llm/run_llm_hf.py +++ b/backends/mlx/examples/llm/run_llm_hf.py @@ -25,14 +25,18 @@ import time import torch - from executorch.backends.mlx.examples.llm.runtime_meta import ( apply_chat_template, chunked_prefill, get_eos_token_ids, load_text_processor, + read_const_int, read_model_limits, ) +from executorch.extension.llm.export.model_metadata import ( + LOGITS_TO_KEEP_MODE_METHOD, + LOGITS_TO_KEEP_MODES, +) from executorch.runtime import Runtime, Verification FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" @@ -40,11 +44,12 @@ logger = logging.getLogger(__name__) -def _get_max_input_seq_len(program) -> int: - """Inspect the .pte program metadata to determine the max input_ids seq len. +def _forward_input_seq_len(program) -> int: + """The forward's traced token-input width -- what set_inputs will accept. - Fallback for .pte files exported before get_prefill_chunk_size existed. - Returns the static seq-len dimension of the first input tensor (input_ids). + 1 for a static token-by-token export (e.g. optimum's static cache), or the + dynamic upper bound for a chunked-prefill export. This is authoritative: + feeding more tokens than this per step fails set_inputs. """ meta = program.metadata("forward") input_ids_info = meta.input_tensor_meta(0) @@ -67,9 +72,31 @@ def run_inference( et_runtime = Runtime.get() program = et_runtime.load_program(pte_path, verification=Verification.Minimal) - max_ctx_len, prefill_chunk_size = read_model_limits(program) - if prefill_chunk_size is None: - prefill_chunk_size = _get_max_input_seq_len(program) + # This pybindings runner only feeds tokens and positions. A model exported + # with --logits-to-keep selected takes a third runtime selector input, so + # its forward cannot be invoked here; use the C++ runner (mlx_run_llm_hf). + if ( + read_const_int(program, LOGITS_TO_KEEP_MODE_METHOD) + == LOGITS_TO_KEEP_MODES["selected"] + ): + raise ValueError( + "This .pte was exported with --logits-to-keep selected, which needs " + "a runtime-supplied logits selector input that run_llm_hf.py does " + "not provide. Run it with the C++ runner mlx_run_llm_hf, or " + "re-export with --logits-to-keep full or last." + ) + + max_ctx_len, declared_max_seq_len = read_model_limits(program) + # The forward only accepts up to its traced token width, so clamp the + # declared step to it: optimum's static export takes 1 token/forward while + # its get_max_seq_len is the context length, and feeding more crashes + # set_inputs. A chunked-prefill export reports the two as equal. + input_seq_len = _forward_input_seq_len(program) + prefill_chunk_size = ( + min(declared_max_seq_len, input_seq_len) + if declared_max_seq_len is not None + else input_seq_len + ) logger.info( f"Model limits: max_ctx_len={max_ctx_len}, " f"prefill_chunk_size={prefill_chunk_size}" diff --git a/backends/mlx/examples/llm/runner_utils.h b/backends/mlx/examples/llm/runner_utils.h new file mode 100644 index 00000000000..6455dff7830 --- /dev/null +++ b/backends/mlx/examples/llm/runner_utils.h @@ -0,0 +1,118 @@ +/* + * 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 + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace executorch { +namespace backends { +namespace mlx { +namespace examples { +namespace llm { + +struct StopTokens { + std::unordered_set ids; + std::optional turn_end_id; +}; + +inline const char* turn_end_piece(const std::string& chat) { + if (chat == "llama3") { + return "<|eot_id|>"; + } + if (chat == "gemma") { + return ""; + } + if (chat == "gemma4") { + return ""; + } + return nullptr; +} + +inline bool resolve_stop_tokens( + tokenizers::Tokenizer& tokenizer, + ::executorch::extension::Module& module, + const std::string& chat, + StopTokens& out) { + out.ids = ::executorch::extension::llm::get_eos_ids(&tokenizer, &module); + out.turn_end_id.reset(); + if (chat == "0") { + return true; + } + const char* piece = turn_end_piece(chat); + if (piece == nullptr) { + return false; + } + auto id = tokenizer.piece_to_id(piece); + if (!id.ok()) { + return false; + } + out.turn_end_id = *id; + out.ids.insert(*id); + return true; +} + +inline bool wrap_turn( + const std::string& chat, + const std::string& prompt, + bool with_bos, + std::string& out) { + if (chat == "0") { + out = prompt; + } else if (chat == "llama3") { + out = std::string(with_bos ? "<|begin_of_text|>" : "") + + "<|start_header_id|>user<|end_header_id|>\n\n" + prompt + + "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"; + } else if (chat == "gemma") { + out = std::string(with_bos ? "" : "") + "user\n" + + prompt + "\nmodel\n"; + } else if (chat == "gemma4") { + out = std::string(with_bos ? "" : "") + "<|turn>user\n" + prompt + + "\n<|turn>model\n"; + } else { + return false; + } + return true; +} + +inline int storage_dtype(const std::string& name) { + using ScalarType = ::executorch::runtime::etensor::ScalarType; + if (name == "bf16") { + return static_cast(ScalarType::BFloat16); + } + if (name == "fp16") { + return static_cast(ScalarType::Half); + } + if (name == "fp32") { + return static_cast(ScalarType::Float); + } + return -1; +} + +inline int resolve_kv_storage_dtype( + const std::string& override_name, + ::executorch::aten::ScalarType activation_dtype) { + if (!override_name.empty()) { + return storage_dtype(override_name); + } + return static_cast(activation_dtype); +} + +} // namespace llm +} // namespace examples +} // namespace mlx +} // namespace backends +} // namespace executorch diff --git a/backends/mlx/examples/llm/runtime_meta.py b/backends/mlx/examples/llm/runtime_meta.py index 7e50b198861..d38c1d8eee2 100644 --- a/backends/mlx/examples/llm/runtime_meta.py +++ b/backends/mlx/examples/llm/runtime_meta.py @@ -6,8 +6,8 @@ """Shared runtime helpers for the MLX LLM example runners. -Exports publish their limits as constant methods (``get_max_ctx_len``, -``get_prefill_chunk_size``) so runners do not have to be told what a .pte +Exports publish their limits as constant methods (``get_max_context_len``, +``get_max_seq_len``) so runners do not have to be told what a .pte supports. This mirrors ``const_int`` in run_llm_hf.cpp. Prompt handling (processor loading, chat templating, EOS lookup) lives here too: @@ -21,6 +21,11 @@ import torch +from executorch.extension.llm.export.model_metadata import ( + MAX_CONTEXT_LEN_METHOD, + MAX_SEQ_LEN_METHOD, +) + logger = logging.getLogger(__name__) @@ -37,10 +42,10 @@ def read_const_int(program, name: str) -> Optional[int]: def read_model_limits(program) -> Tuple[Optional[int], Optional[int]]: - """Return (max_ctx_len, prefill_chunk_size) as published by the export.""" + """Return (max_context_len, max_seq_len) as published by the export.""" return ( - read_const_int(program, "get_max_ctx_len"), - read_const_int(program, "get_prefill_chunk_size"), + read_const_int(program, MAX_CONTEXT_LEN_METHOD), + read_const_int(program, MAX_SEQ_LEN_METHOD), ) diff --git a/backends/mlx/examples/whisper/export_whisper.py b/backends/mlx/examples/whisper/export_whisper.py index 97d3a22bc79..84b45042ef5 100644 --- a/backends/mlx/examples/whisper/export_whisper.py +++ b/backends/mlx/examples/whisper/export_whisper.py @@ -122,14 +122,17 @@ def forward( # Update KV cache k_cache, v_cache = self.kv_cache.update(pos_int, k, v) - # Explicit windowing: slice cache to valid positions - end_pos = pos_int + T - k_win = k_cache[:, :, :end_pos, :] - v_win = v_cache[:, :, :end_pos, :] - - # SDPA with causal mask - attn_out = F.scaled_dot_product_attention( - q, k_win, v_win, attn_mask=None, is_causal=True, scale=self.scale + # The cache-aware op slices the cache to start_pos + query length itself, and + # applies the causal mask the kernel wants for a cached decode step. Passing + # a hand-sliced window with is_causal instead means something different: torch + # would let the new token see only the first cached key. + attn_out = torch.ops.mlx.custom_sdpa( + q, + k_cache, + v_cache, + start_pos=pos_int, + is_causal=True, + scale=self.scale, ) # Reshape back diff --git a/backends/mlx/llm/exportable.py b/backends/mlx/llm/exportable.py index 4660247e27b..174d5ae7ea9 100644 --- a/backends/mlx/llm/exportable.py +++ b/backends/mlx/llm/exportable.py @@ -21,7 +21,8 @@ """ import logging -from typing import List, Optional, Sequence +from enum import IntEnum +from typing import List, Optional, Sequence, Union import torch from transformers.integrations.executorch import ( @@ -32,6 +33,50 @@ logger = logging.getLogger(__name__) +class LogitsToKeepMode(IntEnum): + FULL = 0 + LAST = 1 + SELECTED = 2 + + @classmethod + def from_value(cls, value: Union["LogitsToKeepMode", str, int]): + if isinstance(value, str): + try: + return cls[value.upper()] + except KeyError as error: + raise ValueError(f"Unsupported logits-to-keep mode: {value}") from error + return cls(value) + + +class _LogitsToKeepMixin: + logits_to_keep_mode: LogitsToKeepMode + + def _resolve_logits_to_keep( + self, logits_to_keep: Optional[torch.LongTensor] + ) -> Union[int, torch.LongTensor]: + if self.logits_to_keep_mode == LogitsToKeepMode.SELECTED: + if logits_to_keep is None: + raise ValueError("selected logits-to-keep requires an index tensor") + if logits_to_keep.dtype != torch.int64 or logits_to_keep.dim() != 1: + raise ValueError("logits_to_keep must be an int64[K] tensor") + return logits_to_keep + return int(self.logits_to_keep_mode) + + def _logits_to_keep_kwargs( + self, logits_to_keep: Optional[torch.LongTensor] + ) -> dict: + if self.logits_to_keep_mode == LogitsToKeepMode.FULL: + return {} + return {"logits_to_keep": self._resolve_logits_to_keep(logits_to_keep)} + + def _sync_cache_position(self, cache, cache_position) -> None: + if cache_position is None or not hasattr(cache, "layers"): + return + for layer in cache.layers: + if hasattr(layer, "cumulative_length"): + layer.cumulative_length.copy_(cache_position[0]) + + class _HiddenTapMixin: """Shared tapping logic - expects self.layer_ids and self.model to exist.""" @@ -41,8 +86,78 @@ def _tap_hidden(self, outs): return torch.cat(captured, dim=-1) +class TorchExportableModuleWithStaticCacheAndLogitsToKeep( + _LogitsToKeepMixin, TorchExportableModuleWithStaticCache +): + def __init__( + self, + model, + batch_size: Optional[int] = None, + max_cache_len: Optional[int] = None, + device: Optional[torch.device] = None, + logits_to_keep_mode: LogitsToKeepMode = LogitsToKeepMode.FULL, + ): + super().__init__( + model, batch_size=batch_size, max_cache_len=max_cache_len, device=device + ) + self.logits_to_keep_mode = LogitsToKeepMode.from_value(logits_to_keep_mode) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + cache_position: Optional[torch.Tensor] = None, + logits_to_keep: Optional[torch.LongTensor] = None, + ): + self._sync_cache_position(self.static_cache, cache_position) + return self.model( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + cache_position=cache_position, + attention_mask=None, + past_key_values=self.static_cache, + use_cache=True, + **self._logits_to_keep_kwargs(logits_to_keep), + ).logits + + +class TorchExportableModuleWithHybridCacheAndLogitsToKeep( + _LogitsToKeepMixin, TorchExportableModuleWithHybridCache +): + def __init__( + self, + model, + batch_size: Optional[int] = None, + max_cache_len: Optional[int] = None, + device: Optional[torch.device] = None, + logits_to_keep_mode: LogitsToKeepMode = LogitsToKeepMode.FULL, + ): + super().__init__( + model, batch_size=batch_size, max_cache_len=max_cache_len, device=device + ) + self.logits_to_keep_mode = LogitsToKeepMode.from_value(logits_to_keep_mode) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + cache_position: Optional[torch.Tensor] = None, + logits_to_keep: Optional[torch.LongTensor] = None, + ): + self._sync_cache_position(self.cache, cache_position) + return self.model( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + cache_position=cache_position, + attention_mask=None, + past_key_values=self.cache, + use_cache=True, + **self._logits_to_keep_kwargs(logits_to_keep), + ).logits + + class TorchExportableModuleWithStaticCacheAndHidden( - _HiddenTapMixin, TorchExportableModuleWithStaticCache + _HiddenTapMixin, _LogitsToKeepMixin, TorchExportableModuleWithStaticCache ): def __init__( self, @@ -51,10 +166,12 @@ def __init__( max_cache_len: Optional[int] = None, device: Optional[torch.device] = None, layer_ids: Sequence[int] = (), + logits_to_keep_mode: LogitsToKeepMode = LogitsToKeepMode.FULL, ): super().__init__( model, batch_size=batch_size, max_cache_len=max_cache_len, device=device ) + self.logits_to_keep_mode = LogitsToKeepMode.from_value(logits_to_keep_mode) if not layer_ids: raise ValueError("layer_ids must be non-empty") self.layer_ids: List[int] = list(layer_ids) @@ -64,7 +181,9 @@ def forward( input_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.Tensor] = None, cache_position: Optional[torch.Tensor] = None, + logits_to_keep: Optional[torch.LongTensor] = None, ): + self._sync_cache_position(self.static_cache, cache_position) outs = self.model( input_ids=input_ids, inputs_embeds=inputs_embeds, @@ -73,6 +192,7 @@ def forward( past_key_values=self.static_cache, use_cache=True, output_hidden_states=True, + **self._logits_to_keep_kwargs(logits_to_keep), ) hidden = self._tap_hidden(outs) if hasattr(outs, "logits"): @@ -81,7 +201,7 @@ def forward( class TorchExportableModuleWithHybridCacheAndHidden( - _HiddenTapMixin, TorchExportableModuleWithHybridCache + _HiddenTapMixin, _LogitsToKeepMixin, TorchExportableModuleWithHybridCache ): def __init__( self, @@ -90,10 +210,12 @@ def __init__( max_cache_len: Optional[int] = None, device: Optional[torch.device] = None, layer_ids: Sequence[int] = (), + logits_to_keep_mode: LogitsToKeepMode = LogitsToKeepMode.FULL, ): super().__init__( model, batch_size=batch_size, max_cache_len=max_cache_len, device=device ) + self.logits_to_keep_mode = LogitsToKeepMode.from_value(logits_to_keep_mode) if not layer_ids: raise ValueError("layer_ids must be non-empty") self.layer_ids: List[int] = list(layer_ids) @@ -103,7 +225,9 @@ def forward( input_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.Tensor] = None, cache_position: Optional[torch.Tensor] = None, + logits_to_keep: Optional[torch.LongTensor] = None, ): + self._sync_cache_position(self.cache, cache_position) outs = self.model( input_ids=input_ids, inputs_embeds=inputs_embeds, @@ -112,6 +236,7 @@ def forward( past_key_values=self.cache, use_cache=True, output_hidden_states=True, + **self._logits_to_keep_kwargs(logits_to_keep), ) hidden = self._tap_hidden(outs) if hasattr(outs, "logits"): @@ -124,6 +249,7 @@ def create_hf_exportable( max_cache_len: int, tap_layers: Optional[Sequence[int]] = None, batch_size: int = 1, + logits_to_keep_mode: Union[LogitsToKeepMode, str, int] = LogitsToKeepMode.FULL, ): """Factory: picks static vs hybrid and hidden-tapping vs plain. @@ -132,6 +258,7 @@ def create_hf_exportable( max_cache_len: cache capacity tap_layers: optional layer indices to tap and concat as second output batch_size: batch size for cache init + logits_to_keep_mode: full, last, or selected logits selection Returns: An exportable module with .model attribute pointing to HF model @@ -139,6 +266,7 @@ def create_hf_exportable( """ text_config = model.config.get_text_config() sliding_window = getattr(text_config, "sliding_window", None) + logits_to_keep_mode = LogitsToKeepMode.from_value(logits_to_keep_mode) if sliding_window is not None: if tap_layers is not None: @@ -150,12 +278,23 @@ def create_hf_exportable( batch_size=batch_size, max_cache_len=max_cache_len, layer_ids=tap_layers, + logits_to_keep_mode=logits_to_keep_mode, + ) + if logits_to_keep_mode == LogitsToKeepMode.FULL: + logger.info("Creating TorchExportableModuleWithHybridCache wrapper...") + return TorchExportableModuleWithHybridCache( + model=model, + batch_size=batch_size, + max_cache_len=max_cache_len, ) - logger.info("Creating TorchExportableModuleWithHybridCache wrapper...") - return TorchExportableModuleWithHybridCache( + logger.info( + f"Creating hybrid-cache wrapper with {logits_to_keep_mode.name.lower()} logits..." + ) + return TorchExportableModuleWithHybridCacheAndLogitsToKeep( model=model, batch_size=batch_size, max_cache_len=max_cache_len, + logits_to_keep_mode=logits_to_keep_mode, ) else: if tap_layers is not None: @@ -167,12 +306,23 @@ def create_hf_exportable( batch_size=batch_size, max_cache_len=max_cache_len, layer_ids=tap_layers, + logits_to_keep_mode=logits_to_keep_mode, ) - logger.info("Creating TorchExportableModuleWithStaticCache wrapper...") - return TorchExportableModuleWithStaticCache( + if logits_to_keep_mode == LogitsToKeepMode.FULL: + logger.info("Creating TorchExportableModuleWithStaticCache wrapper...") + return TorchExportableModuleWithStaticCache( + model=model, + batch_size=batch_size, + max_cache_len=max_cache_len, + ) + logger.info( + f"Creating static-cache wrapper with {logits_to_keep_mode.name.lower()} logits..." + ) + return TorchExportableModuleWithStaticCacheAndLogitsToKeep( model=model, batch_size=batch_size, max_cache_len=max_cache_len, + logits_to_keep_mode=logits_to_keep_mode, ) diff --git a/backends/mlx/llm/hf_attention.py b/backends/mlx/llm/hf_attention.py index c8f34b11d5e..ff8fc3ff9a3 100644 --- a/backends/mlx/llm/hf_attention.py +++ b/backends/mlx/llm/hf_attention.py @@ -41,8 +41,8 @@ from typing import Callable, Optional, Tuple, Union import executorch.backends.mlx.custom_ops as _mlx_custom_ops # noqa: F401 - import torch +from executorch.backends.mlx.llm.exportable import _LogitsToKeepMixin, LogitsToKeepMode def mlx_sdpa_with_start_pos_forward( @@ -219,8 +219,8 @@ def register_mlx_attention(name: str = "mlx") -> None: ) -class OffGraphExportWrapper(torch.nn.Module): - """forward(input_ids, cache_position) -> logits, with no in-graph cache. +class OffGraphExportWrapper(_LogitsToKeepMixin, torch.nn.Module): + """Export logits with no in-graph cache. The analog of TorchExportableModuleWithStaticCache for the off-graph op: runs the model with use_cache=False so each attention layer sees only this @@ -228,12 +228,16 @@ class OffGraphExportWrapper(torch.nn.Module): signature the runner drives. """ - def __init__(self, model: torch.nn.Module): + def __init__(self, model: torch.nn.Module, logits_to_keep_mode="full"): super().__init__() self.model = model + self.logits_to_keep_mode = LogitsToKeepMode.from_value(logits_to_keep_mode) def forward( - self, input_ids: torch.Tensor, cache_position: torch.Tensor + self, + input_ids: torch.Tensor, + cache_position: torch.Tensor, + logits_to_keep: Optional[torch.LongTensor] = None, ) -> torch.Tensor: # Single sequence: the op takes [q_len, n_dims] positions and the # attention function reads position_ids[0], so a batch would be placed @@ -249,6 +253,7 @@ def forward( position_ids=cache_position.unsqueeze(0), use_cache=False, past_key_values=None, + **self._logits_to_keep_kwargs(logits_to_keep), ).logits diff --git a/backends/mlx/partitioner.py b/backends/mlx/partitioner.py index 7814e883588..a82f54cda00 100644 --- a/backends/mlx/partitioner.py +++ b/backends/mlx/partitioner.py @@ -133,6 +133,18 @@ def ops_to_not_decompose( handler that rejects the 6-arg edge form, for instance). Preserving an op the handler then rejects is worse than not preserving it, because the op neither decomposes into something delegatable nor lowers itself. + + A target is only preserved when every node carrying it is supported. One + unsupported node is enough to give the whole operator back to decomposition, + because keeping it would leave that node neither lowered nor decomposed and + export would stop. + + The second return value is a per-node filter, which would keep the supported + calls fused and decompose only the rest. It is deliberately not used: it puts + the program on exir's EDGE_DO_NOT_DECOMP path, which fails on an ordinary + attention block that reshapes its output, a shape this backend has to lower. + The cost of the coarser choice is that one declined call also unfuses the + operator's other calls in that graph. """ from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder @@ -157,6 +169,7 @@ def ops_to_not_decompose( # Collect ops for nodes that are actually supported do_not_decompose: list[torch._ops.OpOverload] = [] + declined: set[torch._ops.OpOverload] = set() for node in ep.graph.nodes: if node.op == "call_function" and isinstance( @@ -166,6 +179,10 @@ def ops_to_not_decompose( if info is not None and info.supported: if node.target not in do_not_decompose: do_not_decompose.append(node.target) + else: + declined.add(node.target) + + do_not_decompose = [op for op in do_not_decompose if op not in declined] self._not_decompose_cache = (weakref.ref(ep), do_not_decompose) diff --git a/backends/mlx/patches/apply.sh b/backends/mlx/patches/apply.sh old mode 100644 new mode 100755 diff --git a/backends/mlx/patches/mlx_gather_mm_rhs_lda.patch b/backends/mlx/patches/mlx_gather_mm_rhs_lda.patch deleted file mode 100644 index 8430d83e9cd..00000000000 --- a/backends/mlx/patches/mlx_gather_mm_rhs_lda.patch +++ /dev/null @@ -1,34 +0,0 @@ -Fix the activation row stride in sorted RHS gather_mm. - -The specialized gather_mm_rhs paths flatten all leading activation dimensions -into M, but derive lda from the original second-to-last dimension. For a valid -[N, 1, K] view produced by expand_dims, that singleton dimension can have -stride 1. The kernel then reads row n from A + n instead of A + n * K. - -The activation is made row-contiguous before this calculation, so flattened -rows are K elements apart. Use K as lda in both the Steel and NAX paths. - -Upstream candidate. - -diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp -index 87d2bf52..4fc2c1de 100644 ---- a/mlx/backend/metal/matmul.cpp -+++ b/mlx/backend/metal/matmul.cpp -@@ -1897,7 +1897,7 @@ void gather_mm_rhs( - int K = a.shape(-1); - int M = a.size() / K; - int N = b.shape(-1); -- int lda = a.strides()[a.ndim() - 2]; // should be K -+ int lda = K; - - // Define the dispatch blocks - int bm = 16, bn = 64, bk = 16; -@@ -2030,7 +2030,7 @@ void gather_mm_rhs_nax( - int K = a.shape(-1); - int M = a.size() / K; - int N = b.shape(-1); -- int lda = a.strides()[a.ndim() - 2]; // should be K -+ int lda = K; - int E = b.shape(0); - - // Define the dispatch blocks diff --git a/backends/mlx/patches/mlx_metal_remove_addrspace_compat.patch b/backends/mlx/patches/mlx_metal_remove_addrspace_compat.patch new file mode 100644 index 00000000000..6e1ad402c65 --- /dev/null +++ b/backends/mlx/patches/mlx_metal_remove_addrspace_compat.patch @@ -0,0 +1,39 @@ +Keep Steel integral constants compatible with older Metal frontends. + +MLX supports macOS 14 and selects Metal language 3.1 there. The Metal 4.1 +address-space fix uses metal::remove_addrspace_t, but that trait is unavailable +in the Xcode 15.4 Metal standard library. Because integral_constant.h is embedded +in runtime-JIT sources, this prevents kernels from compiling on that supported +toolchain. + +Use remove_addrspace_t on Metal 4.1+, where explicit thread qualifiers make it +necessary and the trait is available. For older language versions, preserve the +previous decltype behavior with an identity alias. + +Upstream candidate; carried locally until MLX includes the compatibility fix. + +diff --git a/mlx/backend/metal/kernels/steel/utils/integral_constant.h b/mlx/backend/metal/kernels/steel/utils/integral_constant.h +index 2f153f48..97022416 100644 +--- a/mlx/backend/metal/kernels/steel/utils/integral_constant.h ++++ b/mlx/backend/metal/kernels/steel/utils/integral_constant.h +@@ -47,11 +47,19 @@ using Int = integral_constant; + // Binary Operators on Integral constants + /////////////////////////////////////////////////////////////////////////////// + ++#if __METAL_VERSION__ >= 410 ++template ++using remove_addrspace_t = metal::remove_addrspace_t; ++#else ++template ++using remove_addrspace_t = T; ++#endif ++ + #define integral_const_binop(__op__, __operator__) \ + template \ + METAL_FUNC constexpr auto __operator__( \ + integral_constant, integral_constant) { \ + constexpr auto res = tv __op__ uv; \ +- using res_t = metal::remove_addrspace_t; \ ++ using res_t = remove_addrspace_t; \ + return integral_constant{}; \ + } diff --git a/backends/mlx/patches/mlx_metal_sdk_per_platform.patch b/backends/mlx/patches/mlx_metal_sdk_per_platform.patch index e8983d61844..595b452feca 100644 --- a/backends/mlx/patches/mlx_metal_sdk_per_platform.patch +++ b/backends/mlx/patches/mlx_metal_sdk_per_platform.patch @@ -1,45 +1,56 @@ -Select the Metal SDK per target platform. +Select the Metal SDK from the CMake target platform. -MLX compiles and links its Metal shader library with a hardcoded -`xcrun -sdk macosx metal` and a hardcoded `-mmacosx-version-min` flag. ExecuTorch -builds MLX for iOS device, iOS simulator, and macOS from one source tree, so the -hardcoded macOS SDK produces a metallib built for the wrong platform on the iOS -and simulator slices. +MLX supports both macOS and iOS CMake targets, but its Metal compile and link +commands use the macOS SDK and deployment flag unconditionally. That produces a +metallib for the wrong platform when cross-compiling for an iOS device or +simulator. -Derive the Metal SDK and the deployment-version flag from PLATFORM (which -ExecuTorch already passes to this build): iphoneos for OS64, iphonesimulator for -SIMULATORARM64, macosx otherwise. This mirrors how the rest of the Apple build -selects its SDK. +Use CMAKE_SYSTEM_NAME to distinguish macOS from iOS and CMAKE_OSX_SYSROOT to +distinguish iOS device and simulator SDKs. CMAKE_OSX_SYSROOT may be either a +short SDK name or an absolute, versioned SDK path, so match it case-insensitively. +Fail ambiguous or unsupported configurations instead of silently producing a +metallib for the wrong platform. -Upstream candidate; carried locally until MLX selects the Metal SDK by platform. +Upstream candidate; carried locally until MLX selects the Metal SDK from its +CMake target platform. diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt -index edc169ee..94154f80 100644 +index ecabbda5..d0b86d3b 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt -@@ -9,6 +9,21 @@ set(BASE_HEADERS +@@ -9,6 +9,31 @@ set(BASE_HEADERS logging.h utils.h) -+# The Metal SDK and deployment flag must follow the target platform. ExecuTorch -+# builds this for iOS device, iOS simulator, and macOS from one source tree and -+# passes PLATFORM in for each. Without this the shaders are always built against -+# the macOS SDK, so the iOS and simulator metallibs are wrong for their slice. -+if(PLATFORM STREQUAL "OS64") -+ set(MLX_METAL_SDK iphoneos) -+ set(MLX_METAL_VERSION_MIN_FLAG "-mios-version-min") -+elseif(PLATFORM STREQUAL "SIMULATORARM64") -+ set(MLX_METAL_SDK iphonesimulator) -+ set(MLX_METAL_VERSION_MIN_FLAG "-mios-simulator-version-min") -+else() ++# Match the Metal SDK and deployment flag to the CMake target platform. ++if(CMAKE_SYSTEM_NAME STREQUAL "Darwin") + set(MLX_METAL_SDK macosx) + set(MLX_METAL_VERSION_MIN_FLAG "-mmacosx-version-min") ++elseif(CMAKE_SYSTEM_NAME STREQUAL "iOS") ++ string(TOLOWER "${CMAKE_OSX_SYSROOT}" _mlx_osx_sysroot) ++ if(_mlx_osx_sysroot MATCHES "iphonesimulator") ++ set(MLX_METAL_SDK iphonesimulator) ++ set(MLX_METAL_VERSION_MIN_FLAG "-mios-simulator-version-min") ++ elseif(_mlx_osx_sysroot MATCHES "iphoneos") ++ set(MLX_METAL_SDK iphoneos) ++ set(MLX_METAL_VERSION_MIN_FLAG "-mios-version-min") ++ else() ++ message( ++ FATAL_ERROR ++ "Unable to select the Metal SDK for iOS from CMAKE_OSX_SYSROOT='${CMAKE_OSX_SYSROOT}'" ++ ) ++ endif() ++else() ++ message( ++ FATAL_ERROR ++ "MLX Metal supports only macOS and iOS, got CMAKE_SYSTEM_NAME='${CMAKE_SYSTEM_NAME}'" ++ ) +endif() + function(build_kernel_base TARGET SRCFILE DEPS) set(METAL_FLAGS -x -@@ -26,10 +41,10 @@ function(build_kernel_base TARGET SRCFILE DEPS) +@@ -27,10 +52,10 @@ function(build_kernel_base TARGET SRCFILE DEPS) endif() if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "") set(METAL_FLAGS ${METAL_FLAGS} @@ -52,7 +63,7 @@ index edc169ee..94154f80 100644 -I${PROJECT_SOURCE_DIR} -o ${TARGET}.air DEPENDS ${SRCFILE} ${DEPS} ${BASE_HEADERS} OUTPUT ${TARGET}.air -@@ -180,12 +195,12 @@ endif() +@@ -188,12 +213,12 @@ endif() set(METAL_LINK_FLAGS) if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "") diff --git a/backends/mlx/patches/mlx_nax_jit_sdk_gate.patch b/backends/mlx/patches/mlx_nax_jit_sdk_gate.patch deleted file mode 100644 index b1b1fcfd820..00000000000 --- a/backends/mlx/patches/mlx_nax_jit_sdk_gate.patch +++ /dev/null @@ -1,112 +0,0 @@ -Gate the NAX JIT kernel sources behind the SDK requirement. - -MLX's NAX kernels (GEMM and attention) include -, a framework that only -ships in the macOS 26 / Xcode 26 SDK. With MLX_METAL_JIT=ON (which ExecuTorch -uses) on an older SDK, the JIT preamble generator (make_compiled_preamble.sh) -runs `metal -E` over these headers, which fatals on the missing include and -fails the build. MLX already gates NAX on the non-JIT metallib path -(kernels/CMakeLists.txt), but the JIT path was ungated. - -Instead of guarding the includes with __has_include, gate the NAX -make_jit_source() calls behind the same -MLX_METAL_VERSION/MACOS_SDK_VERSION/CMAKE_OSX_DEPLOYMENT_TARGET check the -metallib path uses, and define MLX_METAL_NO_NAX when the requirement is unmet. -On those SDKs NAX is already runtime-gated via is_nax_available() -(device.cpp), so the get_*_nax_kernel entry points in jit_kernels.cpp are -unreachable; guarded empty preamble getters keep that translation unit linking. -macOS 26+ SDKs still build NAX. - -Upstream candidate; carried locally until an MLX release gates the JIT path. - -diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt ---- a/mlx/backend/metal/CMakeLists.txt -+++ b/mlx/backend/metal/CMakeLists.txt -@@ -84,19 +84,32 @@ if(MLX_METAL_JIT) - - make_jit_source(steel/attn/kernels/steel_attention) - -- make_jit_source( -- steel/gemm/gemm_nax kernels/steel/utils.h kernels/steel/gemm/nax.h -- kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h) -- make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax) -- make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax) -- make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax) -- make_jit_source(steel/gemm/kernels/steel_gemm_segmented_nax) -+ if(MLX_METAL_VERSION GREATER_EQUAL 400 -+ AND MACOS_SDK_VERSION VERSION_GREATER_EQUAL 26.2 -+ AND CMAKE_OSX_DEPLOYMENT_TARGET VERSION_GREATER_EQUAL 26.2) - -- make_jit_source(quantized_nax kernels/quantized_utils.h) -- make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h -- kernels/fp4.h) -+ make_jit_source( -+ steel/gemm/gemm_nax kernels/steel/utils.h kernels/steel/gemm/nax.h -+ kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h) -+ make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax) -+ make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax) -+ make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax) -+ make_jit_source(steel/gemm/kernels/steel_gemm_segmented_nax) -+ -+ make_jit_source(quantized_nax kernels/quantized_utils.h) -+ make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h -+ kernels/fp4.h) -+ -+ make_jit_source(steel/attn/kernels/steel_attention_nax) - -- make_jit_source(steel/attn/kernels/steel_attention_nax) -+ else() -+ message( -+ WARNING "NAX kernels require Metal 4, macOS SDK >= 26.2, and " -+ "MACOSX_DEPLOYMENT_TARGET >= 26.2 (SDK ${MACOS_SDK_VERSION}, " -+ "CMAKE_OSX_DEPLOYMENT_TARGET=${CMAKE_OSX_DEPLOYMENT_TARGET}). " -+ "Building without NAX kernels.") -+ target_compile_definitions(mlx PRIVATE MLX_METAL_NO_NAX) -+ endif() - - else() - target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nojit_kernels.cpp) -diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp ---- a/mlx/backend/metal/jit_kernels.cpp -+++ b/mlx/backend/metal/jit_kernels.cpp -@@ -8,6 +8,40 @@ using namespace fmt::literals; - - namespace mlx::core { - -+#ifdef MLX_METAL_NO_NAX -+// NAX JIT preambles are only generated (via make_jit_source) when the SDK -+// requirement is met. On older SDKs they are skipped and MLX_METAL_NO_NAX is -+// defined, so is_nax_available() returns false and the get_*_nax_kernel entry -+// points below are never reached. These empty definitions only exist to satisfy -+// the linker for this translation unit. -+namespace metal { -+const char* gemm_nax() { -+ return ""; -+} -+const char* steel_gemm_fused_nax() { -+ return ""; -+} -+const char* steel_gemm_gather_nax() { -+ return ""; -+} -+const char* steel_gemm_splitk_nax() { -+ return ""; -+} -+const char* steel_gemm_segmented_nax() { -+ return ""; -+} -+const char* quantized_nax() { -+ return ""; -+} -+const char* fp_quantized_nax() { -+ return ""; -+} -+const char* steel_attention_nax() { -+ return ""; -+} -+} // namespace metal -+#endif // MLX_METAL_NO_NAX -+ - MTL::ComputePipelineState* get_arange_kernel( - metal::Device& d, - const std::string& kernel_name, diff --git a/backends/mlx/patches/mlx_qmm_splitk_bk_align.patch b/backends/mlx/patches/mlx_qmm_splitk_bk_align.patch deleted file mode 100644 index f4405b5ae9f..00000000000 --- a/backends/mlx/patches/mlx_qmm_splitk_bk_align.patch +++ /dev/null @@ -1,42 +0,0 @@ -Align split-K partitions to the qmm K-tile (BK=32), fixing nvfp4. - -MLX v0.32.0's qmm_splitk caps split_k only by the quantization group count -(K / group_size), not by the kernel's K-tile width BK (=32). For nvfp4 -(group_size=16) this yields a per-partition K of 16 < BK, and fp_qmm_t_splitk's -tile load reads a full BK-wide K-tile with no K bound -- spilling 16 columns -past the partition into the next group's packed weights and fp8 scales. That -corrupts every partial (non-uniform ~2x error) and reads past the buffer on the -last partition (NaN/inf). Affine (group_size >= 32) is unaffected because its -partitions are already >= BK. - -Fix in the dispatch (MLX's pattern for tile-alignment constraints): require each -K partition to be a whole number of BK-wide tiles as well as whole groups, i.e. -align split_k to max(group_size, BK). Only changes behavior for group_size < 32. - -Upstream candidate. - -diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp -index 62d48714..94c56307 100644 ---- a/mlx/backend/metal/quantized.cpp -+++ b/mlx/backend/metal/quantized.cpp -@@ -884,11 +884,15 @@ void qmm_splitk( - int current_tgs = n_tiles * m_tiles; - int split_k = std::max(1, 512 / current_tgs); - -- // Cap split_k by the number of quantization groups -- split_k = std::min(split_k, K / group_size); -- -- // Ensure K divides evenly by split_k * group_size -- while (split_k > 1 && (K % (split_k * group_size) != 0)) { -+ // Each K partition must be a whole number of BK-wide (32) K-tiles as well as -+ // whole quantization groups. The qmm_t_splitk kernels tile K by BK=32 and do -+ // not bound the K dimension, so a partition smaller than BK (e.g. nvfp4's -+ // group_size=16) would over-read into the next group's weights/scales. -+ int k_align = group_size > 32 ? group_size : 32; -+ split_k = std::min(split_k, K / k_align); -+ -+ // Ensure K divides evenly by split_k * k_align -+ while (split_k > 1 && (K % (split_k * k_align) != 0)) { - split_k--; - } - if (split_k <= 1) { diff --git a/backends/mlx/patches/mlx_swiftpm_metallib_name.patch b/backends/mlx/patches/mlx_swiftpm_metallib_name.patch deleted file mode 100644 index 6f08cb6578c..00000000000 --- a/backends/mlx/patches/mlx_swiftpm_metallib_name.patch +++ /dev/null @@ -1,35 +0,0 @@ -Make the SwiftPM metallib name a build-time define. - -MLX loads its Metal library from a SwiftPM resource bundle by a fixed name, -"default". ExecuTorch ships one resource bundle that holds a separate metallib for -each Apple platform slice (device, simulator, macOS), because a metallib built for -one slice does not load on another. Each slice's binary therefore has to ask for -its own file. - -Read the name from MLX_SWIFTPM_METALLIB_NAME when defined, falling back to -"default" so a plain MLX build is unchanged. - -Upstream candidate; carried locally until MLX supports a per-slice bundle name. - -diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp -index 29eecf55..61a1180f 100644 ---- a/mlx/backend/metal/device.cpp -+++ b/mlx/backend/metal/device.cpp -@@ -217,8 +217,15 @@ MTL::Library* load_default_library(MTL::Device* device) { - return lib; - } - -- // Then try default.metallib in a SwiftPM bundle if we have one -- std::tie(lib, error[2]) = load_swiftpm_library(device, "default"); -+ // Then try the metallib in a SwiftPM bundle if we have one. The name is a -+ // build-time define because ExecuTorch ships one bundle holding a metallib per -+ // platform slice, so each slice's binary must ask for its own file rather than a -+ // single shared "default". Falls back to "default" for a plain MLX build. -+#ifndef MLX_SWIFTPM_METALLIB_NAME -+#define MLX_SWIFTPM_METALLIB_NAME "default" -+#endif -+ std::tie(lib, error[2]) = -+ load_swiftpm_library(device, MLX_SWIFTPM_METALLIB_NAME); - if (lib) { - return lib; - } diff --git a/backends/mlx/patterns.py b/backends/mlx/patterns.py index 1760dc63b9a..9ff360fdb8f 100644 --- a/backends/mlx/patterns.py +++ b/backends/mlx/patterns.py @@ -21,6 +21,7 @@ from executorch.backends.mlx.builder.op_helpers import ( emit_quantized_biases, emit_quantized_gather, + emit_shape, emit_stop_position, mlx_qparams_supported, parse_dequant_int4_node, @@ -44,6 +45,7 @@ AddIntNode, AddNode, AsTypeNode, + ExpandDimsNode, IndexCopyNode, IntOrVid, ModIntNode, @@ -52,10 +54,12 @@ SdpaNode, SliceNode, SliceUpdateNode, + SqueezeNode, SubtractIntNode, SymSizeNode, ) from torch.export.exported_program import ExportedProgram +from torch.fx.experimental.symbolic_shapes import statically_known_true from torch.fx.node import Node @@ -527,6 +531,71 @@ def _try_unwrap_repeat_kv(cls, node: Node) -> Optional[Tuple[Node, List[Node]]]: body = [e for e in entries if e is not None] return base, body + @classmethod + def _kernel_can_compute(cls, sdpa_node: Node) -> bool: + """Whether the fused kernel can compute this call faithfully. + + Its preconditions are narrower than what PyTorch accepts, and a claimed call + that the kernel cannot compute fails loudly at execute rather than falling + back, so declining it here is what sends it to decomposition instead. + """ + q, k, v, attn_mask, _, is_causal, _, _ = cls._parse_sdpa_args_and_kwargs( + sdpa_node + ) + operand_vals = [ + operand.meta.get("val") if isinstance(operand, Node) else None + for operand in (q, k, v) + ] + if any(val is None for val in operand_vals): + return False + + # Ranks 2 and 3 are lifted to 4 on emission. Beyond 4 the leading dimensions + # would have to fold together, and a fold pairs the wrong operands as soon as + # one of them broadcasts a batch the others do not. + ranks = {val.dim() for val in operand_vals} + if len(ranks) != 1 or not 2 <= next(iter(ranks)) <= 4: + return False + + # Torch anchors a causal mask at the top left and MLX at the bottom right, so + # the two agree only when the query and key lengths are equal. Rank 4 already + # reaches the kernel today and is left alone; the lower ranks are newly lifted + # here, so do not open a path that returns wrong values with no error. + if ( + next(iter(ranks)) < 4 + and is_causal + and not statically_known_true( + operand_vals[0].shape[-2] == operand_vals[1].shape[-2] + ) + ): + return False + + # The kernel requires the batch sizes to match and rejects the call otherwise, + # so a broadcast batch has to be decomposed rather than fused. + if next(iter(ranks)) == 4 and any( + not statically_known_true(val.shape[0] == operand_vals[0].shape[0]) + for val in operand_vals[1:] + ): + return False + + if attn_mask is not None: + mask_val = ( + attn_mask.meta.get("val") if isinstance(attn_mask, Node) else None + ) + if mask_val is None or mask_val.dim() > 4: + return False + + # Causal attention is only taken when the query is known to be no longer than + # the keys. Torch clamps a longer query to the keys that exist, and neither the + # kernel's own mask nor slicing the keys reproduces that. A query length that + # cannot be compared at build time, two unrelated dynamic dimensions for + # instance, is also declined, because the relation has to hold for every call. + if is_causal and not statically_known_true( + operand_vals[0].shape[-2] <= operand_vals[1].shape[-2] + ): + return False + + return True + @classmethod def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler"]: sdpa_node = head @@ -535,15 +604,23 @@ def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler" ): return None + if not cls._kernel_can_compute(sdpa_node): + return None + q, k, v, _, _, _, _, _ = cls._parse_sdpa_args_and_kwargs(sdpa_node) - # Detect grouped kv attention pattern with repeat_interleave before SDPA + # Detect grouped kv attention pattern with repeat_interleave before SDPA. + # Both unwraps below key on dim 1, which is the head dimension only at rank 4. + # At a lower rank dim 1 is the key sequence, and absorbing a repeat there + # drops keys, which a causal mask then turns into a wrong answer. + is_rank4 = q.meta["val"].dim() == 4 if isinstance(q, Node) else False is_grouped_kv = False k_base = k v_base = v body: List[Node] = [] if ( - match_target(k, torch.ops.aten.repeat_interleave.self_int) + is_rank4 + and match_target(k, torch.ops.aten.repeat_interleave.self_int) and has_single_user(k) and (len(k.args) == 3) and (len(k.kwargs) == 0) @@ -563,7 +640,7 @@ def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler" # Detect HuggingFace repeat_kv pattern: # unsqueeze(dim=2) → expand → clone → view - if not is_grouped_kv: + if is_rank4 and not is_grouped_kv: k_unwrap = cls._try_unwrap_repeat_kv(k) v_unwrap = cls._try_unwrap_repeat_kv(v) if k_unwrap is not None and v_unwrap is not None: @@ -572,6 +649,27 @@ def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler" is_grouped_kv = True body = k_body + v_body + # Checked after the unwrapping above, because grouped-query attention reaches + # the kernel with its original head counts. MLX pairs heads only when the key + # and value agree and the query is a whole multiple of them. + kernel_vals = [ + node.meta.get("val") if isinstance(node, Node) else None + for node in (q, k_base, v_base) + ] + if any(val is None for val in kernel_vals): + return None + q_heads, k_heads, v_heads = ( + 1 if val.dim() == 2 else val.shape[-3] for val in kernel_vals + ) + # A zero head count would make the multiple test below divide by zero, and a + # raise here aborts the whole export rather than declining this one node. + if not statically_known_true(k_heads > 0): + return None + if not statically_known_true(k_heads == v_heads): + return None + if not statically_known_true(q_heads % k_heads == 0): + return None + head = sdpa_node if not is_grouped_kv: body = [] @@ -593,19 +691,73 @@ def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: assert dropout_p == 0.0, "SDPA with dropout is not supported" q, k, v, attn_mask = P.slot_map([q, k, v, attn_mask]) + # Add the dimensions the kernel is missing at the front, never in the middle. + # For a rank-3 input the first dimension is already the head one, so inserting + # there would move it into the batch slot and misalign masks and grouped heads. + input_nodes = (self.q_node, self.k_node, self.v_node) + inputs = [q, k, v] + for i, input_node in enumerate(input_nodes): + for _ in range(4 - input_node.meta["val"].dim()): + _, expanded = P.make_tmp_slot() + P.emit( + ExpandDimsNode( + x=P.slot_to_tid(inputs[i]), + out=P.slot_to_tid(expanded), + axis=0, + ) + ) + inputs[i] = expanded + + output_rank = n.meta["val"].dim() + out = P.make_or_get_slot(n) + sdpa_out = out + if output_rank < 4: + _, sdpa_out = P.make_tmp_slot() + + # MLX anchors its causal mask at the bottom right and torch at the top left, so + # the flag alone only means the same thing when the lengths are equal. Slicing + # the keys and values to the query length makes the problem square, where the + # two conventions agree, and it is what torch computes: a query at row i attends + # to keys 0..i, so keys past the last query row are never read. + q_len = self.q_node.meta["val"].shape[-2] + k_len = self.k_node.meta["val"].shape[-2] + if is_causal and not statically_known_true(q_len == k_len): + # A literal when the query length is known at build time, which is the + # decode case this exists for, and a size node only when it is symbolic. + rows = emit_shape(P, self.q_node, inputs[0])[-2] + for i in (1, 2): + _, sliced = P.make_tmp_slot() + P.emit( + SliceNode( + x=P.slot_to_tid(inputs[i]), + out=P.slot_to_tid(sliced), + axis=IntOrVid.from_literal(2), + start=IntOrVid.from_literal(0), + stop=rows, + ) + ) + inputs[i] = sliced P.emit( SdpaNode( - q=P.slot_to_tid(q), - k=P.slot_to_tid(k), - v=P.slot_to_tid(v), - out=P.slot_to_tid(out), + q=P.slot_to_tid(inputs[0]), + k=P.slot_to_tid(inputs[1]), + v=P.slot_to_tid(inputs[2]), + out=P.slot_to_tid(sdpa_out), scale=scale, mask=P.slot_to_tid(attn_mask) if attn_mask else None, causal=is_causal, ) ) + if output_rank < 4: + P.emit( + SqueezeNode( + x=P.slot_to_tid(sdpa_out), + out=P.slot_to_tid(out), + dims=list(range(4 - output_rank)), + ) + ) return out diff --git a/backends/mlx/runtime/MLXBackend.cpp b/backends/mlx/runtime/MLXBackend.cpp index 615c4f1c6ce..f5a79a9402d 100644 --- a/backends/mlx/runtime/MLXBackend.cpp +++ b/backends/mlx/runtime/MLXBackend.cpp @@ -14,6 +14,10 @@ #include "MLXSequenceCache.h" #include "mlx_mutable_state.h" +#ifdef EXECUTORCH_MLX_SWIFTPM_RESOURCES +#include "SwiftPMMetallibPath.h" +#endif + #include #include @@ -21,6 +25,7 @@ #include #include #include +#include #include @@ -186,8 +191,8 @@ struct MLXHandle { // Keep-alive for the off-graph KV cache bound in init(). state.cache is a // non-owning view of the same object, so the cache must outlive the handle - // even if the runner's session is torn down first. - std::shared_ptr<::executorch::extension::llm::cache::CacheBase> cache_shared; + // even if the runner drops its InstallGuard first. + std::shared_ptr<::executorch::extension::llm::cache::Cache> cache_shared; // Keep the constant buffers alive for zero-copy constants // Each FreeableBuffer must outlive the MLX arrays that reference it @@ -214,6 +219,23 @@ static std::mutex& mlx_global_mutex() { return m; } +#ifdef EXECUTORCH_MLX_SWIFTPM_RESOURCES +// Must be called while holding mlx_global_mutex() and before MLX initializes +// its Metal device. An application-provided path always takes precedence. +static bool configure_metallib_path_locked() { + if (!::mlx::core::metal::get_metallib_path().empty()) { + return true; + } + + const auto path = resolve_swiftpm_metallib_path(); + if (!path.has_value()) { + return false; + } + ::mlx::core::metal::set_metallib_path(*path); + return true; +} +#endif + class MLXBackend final : public ::executorch::runtime::BackendInterface { public: ~MLXBackend() override = default; @@ -236,6 +258,17 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { FreeableBuffer* processed, ArrayRef compile_specs) const override { std::lock_guard lock(mlx_global_mutex()); +#ifdef EXECUTORCH_MLX_SWIFTPM_RESOURCES + if (!configure_metallib_path_locked()) { + ET_LOG( + Error, + "Failed to find the MLX metallib in the SwiftPM resource bundle"); + if (processed != nullptr) { + processed->Free(); + } + return Error::NotFound; + } +#endif auto* handle = context.get_runtime_allocator()->allocateInstance(); if (handle == nullptr) { @@ -339,7 +372,8 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { // Bind the off-graph KV cache, if the runner installed one under a key it // passed as a runtime spec. Bound before the init chain runs so an // update_and_attend node there sees the same cache execute() will. - if (auto spec = context.get_runtime_spec(kCacheKeyKey); + if (auto spec = + context.get_runtime_spec(cache::kCacheKeyOption); spec.ok() && spec.get() != nullptr && *spec.get() != '\0') { const char* cache_key = spec.get(); handle->cache_shared = @@ -349,12 +383,10 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { std::string("init: cache_key '") + cache_key + "' is not installed in the CacheRegistry"); } - // Cross-cast from the neutral ownership anchor to this backend's - // tensor-typed op face; the two are deliberately unrelated bases (see - // MLXCache.h), so nullptr here means the key names another backend's - // cache. - handle->state.cache = - dynamic_cast(handle->cache_shared.get()); + // Ask the neutral ownership anchor for this backend's tensor-typed op + // face. It is named by MLXCache itself rather than by cache.h, so + // nullptr here means the key names another backend's cache. + handle->state.cache = handle->cache_shared->as(); if (handle->state.cache == nullptr) { throw std::runtime_error( std::string("init: cache under key '") + cache_key + @@ -571,18 +603,30 @@ static auto success_with_compiler = register_backend(backend); // Cache kind is named by the builder tag rather than an enum on the config: a // runner asks the registry for (backend_id, kind) and gets back a neutral -// CacheBase it installs under a cache_key. Adding a kind is a new builder here. +// Cache it installs under a cache_key. Adding a kind is a new builder here. const int cache_builders_registered = [] { - cache::CacheBuilderRegistry::global().register_builder( - kMLXBackendId, "seq", [](const cache::CacheConfig& cfg) { - return std::shared_ptr( + const Error single = cache::CacheFactory::global().register_builder( + kMLXBackendId, cache::kind::kSingle, [](const cache::CacheConfig& cfg) { + return std::shared_ptr( std::make_shared(cfg)); }); - cache::CacheBuilderRegistry::global().register_builder( - kMLXBackendId, "cell", [](const cache::CacheConfig& cfg) { - return std::shared_ptr( + ET_CHECK_MSG( + single == Error::Ok, + "Failed to register cache builder for %s:%s", + kMLXBackendId, + cache::kind::kSingle); + const Error batched_cell = cache::CacheFactory::global().register_builder( + kMLXBackendId, + cache::kind::kBatchedCell, + [](const cache::CacheConfig& cfg) { + return std::shared_ptr( std::make_shared(cfg)); }); + ET_CHECK_MSG( + batched_cell == Error::Ok, + "Failed to register cache builder for %s:%s", + kMLXBackendId, + cache::kind::kBatchedCell); return 0; }(); } // namespace diff --git a/backends/mlx/runtime/MLXCache.h b/backends/mlx/runtime/MLXCache.h index 523d3054712..40da8e94110 100644 --- a/backends/mlx/runtime/MLXCache.h +++ b/backends/mlx/runtime/MLXCache.h @@ -30,12 +30,16 @@ struct AttendSpec { }; // Tensor-typed op face of the off-graph KV cache, kept separate from the -// neutral CacheBase (which is tensor-free) so a cache can expose both without a +// neutral Cache (which is tensor-free) so a cache can expose both without a // diamond. ExecutionState holds one; nothing assigns it yet -- the registry // that owns the cache and hands this pointer to the executor lands in a // follow-up, until which exec_update_and_attend is unreachable. class MLXCache { public: + // Named here, not in cache.h: a backend face is tensor-typed and the + // neutral header cannot know about it. + static constexpr const char* kFaceName = "mlx.MLXCache"; + virtual ~MLXCache() = default; // Write this step's K/V for `layer` at `positions`, one host int per query diff --git a/backends/mlx/runtime/MLXCellCache.h b/backends/mlx/runtime/MLXCellCache.h index 566fc6f7bfa..91719fa0526 100644 --- a/backends/mlx/runtime/MLXCellCache.h +++ b/backends/mlx/runtime/MLXCellCache.h @@ -84,6 +84,14 @@ class MLXCellCache : public cache::CellCache, public MLXCache { mask(*step)}; } + protected: + void* face(cache::FaceId id) override { + if (void* p = cache::CellCache::face(id)) { + return p; + } + return cache::expose(this, id); + } + private: // The step's bits as SDPA wants them: [1, 1, length, read_len], one row per // query token. diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index d9b77e8f116..001c3fd4bf2 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -293,7 +293,15 @@ inline void exec_sdpa(const SdpaNode& n, ExecutionState& st, StreamOrDevice s) { } array out = fast::scaled_dot_product_attention( - Q, K, V, static_cast(n.scale), mask_mode, mask_arr, sinks, s); + Q, + K, + V, + static_cast(n.scale), + mask_mode, + mask_arr, + sinks, + false, + s); st.set_tensor(n.out, std::move(out)); } @@ -382,6 +390,7 @@ inline void exec_update_and_attend( mask_mode, spec.mask, std::nullopt, + false, s); // Honor the op's output-dtype contract (unset -> SDPA's native output). if (n.out_dtype) { diff --git a/backends/mlx/runtime/MLXSequenceCache.h b/backends/mlx/runtime/MLXSequenceCache.h index 79beeeb846e..3680b16087c 100644 --- a/backends/mlx/runtime/MLXSequenceCache.h +++ b/backends/mlx/runtime/MLXSequenceCache.h @@ -125,6 +125,14 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache { return AttendSpec{K, V, AttendSpec::Mask::Causal, std::nullopt}; } + protected: + void* face(cache::FaceId id) override { + if (void* p = cache::SequenceCache::face(id)) { + return p; + } + return cache::expose(this, id); + } + private: // A sequence cache holds one run of one sequence, so the step is described by // where it starts; the remaining positions carry no information beyond diff --git a/backends/mlx/runtime/SwiftPMMetallibPath.h b/backends/mlx/runtime/SwiftPMMetallibPath.h new file mode 100644 index 00000000000..8bc838cf93d --- /dev/null +++ b/backends/mlx/runtime/SwiftPMMetallibPath.h @@ -0,0 +1,24 @@ +// 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 + +#include +#include +#include + +namespace executorch::backends::mlx { + +// Finds the current platform's metallib in the named SwiftPM resource bundle. +// Each path may be a containing directory, a code bundle, or the resource +// bundle itself. Exposed for focused path tests. +std::optional find_swiftpm_metallib_path( + const std::vector& container_paths); + +// Resolves the current platform's metallib from loaded Apple bundles. +std::optional resolve_swiftpm_metallib_path(); + +} // namespace executorch::backends::mlx diff --git a/backends/mlx/runtime/SwiftPMMetallibPath.mm b/backends/mlx/runtime/SwiftPMMetallibPath.mm new file mode 100644 index 00000000000..977a35b4eee --- /dev/null +++ b/backends/mlx/runtime/SwiftPMMetallibPath.mm @@ -0,0 +1,152 @@ +// 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. + +#import +#import + +#include "SwiftPMMetallibPath.h" + +#include +#include + +namespace executorch::backends::mlx { +namespace { + +constexpr const char* kResourceBundleName = + "executorch_backend_mlx_resources.bundle"; + +const char* metallib_filename() { +#if TARGET_OS_SIMULATOR + return "mlx-ios-simulator.metallib"; +#elif TARGET_OS_IOS + return "mlx-ios.metallib"; +#elif TARGET_OS_OSX + return "mlx-macos.metallib"; +#else + return nullptr; +#endif +} + +std::optional regular_file_path(NSURL* url) { + if (url == nil || !url.fileURL || url.path == nil) { + return std::nullopt; + } + + std::error_code error; + const std::filesystem::path path(url.fileSystemRepresentation); + if (!std::filesystem::is_regular_file(path, error)) { + return std::nullopt; + } + return path.string(); +} + +std::optional find_in_resource_bundle(NSURL* bundle_url) { + if (bundle_url == nil || !bundle_url.fileURL) { + return std::nullopt; + } + + NSBundle* bundle = [NSBundle bundleWithURL:bundle_url]; + if (bundle == nil) { + return std::nullopt; + } + + NSString* filename = [NSString stringWithUTF8String:metallib_filename()]; + if (filename == nil) { + return std::nullopt; + } + + if (auto path = regular_file_path( + [bundle URLForResource:filename.stringByDeletingPathExtension + withExtension:filename.pathExtension])) { + return path; + } + + // SwiftPM's native build system can emit a flat resource bundle. Check the + // bundle root explicitly in addition to Foundation's platform resource URL. + return regular_file_path([bundle_url URLByAppendingPathComponent:filename]); +} + +std::optional find_from_container(NSURL* container_url) { + if (container_url == nil || !container_url.fileURL) { + return std::nullopt; + } + + NSString* resource_bundle_name = + [NSString stringWithUTF8String:kResourceBundleName]; + if ([container_url.lastPathComponent isEqualToString:resource_bundle_name]) { + return find_in_resource_bundle(container_url); + } + + NSBundle* container_bundle = [NSBundle bundleWithURL:container_url]; + if (container_bundle != nil) { + NSURL* resource_bundle_url = [container_bundle + URLForResource:resource_bundle_name.stringByDeletingPathExtension + withExtension:resource_bundle_name.pathExtension]; + if (auto path = find_in_resource_bundle(resource_bundle_url)) { + return path; + } + } + + return find_in_resource_bundle( + [container_url URLByAppendingPathComponent:resource_bundle_name]); +} + +void append_path(NSMutableOrderedSet* paths, NSURL* url) { + if (url != nil && url.fileURL && url.path != nil) { + [paths addObject:url.path]; + } +} + +} // namespace + +std::optional find_swiftpm_metallib_path( + const std::vector& container_paths) { + if (metallib_filename() == nullptr) { + return std::nullopt; + } + + @autoreleasepool { + for (const auto& container_path : container_paths) { + NSString* path = [NSString stringWithUTF8String:container_path.c_str()]; + if (path == nil) { + continue; + } + if (auto metallib_path = + find_from_container([NSURL fileURLWithPath:path])) { + return metallib_path; + } + } + } + + return std::nullopt; +} + +std::optional resolve_swiftpm_metallib_path() { + @autoreleasepool { + NSMutableOrderedSet* paths = [NSMutableOrderedSet orderedSet]; + NSBundle* main_bundle = NSBundle.mainBundle; + append_path(paths, main_bundle.bundleURL); + append_path(paths, main_bundle.resourceURL); + + for (NSBundle* bundle in NSBundle.allBundles) { + append_path(paths, bundle.bundleURL); + append_path(paths, bundle.resourceURL); + } + for (NSBundle* framework in NSBundle.allFrameworks) { + append_path(paths, framework.bundleURL); + append_path(paths, framework.resourceURL); + } + + std::vector container_paths; + container_paths.reserve(paths.count); + for (NSString* path in paths) { + container_paths.emplace_back(path.fileSystemRepresentation); + } + return find_swiftpm_metallib_path(container_paths); + } +} + +} // namespace executorch::backends::mlx diff --git a/backends/mlx/runtime/backend_options.h b/backends/mlx/runtime/backend_options.h index af9a993ce7e..2f96ce93525 100644 --- a/backends/mlx/runtime/backend_options.h +++ b/backends/mlx/runtime/backend_options.h @@ -42,14 +42,6 @@ inline constexpr char kClearCacheIntervalKey[] = "clear_cache_interval"; // errors otherwise). Saves one full mutable-buffer (KV-cache) copy per handle. inline constexpr char kSkipMutableBufferInitKey[] = "skip_mutable_buffer_init"; -// Per-model runtime-spec key (string). Names the off-graph KV cache this handle -// binds to: the runner creates the cache, installs it in the process-global -// CacheRegistry under this key, and the delegate looks it up in init(). The -// DelegateHandle is opaque to the host, so the key is the only rendezvous -// channel. Unset means no cache, and any update_and_attend node then fails at -// execute() rather than silently attending nothing. -inline constexpr char kCacheKeyKey[] = "cache_key"; - } // namespace mlx } // namespace backends } // namespace executorch diff --git a/backends/mlx/test/CMakeLists.txt b/backends/mlx/test/CMakeLists.txt index 377818ec75a..44693ad8c66 100644 --- a/backends/mlx/test/CMakeLists.txt +++ b/backends/mlx/test/CMakeLists.txt @@ -53,6 +53,32 @@ add_dependencies(op_test_runner strict_compile_test) # Multi-threaded inference test include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) +if(APPLE) + et_cxx_test( + mlx_metallib_path_test SOURCES + ${CMAKE_CURRENT_LIST_DIR}/mlx_metallib_path_test.mm + ${CMAKE_CURRENT_LIST_DIR}/../runtime/SwiftPMMetallibPath.mm + ) + target_include_directories( + mlx_metallib_path_test PRIVATE ${CMAKE_CURRENT_LIST_DIR}/../runtime + ) + set_source_files_properties( + ${CMAKE_CURRENT_LIST_DIR}/mlx_metallib_path_test.mm + ${CMAKE_CURRENT_LIST_DIR}/../runtime/SwiftPMMetallibPath.mm + PROPERTIES COMPILE_FLAGS "-fobjc-arc" + ) + if(EXECUTORCH_MLX_ENABLE_SANITIZERS) + target_compile_options( + mlx_metallib_path_test PRIVATE -fsanitize=address,undefined + -fno-omit-frame-pointer + ) + target_link_options( + mlx_metallib_path_test PRIVATE ${_mlx_sanitizer_link_options} + ) + endif() + target_link_libraries(mlx_metallib_path_test ${FOUNDATION_FRAMEWORK}) +endif() + et_cxx_test( multi_thread_test_runner SOURCES diff --git a/backends/mlx/test/mlx_cell_cache_test.cpp b/backends/mlx/test/mlx_cell_cache_test.cpp index a93590b453e..f808965d1a5 100644 --- a/backends/mlx/test/mlx_cell_cache_test.cpp +++ b/backends/mlx/test/mlx_cell_cache_test.cpp @@ -224,8 +224,9 @@ TEST_F(MLXCellCacheTest, StorageDtypeDiffersCastsOnWrite) { EXPECT_TRUE(allclose(spec.K, k, 1e-2f)); } -// The step verbs are a contract: no declaration, a miscounted call, a repeated -// layer and a position a sequence already holds are all refused. +// The step verbs are a contract: no declaration, a miscounted call, and a +// position a sequence already holds are refused. A repeated layer with the same +// tokens (a KV-shared donor re-serving) is served again idempotently. TEST_F(MLXCellCacheTest, IllFormedStepsThrow) { using namespace ::mlx::core; MLXCellCache c(flat_config(32, 1, H, D, kHalf)); @@ -239,8 +240,9 @@ TEST_F(MLXCellCacheTest, IllFormedStepsThrow) { EXPECT_ANY_THROW(c.update_and_fetch(1, {0, 1}, k, v, s)); // no such layer c.update_and_fetch(0, {0, 1}, k, v, s); - EXPECT_ANY_THROW( - c.update_and_fetch(0, {0, 1}, k, v, s)); // layer served twice + // A KV-shared layer re-serves its donor's id with the same tokens; the repeat + // is idempotent and returns the same step rather than throwing. + EXPECT_NO_THROW(c.update_and_fetch(0, {0, 1}, k, v, s)); EXPECT_TRUE(c.declare_step({a})); array k1 = randn(1), v1 = randn(1); @@ -268,12 +270,16 @@ TEST_F(MLXCellCacheTest, InvalidConfigThrows) { // A runner reaches a layout by (backend_id, kind), so the builder registration // is as much a part of the layout as the class. TEST_F(MLXCellCacheTest, RegistryBuildsCellLayout) { - auto built = cache::CacheBuilderRegistry::global().build( - kMLXBackendId, "cell", flat_config(32, 1, H, D, kHalf)); + auto built = cache::CacheFactory::global().build( + kMLXBackendId, + cache::kind::kBatchedCell, + flat_config(32, 1, H, D, kHalf)); ASSERT_TRUE(built.ok()); - const std::shared_ptr& c = *built; - EXPECT_NE(c->as_batch_control(), nullptr); - EXPECT_EQ(c->as_control(), nullptr); + const std::shared_ptr& c = *built; + EXPECT_NE(c->as(), nullptr); + EXPECT_NE(c->as(), nullptr) << "the backend face comes back too"; + // A cell layout is multi-sequence, so it offers no single-sequence face. + EXPECT_EQ(c->as(), nullptr); } } // namespace diff --git a/backends/mlx/test/mlx_metallib_path_test.mm b/backends/mlx/test/mlx_metallib_path_test.mm new file mode 100644 index 00000000000..17f15298636 --- /dev/null +++ b/backends/mlx/test/mlx_metallib_path_test.mm @@ -0,0 +1,123 @@ +// 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. + +#import +#import + +#include "SwiftPMMetallibPath.h" + +#include + +#include +#include +#include +#include + +namespace executorch::backends::mlx { +namespace { + +const char* expected_metallib_filename() { +#if TARGET_OS_SIMULATOR + return "mlx-ios-simulator.metallib"; +#elif TARGET_OS_IOS + return "mlx-ios.metallib"; +#else + return "mlx-macos.metallib"; +#endif +} + +const char* wrong_metallib_filename() { +#if TARGET_OS_OSX + return "mlx-ios.metallib"; +#else + return "mlx-macos.metallib"; +#endif +} + +class MLXMetallibPathTest : public ::testing::Test { + protected: + void SetUp() override { + NSString* name = [NSString + stringWithFormat:@"executorch_mlx_metallib_path_%@", + NSUUID.UUID.UUIDString]; + root_ = std::filesystem::temp_directory_path() / + std::string(name.fileSystemRepresentation); + ASSERT_TRUE(std::filesystem::create_directories(root_)); + } + + void TearDown() override { + std::error_code error; + std::filesystem::remove_all(root_, error); + } + + std::filesystem::path create_bundle(bool deep) { + const auto bundle = root_ / "executorch_backend_mlx_resources.bundle"; + const auto contents = deep ? bundle / "Contents" : bundle; + const auto resources = deep ? contents / "Resources" : bundle; + EXPECT_TRUE(std::filesystem::create_directories(resources)); + + std::ofstream(contents / "Info.plist") + << "\n" + << "\n" + << "" + << "CFBundleIdentifier" + << "org.pytorch.executorch.mlx-test-resources" + << "CFBundlePackageTypeBNDL" + << "\n"; + return resources; + } + + std::filesystem::path root_; +}; + +TEST_F(MLXMetallibPathTest, MissingBundleReturnsNoPath) { + EXPECT_FALSE(find_swiftpm_metallib_path({root_.string()}).has_value()); +} + +TEST_F(MLXMetallibPathTest, ProcessWithoutSwiftPMBundleReturnsNoPath) { + EXPECT_FALSE(resolve_swiftpm_metallib_path().has_value()); +} + +TEST_F(MLXMetallibPathTest, FindsFlatBundleResource) { + const auto resources = create_bundle(/*deep=*/false); + const auto metallib = resources / expected_metallib_filename(); + std::ofstream(metallib) << "fixture"; + + EXPECT_EQ(find_swiftpm_metallib_path({root_.string()}), metallib.string()); +} + +TEST_F(MLXMetallibPathTest, FindsMacOSDeepBundleResource) { + const auto resources = create_bundle(/*deep=*/true); + const auto metallib = resources / expected_metallib_filename(); + std::ofstream(metallib) << "fixture"; + + EXPECT_EQ(find_swiftpm_metallib_path({root_.string()}), metallib.string()); + EXPECT_EQ( + find_swiftpm_metallib_path( + {(root_ / "executorch_backend_mlx_resources.bundle").string()}), + metallib.string()); +} + +TEST_F(MLXMetallibPathTest, SelectsCurrentPlatformSlice) { + const auto resources = create_bundle(/*deep=*/true); + const auto expected = resources / expected_metallib_filename(); + const auto wrong = resources / wrong_metallib_filename(); + std::ofstream(expected) << "expected"; + std::ofstream(wrong) << "wrong"; + + EXPECT_EQ(find_swiftpm_metallib_path({root_.string()}), expected.string()); +} + +TEST_F(MLXMetallibPathTest, IgnoresWrongPlatformSlice) { + const auto resources = create_bundle(/*deep=*/true); + std::ofstream(resources / wrong_metallib_filename()) << "fixture"; + + EXPECT_FALSE(find_swiftpm_metallib_path({root_.string()}).has_value()); +} + +} // namespace +} // namespace executorch::backends::mlx diff --git a/backends/mlx/test/op_test_runner.cpp b/backends/mlx/test/op_test_runner.cpp index 53291e41064..12eb7235ce5 100644 --- a/backends/mlx/test/op_test_runner.cpp +++ b/backends/mlx/test/op_test_runner.cpp @@ -300,38 +300,33 @@ int main(int argc, char* argv[]) { namespace cache = ::executorch::extension::llm::cache; - // Build and install the off-graph KV cache before the Module, so the - // registry entry exists by the time the delegate's init() looks it up. - // Declared here so the session outlives the module. - std::optional cache_session; + // Publish the off-graph KV cache until the delegate resolves its key while + // loading the method. + std::optional cache_install_guard; if (!kv_cache_spec.empty()) { cache::CacheConfig cfg{}; if (!parse_kv_cache_spec(kv_cache_spec, cfg)) { std::cerr << "Invalid --kv-cache spec: " << kv_cache_spec << std::endl; return 1; } - auto built = cache::CacheBuilderRegistry::global().build( - ::executorch::backends::mlx::kMLXBackendId, "seq", cfg); + auto built = cache::CacheFactory::global().build( + ::executorch::backends::mlx::kMLXBackendId, + cache::kind::kSingle, + cfg); if (!built.ok()) { std::cerr << "Failed to build KV cache: " << static_cast(built.error()) << std::endl; return 1; } - cache_session.emplace(cache::make_unique_key(), built.get()); - if (verbose) { - std::cout << "Installed KV cache under key " << cache_session->key() - << std::endl; - } + cache_install_guard.emplace(built.get()); } Module module(pte_path); Error load_error = Error::Ok; - if (cache_session) { + if (cache_install_guard) { ::executorch::runtime::BackendOptions<1> mlx_opts; ::executorch::runtime::LoadBackendOptionsMap options_map; - if (mlx_opts.set_option( - ::executorch::backends::mlx::kCacheKeyKey, - cache_session->key().c_str()) != Error::Ok || + if (cache_install_guard->set_option(mlx_opts) != Error::Ok || options_map.set_options( ::executorch::backends::mlx::kMLXBackendId, mlx_opts.view()) != Error::Ok) { @@ -358,6 +353,7 @@ int main(int argc, char* argv[]) { << static_cast(load_method_error) << std::endl; return 1; } + cache_install_guard.reset(); if (verbose) { std::cout << "Reading inputs from: " << input_path << std::endl; diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py index 452ae37c40b..0d5c72af06b 100644 --- a/backends/mlx/test/test_ops.py +++ b/backends/mlx/test/test_ops.py @@ -6212,6 +6212,7 @@ def __init__( is_causal: bool = False, use_mask: bool = False, use_bool_mask: bool = False, + kv_seq_len: Optional[int] = None, ): self.batch_size = batch_size self.num_heads = num_heads @@ -6221,12 +6222,15 @@ def __init__( self.is_causal = is_causal self.use_mask = use_mask self.use_bool_mask = use_bool_mask + self.kv_seq_len = kv_seq_len if kv_seq_len is not None else seq_len parts = ["sdpa"] if num_kv_heads is not None: parts.append(f"gqa{num_kv_heads}") if is_causal: parts.append("causal") + if self.kv_seq_len != seq_len: + parts.append(f"q{seq_len}kv{self.kv_seq_len}") if use_mask: parts.append("mask") if use_bool_mask: @@ -6241,6 +6245,11 @@ def get_test_configs(cls) -> List["SDPATest"]: cls(num_kv_heads=4), cls(use_mask=True), cls(use_bool_mask=True), # Test boolean mask conversion + # A decode step against a longer key cache. MLX anchors its causal mask at + # the bottom right and torch at the top left, so they only agree when the + # lengths match. + cls(is_causal=True, seq_len=1, kv_seq_len=32), + cls(is_causal=True, seq_len=6, kv_seq_len=32), ] def create_model(self) -> nn.Module: @@ -6256,25 +6265,49 @@ def create_model(self) -> nn.Module: def create_inputs(self) -> Tuple[torch.Tensor, ...]: q = torch.randn(self.batch_size, self.num_heads, self.seq_len, self.head_dim) kv_heads = self.num_kv_heads if self.num_kv_heads else self.num_heads - k = torch.randn(self.batch_size, kv_heads, self.seq_len, self.head_dim) - v = torch.randn(self.batch_size, kv_heads, self.seq_len, self.head_dim) + k = torch.randn(self.batch_size, kv_heads, self.kv_seq_len, self.head_dim) + v = torch.randn(self.batch_size, kv_heads, self.kv_seq_len, self.head_dim) if self.use_mask: # Additive float mask: 0 = attend, -inf = masked - mask = torch.zeros(self.batch_size, 1, self.seq_len, self.seq_len) - mask[:, :, :, : self.seq_len // 4] = float("-inf") + mask = torch.zeros(self.batch_size, 1, self.seq_len, self.kv_seq_len) + mask[:, :, :, : self.kv_seq_len // 4] = float("-inf") return (q, k, v, mask) elif self.use_bool_mask: # Boolean mask: True = attend, False = masked # This tests that the backend correctly converts bool -> additive format mask = torch.ones( - self.batch_size, 1, self.seq_len, self.seq_len, dtype=torch.bool + self.batch_size, 1, self.seq_len, self.kv_seq_len, dtype=torch.bool ) - mask[:, :, :, : self.seq_len // 4] = False # Mask out first quarter + mask[:, :, :, : self.kv_seq_len // 4] = False # Mask out first quarter return (q, k, v, mask) return (q, k, v) +@register_test +class SDPARank3Test(OpTestCase): + """Attention on rank-3 tensors, which PyTorch accepts and the fused kernel does not. + + The node counts are the point of the test: they assert the fused kernel is still + used, rather than the operator having been decomposed into primitives. + """ + + name = "sdpa_rank3" + rtol = 1e-3 + atol = 1e-3 + expected_node_counts = { + "SdpaNode": 1, + "ExpandDimsNode": 3, + "SqueezeNode": 1, + } + + def create_model(self) -> nn.Module: + return SDPAModel() + + def create_inputs(self) -> Tuple[torch.Tensor, ...]: + return tuple(torch.randn(2, 16, 64) for _ in range(3)) + + class CustomSDPAModel(nn.Module): """ Test model for mlx::custom_sdpa with KVCache. @@ -8565,14 +8598,14 @@ def compute_expected_outputs(self, model, test_inputs): # an oracle cache installed for its duration. from executorch.extension.llm.cache.reference_cache import ( CacheConfig, - ContiguousReferenceCache, + SequenceReferenceCache, ) from executorch.extension.llm.cache.update_and_attend import REGISTRY key = f"{self.name}-oracle" REGISTRY.install( key, - ContiguousReferenceCache( + SequenceReferenceCache( CacheConfig( n_layers=self.n_layers, n_kv_heads=self.n_kv_heads, diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index 4a5833aa656..3b82e306b1f 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -9,12 +9,16 @@ Tests for the MLX partitioner. """ +import tempfile import unittest +from pathlib import Path import torch import torch.nn as nn from executorch.backends.mlx.partitioner import MLXPartitioner -from executorch.exir import EdgeCompileConfig, to_edge +from executorch.backends.mlx.test.test_utils import get_mlx_node_counts +from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower +from executorch.runtime import Runtime from torch.export import export @@ -41,5 +45,230 @@ def forward(self, x): self.assertIn("to_edge_transform_and_lower", str(ctx.exception)) +def _lower(model, inputs): + return to_edge_transform_and_lower( + export(model, inputs, strict=False), + partitioner=[MLXPartitioner()], + ).to_executorch() + + +def _delegate_count(program) -> int: + return sum( + 1 + for node in program.exported_program().graph_module.graph.nodes + if node.op == "call_function" and "executorch_call_delegate" in str(node.target) + ) + + +def _run(model, inputs): + """Lower, execute, and return the node counts, the delegate count and the error. + + The delegate count is returned so a test can tell "decomposed onto this backend" + apart from "not lowered here at all", which a node count alone cannot show. + """ + with torch.no_grad(): + ref = model(*inputs) + program = _lower(model, inputs) + delegates = _delegate_count(program) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "model.pte" + path.write_bytes(program.buffer) + counts = get_mlx_node_counts(path) + method = Runtime.get().load_program(path).load_method("forward") + out = method.execute(list(inputs))[0] + return counts, delegates, (out - ref).abs().max().item() + + +class Sdpa(nn.Module): + def __init__(self, is_causal: bool = False): + super().__init__() + self.is_causal = is_causal + + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, k, v, is_causal=self.is_causal + ) + + +class GroupedSdpa(nn.Module): + """Grouped key/value attention, where the repeat is unwrapped before the kernel.""" + + def __init__(self, dim: int, is_causal: bool = False): + super().__init__() + self.dim = dim + self.is_causal = is_causal + + def forward(self, q, k, v): + k = k.repeat_interleave(2, dim=self.dim) + v = v.repeat_interleave(2, dim=self.dim) + return torch.nn.functional.scaled_dot_product_attention( + q, k, v, is_causal=self.is_causal + ) + + +class TestMLXPartitionerSdpaShapes(unittest.TestCase): + """The fused kernel takes rank 4, so other ranks are adapted or left alone.""" + + def test_rank4_is_unchanged(self): + counts, _, err = _run( + Sdpa(), tuple(torch.randn(2, 4, 16, 64) for _ in range(3)) + ) + self.assertEqual(counts.get("SdpaNode", 0), 1) + self.assertEqual(counts.get("ExpandDimsNode", 0), 0) + self.assertEqual(counts.get("SqueezeNode", 0), 0) + self.assertLess(err, 1e-4) + + def test_rank3_is_lifted_once(self): + counts, _, err = _run(Sdpa(), tuple(torch.randn(2, 16, 64) for _ in range(3))) + self.assertEqual(counts.get("SdpaNode", 0), 1) + self.assertEqual(counts.get("ExpandDimsNode", 0), 3) + self.assertEqual(counts.get("SqueezeNode", 0), 1) + self.assertLess(err, 1e-4) + + def test_rank2_is_lifted_twice(self): + counts, _, err = _run(Sdpa(), tuple(torch.randn(16, 64) for _ in range(3))) + self.assertEqual(counts.get("SdpaNode", 0), 1) + self.assertEqual(counts.get("ExpandDimsNode", 0), 6) + self.assertEqual(counts.get("SqueezeNode", 0), 1) + self.assertLess(err, 1e-4) + + def test_rank5_is_decomposed_on_this_backend(self): + # Folding the leading dimensions pairs the wrong operands once one of them + # broadcasts a batch, so this decomposes rather than fusing. + counts, delegates, err = _run( + Sdpa(), tuple(torch.randn(2, 2, 4, 16, 64) for _ in range(3)) + ) + self.assertEqual(counts.get("SdpaNode", 0), 0) + self.assertGreater(delegates, 0) + self.assertLess(err, 1e-4) + + def test_unequal_batch_is_decomposed_on_this_backend(self): + counts, delegates, err = _run( + Sdpa(), + ( + torch.randn(2, 4, 16, 64), + torch.randn(1, 4, 16, 64), + torch.randn(1, 4, 16, 64), + ), + ) + self.assertEqual(counts.get("SdpaNode", 0), 0) + self.assertGreater(delegates, 0) + self.assertLess(err, 1e-4) + + def test_zero_head_count_declines_instead_of_raising(self): + # The head multiple test would divide by zero here, and raising from the + # matcher aborts the whole export rather than declining this one node. Only + # lowering is checked: a zero-size operand is not executable either way. + program = _lower( + Sdpa(), + ( + torch.randn(1, 4, 8, 16), + torch.randn(1, 0, 8, 16), + torch.randn(1, 0, 8, 16), + ), + ) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "model.pte" + path.write_bytes(program.buffer) + self.assertEqual(get_mlx_node_counts(path).get("SdpaNode", 0), 0) + + +class TestMLXPartitionerGroupedKeys(unittest.TestCase): + """The grouped key/value unwrap reads dim 1 as the head, which holds at rank 4.""" + + def test_rank4_head_repeat_is_absorbed(self): + counts, _, err = _run( + GroupedSdpa(dim=1), + ( + torch.randn(2, 4, 16, 64), + torch.randn(2, 2, 16, 64), + torch.randn(2, 2, 16, 64), + ), + ) + self.assertEqual(counts.get("SdpaNode", 0), 1) + self.assertEqual(counts.get("RepeatNode", 0), 0) + self.assertLess(err, 1e-4) + + def test_rank3_sequence_repeat_is_kept(self): + # At rank 3 dim 1 is the key sequence, so absorbing the repeat would drop + # half the keys. Without a mask that still sums correctly, which is what + # makes it easy to miss; with a causal mask it is wrong by whole units. + counts, _, err = _run( + GroupedSdpa(dim=1, is_causal=True), + (torch.randn(2, 16, 64), torch.randn(2, 8, 64), torch.randn(2, 8, 64)), + ) + self.assertEqual(counts.get("RepeatNode", 0), 2) + self.assertLess(err, 1e-4) + + +class TestMLXPartitionerSdpaCausal(unittest.TestCase): + """MLX anchors a causal mask at the bottom right and torch at the top left.""" + + def test_equal_lengths_stay_fused(self): + counts, _, err = _run( + Sdpa(is_causal=True), tuple(torch.randn(1, 4, 16, 64) for _ in range(3)) + ) + self.assertEqual(counts.get("SdpaNode", 0), 1) + self.assertLess(err, 1e-4) + + def test_rank3_equal_lengths_are_lifted(self): + counts, _, err = _run( + Sdpa(is_causal=True), tuple(torch.randn(2, 16, 64) for _ in range(3)) + ) + self.assertEqual(counts.get("SdpaNode", 0), 1) + self.assertEqual(counts.get("ExpandDimsNode", 0), 3) + self.assertLess(err, 1e-4) + + def test_rank3_unequal_lengths_are_not_lifted(self): + # The two conventions disagree here and the disagreement is silent, so a + # shape this backend could not previously reach is not opened up. + counts, delegates, err = _run( + Sdpa(is_causal=True), + (torch.randn(2, 6, 64), torch.randn(2, 16, 64), torch.randn(2, 16, 64)), + ) + self.assertEqual(counts.get("SdpaNode", 0), 0) + self.assertGreater(delegates, 0) + self.assertLess(err, 1e-4) + + def test_rank2_unequal_lengths_are_not_lifted(self): + counts, _, err = _run( + Sdpa(is_causal=True), + (torch.randn(6, 64), torch.randn(16, 64), torch.randn(16, 64)), + ) + self.assertEqual(counts.get("SdpaNode", 0), 0) + self.assertLess(err, 1e-4) + + +class TestMLXPartitionerMixedSupport(unittest.TestCase): + """An operator is preserved from decomposition per operator, not per call.""" + + def test_supported_and_unsupported_calls_in_one_graph(self): + class Mixed(nn.Module): + def forward(self, a, b): + x = torch.nn.functional.scaled_dot_product_attention(a, a, a) + y = torch.nn.functional.scaled_dot_product_attention(b, b, b) + return x.sum() + y.sum() + + # Without giving the whole operator back, the rank-5 call would be neither + # lowered nor decomposed and this would raise a missing out variant. + counts, delegates, err = _run( + Mixed().eval(), + (torch.randn(1, 4, 16, 64), torch.randn(2, 2, 4, 16, 64)), + ) + # The cost of the coarse choice: the supported call is unfused as well. + self.assertEqual(counts.get("SdpaNode", 0), 0) + self.assertGreater(delegates, 0) + self.assertLess(err, 1e-3) + + def test_mixed_support_outside_attention(self): + class TwoRolls(nn.Module): + def forward(self, x): + return torch.roll(x, 1, dims=0).sum() + torch.roll(x, 1).sum() + + _, delegates, err = _run(TwoRolls().eval(), (torch.randn(4, 8),)) + self.assertGreater(delegates, 0) + self.assertLess(err, 1e-4) + + if __name__ == "__main__": unittest.main() diff --git a/backends/mlx/test/test_utils.py b/backends/mlx/test/test_utils.py index eb20095db97..adf4e56e3a2 100644 --- a/backends/mlx/test/test_utils.py +++ b/backends/mlx/test/test_utils.py @@ -587,7 +587,14 @@ def rebuild_op_test_runner(verbose: bool = False) -> bool: print(f"Rebuilding op_test_runner in {build_dir}...") - cmd = ["cmake", "--build", str(build_dir), "--target", "op_test_runner", "-j8"] + cmd = [ + "cmake", + "--build", + str(build_dir), + "--target", + "op_test_runner", + f"-j{(os.cpu_count() or 1) + 1}", + ] if verbose: print(f"Running: {' '.join(cmd)}") diff --git a/backends/mlx/third-party/mlx b/backends/mlx/third-party/mlx index 7a1d4f5c12a..1f8e74e3f12 160000 --- a/backends/mlx/third-party/mlx +++ b/backends/mlx/third-party/mlx @@ -1 +1 @@ -Subproject commit 7a1d4f5c12ac82f4b4d0a6e71538d89ca0605247 +Subproject commit 1f8e74e3f12f31365464a6867c6579f0e9b29d85 diff --git a/backends/native/runtime/Deserialize.cpp b/backends/native/runtime/Deserialize.cpp new file mode 100644 index 00000000000..75410464b7b --- /dev/null +++ b/backends/native/runtime/Deserialize.cpp @@ -0,0 +1,617 @@ +// 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. + +// The deserializer bridge: native_backend::Program (FlatBuffer) -> ptn +// in-memory IR (Method / Graph / Node / Argument / Value). In-graph references +// (SSA names) are resolved to list ValueIds; per-graph namespaces are +// resolved independently (each HOP subgraph rebuilds its own name -> id map). +// +// Everything here runs on a buffer Program::load() has already put through +// flatbuffers::Verifier, so accessors return non-null wherever the schema +// declares the field required and wherever a union discriminator matches; the +// walkers below dereference those results directly. Fields the schema leaves +// optional are still checked, because verification says nothing about whether +// they are present. Vector>::Get() computes an address rather than a +// nullable pointer, and verification bounds-checks every referenced object. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +namespace ptn { +namespace { + +std::string str_of(const flatbuffers::String* s) { + return s != nullptr ? s->str() : std::string(); +} + +bool nonempty(const flatbuffers::String* s) { + return s != nullptr && s->size() > 0; +} + +ScalarType map_scalar_type(fbs::ScalarType t) { + // ptn::ScalarType ids are pinned to the schema's, so the byte maps straight. + return static_cast(static_cast(t)); +} + +OpKind map_op_kind(fbs::OpKind k) { + switch (k) { + case fbs::OpKind::CALL_FUNCTION: + return OpKind::CallFunction; + case fbs::OpKind::PLACEHOLDER: + return OpKind::Placeholder; + case fbs::OpKind::OUTPUT: + return OpKind::Output; + default: + throw std::runtime_error( + "build_graph: unsupported OpKind value " + + std::to_string(static_cast(k))); + } +} + +OutputValueKind map_output_value_kind(fbs::OutputValueKind k) { + switch (k) { + case fbs::OutputValueKind::TENSOR: + return OutputValueKind::Tensor; + case fbs::OutputValueKind::TENSOR_LIST: + return OutputValueKind::TensorList; + case fbs::OutputValueKind::INT: + return OutputValueKind::Int; + case fbs::OutputValueKind::BOOL: + return OutputValueKind::Bool; + case fbs::OutputValueKind::FLOAT: + return OutputValueKind::Float; + default: + throw std::runtime_error( + "build_graph: unsupported OutputValueKind value " + + std::to_string(static_cast(k))); + } +} + +ValueRole map_input_kind(fbs::InputKind k) { + switch (k) { + case fbs::InputKind::USER_INPUT: + return ValueRole::UserInput; + case fbs::InputKind::PARAMETER: + return ValueRole::Parameter; + case fbs::InputKind::BUFFER: + return ValueRole::Buffer; + case fbs::InputKind::CONSTANT_TENSOR: + return ValueRole::ConstantTensor; + default: + throw std::runtime_error( + "build_method: unsupported InputKind value " + + std::to_string(static_cast(k))); + } +} + +OutputKind map_output_kind(fbs::OutputKind k) { + switch (k) { + case fbs::OutputKind::USER_OUTPUT: + return OutputKind::UserOutput; + case fbs::OutputKind::BUFFER_MUTATION: + return OutputKind::BufferMutation; + case fbs::OutputKind::USER_INPUT_MUTATION: + return OutputKind::UserInputMutation; + default: + throw std::runtime_error( + "build_method: unsupported OutputKind value " + + std::to_string(static_cast(k))); + } +} + +// The wire describes a dim as a min..max range, while the IR holds a concrete +// extent. Collapsing a range to its upper bound would run the graph at that +// bound and compute over elements the caller never supplied, so a dim that is +// not a single non-negative extent is refused where it enters the IR. +int64_t static_extent( + const fbs::Dim* d, + const std::string& value_name, + flatbuffers::uoffset_t i) { + if (d->min() == d->max() && d->min() >= 0) { + return d->min(); + } + throw std::runtime_error( + "build_tensor_meta: " + value_name + " dim " + std::to_string(i) + + " is not a static extent (" + std::to_string(d->min()) + ".." + + (d->max() < 0 ? std::string("inf") : std::to_string(d->max())) + + "); this runtime requires static shapes"); +} + +TensorMeta build_tensor_meta( + const fbs::TensorMeta* m, + const std::string& name) { + TensorMeta out; + if (m == nullptr) { + return out; + } + out.dtype = map_scalar_type(m->dtype()); + if (const auto* sizes = m->sizes()) { + out.sizes.reserve(sizes->size()); + for (flatbuffers::uoffset_t i = 0; i < sizes->size(); ++i) { + out.sizes.push_back(static_extent(sizes->Get(i), name, i)); + } + } + if (const auto* dord = m->dim_order()) { + out.dim_order_hint.reserve(dord->size()); + for (flatbuffers::uoffset_t i = 0; i < dord->size(); ++i) { + out.dim_order_hint.push_back(static_cast(dord->Get(i))); + } + } + return out; +} + +// value name -> tensor metadata. +using MetaTable = std::unordered_map; + +// One graph body plus the SSA-name -> value-id map used to build it. +// +// The wire format addresses values two ways: a graph body is positional, and +// the value list keeps that (a ValueId is an index), while the method-level +// side tables -- constants, mutable_buffers, output_specs -- name their targets +// by SSA string, since the serializer writes them independently of the body. +// Only the builder knows how one maps to the other, so it hands the map back +// for build_method to resolve those names against. +// +// Build-time scaffolding: nothing outside build_method sees it, and neither +// Method nor Graph stores it. Once the bindings are resolved to ids it is +// discarded, and the graph is index-addressed from then on. +struct BuiltGraph { + Graph graph; + std::unordered_map name_to_id; +}; + +// `extra_meta` supplies metadata for values the graph's own tensor_values side +// table omits: a constant placeholder's meta rides on the Method's +// NamedTensorRef binding instead, so build_method passes it in to type those +// values on creation. Subgraphs have no such bindings. +BuiltGraph build_graph(const fbs::Graph* g, const MetaTable& extra_meta = {}); + +// Builds one Graph body. The graph under construction, its SSA-name -> id map, +// and the name -> metadata table are shared by every step of the build, so they +// are members rather than threaded through each call. One builder per body: a +// subgraph gets a fresh one, which is what gives each body its own independent +// SSA namespace. +class GraphBuilder { + public: + BuiltGraph run(const fbs::Graph* g, const MetaTable& extra_meta); + + private: + // Resolve a name to its ValueId, creating the Value on first mention. + // A name in the meta side table becomes a Tensor value; otherwise a None + // value (scalar / symbolic outputs, refined once the loader models them). + ValueId id_of(const std::string& name); + + Argument convert_arg(const fbs::Argument* a); + + Graph graph_; + std::unordered_map n2i_; + MetaTable tm_; +}; + +ValueId GraphBuilder::id_of(const std::string& name) { + if (name.empty()) { + return kInvalid; + } + const auto it = n2i_.find(name); + if (it != n2i_.end()) { + return it->second; + } + const ValueId id = static_cast(graph_.values.size()); + const auto mit = tm_.find(name); + if (mit != tm_.end() && mit->second != nullptr) { + graph_.values.emplace_back(name, build_tensor_meta(mit->second, name)); + } else { + graph_.values.emplace_back(name); + } + n2i_[name] = id; + return id; +} + +Argument GraphBuilder::convert_arg(const fbs::Argument* a) { + using AV = fbs::ArgumentValue; + switch (a->value_type()) { + case AV::NONE: + case AV::NoneArg: + return NoneArg{}; + case AV::TensorArg: { + TensorArg t; + t.id = id_of(str_of(a->value_as_TensorArg()->name())); + return t; + } + case AV::IntArg: { + const auto* x = a->value_as_IntArg(); + IntArg r; + r.value = x->value(); + r.id = nonempty(x->ref()) ? id_of(x->ref()->str()) : kInvalid; + return r; + } + case AV::FloatArg: { + const auto* x = a->value_as_FloatArg(); + FloatArg r; + r.value = x->value(); + r.id = nonempty(x->ref()) ? id_of(x->ref()->str()) : kInvalid; + return r; + } + case AV::BoolArg: { + const auto* x = a->value_as_BoolArg(); + BoolArg r; + r.value = x->value(); + r.id = nonempty(x->ref()) ? id_of(x->ref()->str()) : kInvalid; + return r; + } + case AV::StringArg: { + StringArg r; + r.value = str_of(a->value_as_StringArg()->value()); + return r; + } + case AV::ScalarTypeArg: { + ScalarTypeArg r; + r.value = map_scalar_type(a->value_as_ScalarTypeArg()->value()); + return r; + } + case AV::IntListArg: { + const auto* x = a->value_as_IntListArg(); + IntListArg r; + if (const auto* vals = x->values()) { + for (flatbuffers::uoffset_t i = 0; i < vals->size(); ++i) { + r.values.push_back(vals->Get(i)); + } + } + if (const auto* refs = x->refs()) { + for (flatbuffers::uoffset_t i = 0; i < refs->size(); ++i) { + r.ids.push_back( + nonempty(refs->Get(i)) ? id_of(refs->Get(i)->str()) : kInvalid); + } + } + return r; + } + case AV::FloatListArg: { + FloatListArg r; + if (const auto* vals = a->value_as_FloatListArg()->values()) { + for (flatbuffers::uoffset_t i = 0; i < vals->size(); ++i) { + r.values.push_back(vals->Get(i)); + } + } + return r; + } + case AV::BoolListArg: { + BoolListArg r; + if (const auto* vals = a->value_as_BoolListArg()->values()) { + for (flatbuffers::uoffset_t i = 0; i < vals->size(); ++i) { + r.values.push_back(vals->Get(i)); + } + } + return r; + } + case AV::TensorListArg: { + TensorListArg r; + if (const auto* nm = a->value_as_TensorListArg()->names()) { + for (flatbuffers::uoffset_t i = 0; i < nm->size(); ++i) { + r.ids.push_back(id_of(nm->Get(i)->str())); + } + } + return r; + } + case AV::OptionalTensorListArg: { + const auto* oa = a->value_as_OptionalTensorListArg(); + const auto* nm = oa->names(); + const auto* hv = oa->has_value(); + OptionalTensorListArg r; + if (nm != nullptr) { + for (flatbuffers::uoffset_t i = 0; i < nm->size(); ++i) { + const bool present = hv != nullptr && i < hv->size() && hv->Get(i); + r.ids.push_back(present ? id_of(nm->Get(i)->str()) : kInvalid); + } + } + return r; + } + case AV::GraphArg: { + const fbs::GraphArg* ga = a->value_as_GraphArg(); + GraphArg r; + r.name = str_of(ga->name()); + r.subgraph_id = static_cast(graph_.subgraphs.size()); + graph_.subgraphs.push_back(build_graph(ga->graph()).graph); + return r; + } + default: + throw std::runtime_error( + "build_graph: unsupported ArgumentValue value " + + std::to_string(static_cast(a->value_type()))); + } +} + +BuiltGraph GraphBuilder::run(const fbs::Graph* g, const MetaTable& extra_meta) { + if (g == nullptr) { + return {}; + } + + if (const auto* tvs = g->tensor_values()) { + for (flatbuffers::uoffset_t i = 0; i < tvs->size(); ++i) { + const fbs::TensorValue* tv = tvs->Get(i); + tm_[str_of(tv->name())] = tv->meta(); + } + } + for (const auto& entry : extra_meta) { + const fbs::TensorMeta*& slot = tm_[entry.first]; + if (slot == nullptr) { + slot = entry.second; + } + } + + // Pre-create meta-carrying values in a deterministic order (nicer ids). + if (const auto* tvs = g->tensor_values()) { + for (flatbuffers::uoffset_t i = 0; i < tvs->size(); ++i) { + id_of(str_of(tvs->Get(i)->name())); + } + } + + if (const auto* nodes = g->nodes()) { + for (flatbuffers::uoffset_t i = 0; i < nodes->size(); ++i) { + const fbs::Node* nd = nodes->Get(i); + Node node; + node.name = str_of(nd->name()); + node.op_kind = map_op_kind(nd->op_kind()); + node.target = str_of(nd->target()); + + if (const auto* ins = nd->inputs()) { + for (flatbuffers::uoffset_t j = 0; j < ins->size(); ++j) { + const fbs::NamedArgument* na = ins->Get(j); + NamedArgument narg; + narg.name = str_of(na->name()); + narg.mutated = na->mutated(); + narg.arg = convert_arg(na->arg()); + node.inputs.push_back(std::move(narg)); + } + } + + if (const auto* outs = nd->outputs()) { + for (flatbuffers::uoffset_t j = 0; j < outs->size(); ++j) { + const fbs::Output* o = outs->Get(j); + Output out; + out.kind = map_output_value_kind(o->kind()); + if (o->kind() == fbs::OutputValueKind::TENSOR_LIST) { + if (const auto* nm = o->names()) { + for (flatbuffers::uoffset_t k = 0; k < nm->size(); ++k) { + out.elem_ids.push_back(id_of(nm->Get(k)->str())); + } + } + } else { + out.value_id = id_of(str_of(o->name())); + if (nonempty(o->alias_of()) && valid(out.value_id)) { + const ValueId alias_id = id_of(o->alias_of()->str()); + if (alias_id == out.value_id) { + throw std::runtime_error( + "build_graph: output '" + str_of(o->name()) + + "' cannot alias itself"); + } + graph_.values.at(static_cast(out.value_id)).alias_id = + alias_id; + } + } + node.outputs.push_back(std::move(out)); + } + } + + // A placeholder with no explicit Output still produces its named value; + // synthesize one so def-use wiring records the placeholder as producer. + if (node.op_kind == OpKind::Placeholder && node.outputs.empty() && + !node.name.empty()) { + Output out; + out.value_id = id_of(node.name); + node.outputs.push_back(out); + } + + graph_.nodes.push_back(std::move(node)); + } + } + + if (const auto* gi = g->inputs()) { + for (flatbuffers::uoffset_t i = 0; i < gi->size(); ++i) { + graph_.input_ids.push_back(id_of(gi->Get(i)->str())); + } + } + if (const auto* go = g->outputs()) { + for (flatbuffers::uoffset_t i = 0; i < go->size(); ++i) { + graph_.output_ids.push_back(id_of(go->Get(i)->str())); + } + } + + graph_.initialize_schedule(); + graph_.rebuild_def_use(); + return BuiltGraph{std::move(graph_), std::move(n2i_)}; +} + +BuiltGraph build_graph(const fbs::Graph* g, const MetaTable& extra_meta) { + return GraphBuilder().run(g, extra_meta); +} + +// Resolve a method-level binding name (namespace 2) against the top-level +// graph's SSA names, or kInvalid if the graph holds no such value. +ValueId id_of_name( + const std::unordered_map& n2i, + const std::string& name) { + const auto it = n2i.find(name); + return it != n2i.end() ? it->second : kInvalid; +} + +ValueId require_binding_value( + const std::unordered_map& n2i, + const std::string& name) { + const ValueId id = id_of_name(n2i, name); + if (!valid(id)) { + throw std::runtime_error( + "build_method: data binding '" + name + + "' does not name a graph value"); + } + return id; +} + +void stamp_role(Graph& graph, ValueId id, ValueRole role) { + if (in_bounds(id, graph.values.size())) { + graph.values.at(static_cast(id)).role = role; + } +} + +} // namespace + +Method Program::build_method(size_t index) const { + if (program_fb_ == nullptr) { + throw std::runtime_error("build_method: program is not loaded"); + } + const auto* methods = program_fb_->methods(); + if (methods == nullptr || index >= methods->size()) { + throw std::runtime_error("build_method: method index out of range"); + } + const fbs::Method* m = + methods->Get(static_cast(index)); + + Method method; + method.name = str_of(m->name()); + + MetaTable constant_meta; + if (const auto* cs = m->constants()) { + for (flatbuffers::uoffset_t i = 0; i < cs->size(); ++i) { + const fbs::NamedTensorRef* c = cs->Get(i); + constant_meta[str_of(c->name())] = c->meta(); + } + } + + BuiltGraph built = build_graph(m->graph(), constant_meta); + const std::unordered_map& n2i = built.name_to_id; + method.graph = std::move(built.graph); + Graph& graph = method.graph; + + // external-constant / buffer identity (key) -> value, for BufferMutation + // output targets (a namespace-3 fqn, not an SSA name). + std::unordered_map key_to_id; + std::unordered_set bound_ids; + + if (const auto* cs = m->constants()) { + for (flatbuffers::uoffset_t i = 0; i < cs->size(); ++i) { + const fbs::NamedTensorRef* c = cs->Get(i); + DataBinding b; + const std::string name = str_of(c->name()); + b.value_id = require_binding_value(n2i, name); + if (!bound_ids.insert(b.value_id).second) { + throw std::runtime_error( + "build_method: graph value '" + name + + "' has multiple data bindings"); + } + b.role = map_input_kind(c->kind()); + b.key = str_of(c->data_key()); + b.has_data = true; + b.mutated = c->mutated(); + stamp_role(graph, b.value_id, b.role); + if (!b.key.empty()) { + key_to_id[b.key] = b.value_id; + } + method.data_bindings.push_back(std::move(b)); + } + } + + if (const auto* mbs = m->mutable_buffers()) { + for (flatbuffers::uoffset_t i = 0; i < mbs->size(); ++i) { + const fbs::MutableBufferSpec* mb = mbs->Get(i); + DataBinding b; + const std::string name = str_of(mb->name()); + b.value_id = require_binding_value(n2i, name); + if (!bound_ids.insert(b.value_id).second) { + throw std::runtime_error( + "build_method: graph value '" + name + + "' has multiple data bindings"); + } + b.role = ValueRole::Buffer; + b.key = str_of(mb->fqn()); + b.has_data = false; + b.mutated = true; + stamp_role(graph, b.value_id, ValueRole::Buffer); + if (!b.key.empty()) { + key_to_id[b.key] = b.value_id; + } + method.data_bindings.push_back(std::move(b)); + } + } + + // Top-level graph inputs not otherwise bound are user inputs. + for (const ValueId id : graph.input_ids) { + if (in_bounds(id, graph.values.size()) && + graph.values[id].role == ValueRole::Intermediate) { + graph.values[id].role = ValueRole::UserInput; + } + } + + // output_specs are parallel to graph.outputs (same order); each classifies + // graph.output_ids[i]. The mutation target resolves to a placeholder value: + // an fqn (BufferMutation) via key_to_id, else an SSA name + // (UserInputMutation). + if (const auto* os = m->output_specs()) { + if (os->size() != graph.output_ids.size()) { + throw std::runtime_error( + "build_method: output_specs count does not match graph outputs"); + } + for (flatbuffers::uoffset_t i = 0; i < os->size(); ++i) { + const fbs::OutputSpec* o = os->Get(i); + OutputSpec spec; + spec.kind = map_output_kind(o->kind()); + const std::string target = str_of(o->target()); + if (spec.kind != OutputKind::UserOutput) { + if (spec.kind == OutputKind::BufferMutation) { + const auto it = key_to_id.find(target); + spec.target_id = it != key_to_id.end() ? it->second : kInvalid; + } else { + spec.target_id = id_of_name(n2i, target); + } + if (!valid(spec.target_id)) { + throw std::runtime_error( + "build_method: mutation target '" + target + + "' does not name a bound value"); + } + } + method.output_specs.push_back(spec); + } + } + + return method; +} + +// Defined here (rather than in Program.cpp) so it sits next to build_method and +// the deserializer helpers it drives: get_method is the public lazy entry +// point, build_method the private materializer it calls on a cache miss. +const Method& Program::get_method(const std::string& name) const { + const auto it = method_cache_.find(name); + if (it != method_cache_.end()) { + return it->second; + } + if (program_fb_ != nullptr) { + if (const auto* methods = program_fb_->methods()) { + for (flatbuffers::uoffset_t i = 0; i < methods->size(); ++i) { + const flatbuffers::String* method_name = methods->Get(i)->name(); + if (method_name != nullptr && + std::string_view(method_name->c_str(), method_name->size()) == + name) { + auto res = method_cache_.emplace(name, build_method(i)); + return res.first->second; + } + } + } + } + throw std::runtime_error( + "Program::get_method: no method named '" + name + "'"); +} + +} // namespace ptn diff --git a/backends/native/runtime/Program.cpp b/backends/native/runtime/Program.cpp index 00d8de36deb..0ee87e212dc 100644 --- a/backends/native/runtime/Program.cpp +++ b/backends/native/runtime/Program.cpp @@ -8,6 +8,8 @@ #include #include +#include +#include #include #include @@ -42,6 +44,19 @@ Program Program::load(const void* data, size_t size) { } const fbs::Program* program_fb = fbs::GetProgram(bytes.data()); + // Both accessors below are schema-required, so successful verification + // guarantees that they are non-null. + std::unordered_set method_names; + for (const fbs::Method* method : *program_fb->methods()) { + const std::string name = method->name()->str(); + if (name.empty()) { + throw std::runtime_error("native program: method name is empty"); + } + if (!method_names.insert(name).second) { + throw std::runtime_error( + "native program: duplicate method name '" + name + "'"); + } + } return Program(std::move(bytes), program_fb); } @@ -50,4 +65,17 @@ size_t Program::num_methods() const { return methods == nullptr ? 0 : methods->size(); } +std::vector Program::method_names() const { + std::vector names; + const auto* methods = program_fb_->methods(); + if (methods != nullptr) { + names.reserve(methods->size()); + for (flatbuffers::uoffset_t i = 0; i < methods->size(); ++i) { + const auto* nm = methods->Get(i)->name(); + names.push_back(nm != nullptr ? nm->str() : std::string()); + } + } + return names; +} + } // namespace ptn diff --git a/backends/native/runtime/Program.h b/backends/native/runtime/Program.h index 6db8a887f97..e035e0d4d5d 100644 --- a/backends/native/runtime/Program.h +++ b/backends/native/runtime/Program.h @@ -8,8 +8,12 @@ #include #include +#include +#include #include +#include + // Forward-declaration of the generated FlatBuffer root type, included only from // .cpp files so flatbuffers stays an implementation detail of the reader. namespace native_backend { @@ -29,6 +33,10 @@ class Program { // rather than return a null root, so accessors dereference it unchecked. std::vector bytes_; const fbs::Program* program_fb_ = nullptr; + // Lazily materialized methods, keyed by name, populated on get_method(). The + // cache is mutable so lookups work on a const Program; unordered_map keeps + // returned references stable across later insertions. Not thread-safe. + mutable std::unordered_map method_cache_; Program(std::vector bytes, const fbs::Program* program_fb) : bytes_(std::move(bytes)), program_fb_(program_fb) {} @@ -41,7 +49,8 @@ class Program { Program& operator=(const Program&) = delete; // Parse and verify serialized native-graph bytes (a *.ptg buffer). Throws - // std::runtime_error on failure. + // std::runtime_error on failure. Methods are materialized lazily (see + // get_method), not here. static Program load(const void* data, size_t size); const fbs::Program* flatbuffer() const { @@ -49,6 +58,20 @@ class Program { } size_t num_methods() const; + + // Names of the program's methods, in serialized order. + std::vector method_names() const; + + // Materialize (or return the cached) method by name. Builds the in-memory IR + // on first request and caches it; later calls return the same instance. + // Throws std::runtime_error if no method has that name. Impl in + // Deserialize.cpp. + const Method& get_method(const std::string& name) const; + + private: + // Deserialize the fb method at `index` into the in-memory IR (Graph + + // bindings). Impl in Deserialize.cpp. + Method build_method(size_t index) const; }; } // namespace ptn diff --git a/backends/native/runtime/deserialize/BUCK b/backends/native/runtime/deserialize/BUCK new file mode 100644 index 00000000000..0ab35888218 --- /dev/null +++ b/backends/native/runtime/deserialize/BUCK @@ -0,0 +1,11 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +# Any targets that should be shared between fbcode and xplat must be defined in +# targets.bzl. This file can contain cell-only targets. + +non_fbcode_target(_kind = define_common_targets) + +fbcode_target(_kind = define_common_targets) diff --git a/backends/native/runtime/deserialize/ByteSpan.h b/backends/native/runtime/deserialize/ByteSpan.h new file mode 100644 index 00000000000..0b16f75a746 --- /dev/null +++ b/backends/native/runtime/deserialize/ByteSpan.h @@ -0,0 +1,18 @@ +// 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 + +#include +#include + +namespace ptn { + +// Borrowed byte-range views. The producer defines their lifetime. +using ByteSpan = std::span; +using MutableByteSpan = std::span; + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/Json.h b/backends/native/runtime/deserialize/Json.h new file mode 100644 index 00000000000..8cf1f1cd305 --- /dev/null +++ b/backends/native/runtime/deserialize/Json.h @@ -0,0 +1,15 @@ +// 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 + +#include + +namespace ptn { + +using Json = nlohmann::ordered_json; + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/OwnedBytes.cpp b/backends/native/runtime/deserialize/OwnedBytes.cpp new file mode 100644 index 00000000000..7bb919cd172 --- /dev/null +++ b/backends/native/runtime/deserialize/OwnedBytes.cpp @@ -0,0 +1,159 @@ +// 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. + +#include + +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#include +#include +#endif + +namespace ptn { +namespace { + +std::string errno_suffix(int error) { + return ": " + std::error_code(error, std::system_category()).message(); +} + +} // namespace + +void OwnedBytes::Unmap::operator()(void* base) const noexcept { +#if !defined(_WIN32) + // Nothing useful to do if this fails, and it must not throw: unique_ptr calls + // this from its destructor. + ::munmap(base, size); +#endif +} + +OwnedBytes OwnedBytes::from_vector(std::vector bytes) { + return OwnedBytes(std::move(bytes)); +} + +OwnedBytes OwnedBytes::from_file(const std::string& path, bool use_mmap) { + return use_mmap ? map_file(path) : read_file(path); +} + +OwnedBytes OwnedBytes::read_file(const std::string& path) { + std::error_code error; + const std::filesystem::file_status status = + std::filesystem::status(path, error); + if (error) { + throw std::runtime_error("cannot inspect " + path + ": " + error.message()); + } + if (!std::filesystem::is_regular_file(status)) { + throw std::runtime_error("cannot read " + path + ": not a regular file"); + } + const uintmax_t file_size = std::filesystem::file_size(path, error); + if (error) { + throw std::runtime_error("cannot size " + path + ": " + error.message()); + } + if (file_size > std::numeric_limits::max() || + file_size > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error("cannot read " + path + ": file is too large"); + } + + std::ifstream file(path, std::ios::binary); + if (!file) { + throw std::runtime_error("cannot open " + path); + } + const size_t size = static_cast(file_size); + if (size == 0) { + return OwnedBytes(std::vector()); + } + HeapBuffer buffer{std::make_unique_for_overwrite(size), size}; + if (!file.read( + reinterpret_cast(buffer.data.get()), + static_cast(size))) { + throw std::runtime_error("cannot read " + path); + } + return OwnedBytes(std::move(buffer)); +} + +OwnedBytes OwnedBytes::map_file(const std::string& path) { +#if defined(_WIN32) + // TODO: Implement Windows mappings with CreateFileMapping and MapViewOfFile. + throw std::runtime_error("cannot mmap " + path + ": unsupported platform"); +#else + const int fd = ::open(path.c_str(), O_RDONLY); + if (fd < 0) { + const int error = errno; + throw std::runtime_error("cannot open " + path + errno_suffix(error)); + } + + struct stat st = {}; + if (::fstat(fd, &st) < 0) { + const std::string suffix = errno_suffix(errno); + ::close(fd); + throw std::runtime_error("cannot size " + path + suffix); + } + if (!S_ISREG(st.st_mode)) { + ::close(fd); + throw std::runtime_error("cannot mmap " + path + ": not a regular file"); + } + if (st.st_size < 0) { + ::close(fd); + throw std::runtime_error("cannot mmap " + path + ": invalid file size"); + } + if (static_cast(st.st_size) > std::numeric_limits::max()) { + ::close(fd); + throw std::runtime_error("cannot mmap " + path + ": file is too large"); + } + const size_t size = static_cast(st.st_size); + if (size == 0) { + ::close(fd); + return OwnedBytes(std::vector()); + } + + // The whole file from offset 0, so the base is page-aligned and every span + // into it has the same alignment it would have in a heap buffer. MAP_SHARED + // lets other processes mapping this file share the same physical pages; the + // mapping is read-only either way. + void* base = ::mmap(nullptr, size, PROT_READ, MAP_SHARED, fd, 0); + if (base == MAP_FAILED) { + const int error = errno; + ::close(fd); + throw std::runtime_error("cannot mmap " + path + errno_suffix(error)); + } + // The mapping keeps the file alive on its own, so the descriptor is dead + // weight past this point. + ::close(fd); + return OwnedBytes(MappedFile(base, Unmap{size})); +#endif +} + +ByteSpan OwnedBytes::span() const { + if (const HeapBuffer* buffer = std::get_if(&storage_)) { + if (buffer->data == nullptr) { + return {}; + } + return ByteSpan(buffer->data.get(), buffer->size); + } + if (const MappedFile* mapped = std::get_if(&storage_)) { + if (mapped->get() == nullptr) { + return {}; + } + return ByteSpan( + static_cast(mapped->get()), mapped->get_deleter().size); + } + const std::vector& bytes = std::get>(storage_); + return ByteSpan(bytes.data(), bytes.size()); +} + +bool OwnedBytes::is_mapped() const { + return std::holds_alternative(storage_); +} + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/OwnedBytes.h b/backends/native/runtime/deserialize/OwnedBytes.h new file mode 100644 index 00000000000..fc376bfd68e --- /dev/null +++ b/backends/native/runtime/deserialize/OwnedBytes.h @@ -0,0 +1,96 @@ +// 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 + +#include +#include +#include +#include +#include +#include + +#include + +namespace ptn { + +// Owning, read-only bytes: either a heap buffer or a read-only file mapping. +// +// Hands out spans that alias the storage. Each alternative keeps its payload +// address across a move, so spans taken before a move stay valid for as long as +// the OwnedBytes lives. Copy is deleted: this holds a whole model. +// +// The mapped alternative is the one that matters for large packages: nothing is +// copied, the pages are demand-paged, and they are shared with any other +// process mapping the same file. +class OwnedBytes { + private: + // Releases a mapping. Carries the length because that is what munmap needs, + // which lets a unique_ptr supply the whole move-only lifetime — no + // hand-written destructor or move operations. + struct Unmap { + size_t size = 0; + + void operator()(void* base) const noexcept; + }; + + struct HeapBuffer { + std::unique_ptr data; + size_t size = 0; + }; + + // A read-only mapping of an entire file. + using MappedFile = std::unique_ptr; + + std::variant, HeapBuffer, MappedFile> storage_; + + explicit OwnedBytes(std::vector bytes) + : storage_(std::move(bytes)) {} + explicit OwnedBytes(HeapBuffer buffer) : storage_(std::move(buffer)) {} + explicit OwnedBytes(MappedFile mapped_file) + : storage_(std::move(mapped_file)) {} + + public: + // Empty, owning nothing. + OwnedBytes() = default; + + ~OwnedBytes() = default; + OwnedBytes(OwnedBytes&&) noexcept = default; + OwnedBytes& operator=(OwnedBytes&&) noexcept = default; + OwnedBytes(const OwnedBytes&) = delete; + OwnedBytes& operator=(const OwnedBytes&) = delete; + + // The whole payload. Valid for this OwnedBytes' lifetime. + ByteSpan span() const; + + // True when these bytes are a file mapping rather than a heap buffer. + bool is_mapped() const; + + // Take ownership of a buffer the caller already has, without copying it. + static OwnedBytes from_vector(std::vector bytes); + + // Acquire the contents of `path`. Maps it read-only by default: nothing is + // copied, pages arrive on demand, and they are shared with any other process + // mapping the same file. Pass use_mmap=false to read it into the heap + // instead, which is worth it only when the file must outlive edits to it on + // disk. A mapping sees concurrent edits, and accessing pages past a + // concurrent truncation can terminate the process with SIGBUS. + // + // `path` must report its complete size through the filesystem. Virtual files + // such as procfs entries that report zero bytes but produce data are outside + // this API's scope. + // + // Throws std::runtime_error if the file cannot be read, or cannot be mapped + // when mapping was asked for (including on a platform with no mmap). An empty + // file yields empty heap bytes either way, since mmap rejects a zero length. + static OwnedBytes from_file(const std::string& path, bool use_mmap = true); + + private: + static OwnedBytes read_file(const std::string& path); + static OwnedBytes map_file(const std::string& path); +}; + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/Package.cpp b/backends/native/runtime/deserialize/Package.cpp new file mode 100644 index 00000000000..2b151c35b08 --- /dev/null +++ b/backends/native/runtime/deserialize/Package.cpp @@ -0,0 +1,233 @@ +// 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. + +#include + +#include +#include +#include +#include +#include + +#include + +namespace ptn { +namespace { + +// Reserved by safetensors, so it can never name a constant. +constexpr std::string_view kMetadataKey = "__metadata__"; +std::atomic next_package_id{1}; + +std::unordered_map parse_aliases( + ByteSpan member, + const SafeTensorsReader& tensors) { + Json doc; + try { + doc = Json::parse(std::string_view( + reinterpret_cast(member.data()), member.size())); + } catch (const Json::exception& error) { + throw std::runtime_error( + "package: invalid aliases.json: " + std::string(error.what())); + } + if (!doc.is_object()) { + throw std::runtime_error("package: aliases.json is not a JSON object"); + } + + std::unordered_map aliases; + for (auto entry = doc.begin(); entry != doc.end(); ++entry) { + const std::string& key = entry.key(); + if (!entry.value().is_string()) { + throw std::runtime_error( + "package: alias '" + key + "' does not name a string owner"); + } + const std::string& owner = entry.value().get_ref(); + if (key == kMetadataKey) { + throw std::runtime_error( + "package: alias key is reserved by safetensors: " + key); + } + // An owner is always a real safetensors entry and an alias is never one, so + // resolution stays a single lookup. Enforce both rather than trusting it. + if (tensors.find(owner) == nullptr) { + std::string message = "package: alias '"; + message += key; + message += "' names owner '"; + message += owner; + message += "', which has no safetensors entry"; + throw std::runtime_error(message); + } + if (tensors.find(key) != nullptr) { + throw std::runtime_error( + "package: '" + key + "' is both a safetensors owner and an alias"); + } + if (!aliases.emplace(key, owner).second) { + throw std::runtime_error("package: duplicate alias key: " + key); + } + } + return aliases; +} + +} // namespace + +bool Package::looks_like_package(ByteSpan bytes) { + // Every zip record signature begins "PK"; a bare .ptg starts with a + // flatbuffer root offset followed by "NPTG" at offset 4, so this cannot + // collide. + return bytes.size() >= 2 && bytes[0] == 'P' && bytes[1] == 'K'; +} + +Package::Package() + : id_(next_package_id.fetch_add(1, std::memory_order_relaxed)) {} + +Package Package::load(OwnedBytes bytes) { + Package out; + out.archive_bytes_ = std::move(bytes); + out.zip_ = ZipReader::open(out.archive_bytes_.span()); + out.load_metadata(); + return out; +} + +Package Package::load(const std::string& path) { + Package out; + out.zip_ = ZipReader::open(path); + out.load_metadata(); + return out; +} + +Package& Package::operator=(Package&& other) noexcept { + if (this != &other) { + zip_.reset(); + id_ = other.id_; + archive_bytes_ = std::move(other.archive_bytes_); + zip_ = std::move(other.zip_); + program_ = std::move(other.program_); + tensors_ = std::move(other.tensors_); + tensor_data_offset_ = other.tensor_data_offset_; + aliases_ = std::move(other.aliases_); + } + return *this; +} + +void Package::load_metadata() { + if (!zip_->member_size(kProgramEntry)) { + throw std::runtime_error( + std::string("package: missing required member ") + kProgramEntry); + } + program_ = zip_->read(kProgramEntry); + + // Absent whenever the program references no constants, which is normal for a + // graph over user inputs alone. + const std::optional tensor_size = + zip_->member_size(kSafeTensorsEntry); + if (tensor_size) { + std::array prefix{}; + if (*tensor_size < prefix.size()) { + throw std::runtime_error( + "package: safetensors member is shorter than its length prefix"); + } + zip_->read_into(kSafeTensorsEntry, 0, MutableByteSpan(prefix)); + const size_t header_size = SafeTensorsReader::header_size(prefix); + if (header_size > *tensor_size - prefix.size()) { + throw std::runtime_error( + "package: safetensors header exceeds its zip member"); + } + std::vector header(header_size); + zip_->read_into(kSafeTensorsEntry, prefix.size(), MutableByteSpan(header)); + tensor_data_offset_ = prefix.size() + header_size; + tensors_ = SafeTensorsReader::open_header( + ByteSpan(header), *tensor_size - tensor_data_offset_); + } + + if (zip_->member_size(kAliasesEntry)) { + if (!tensors_) { + throw std::runtime_error( + std::string("package: has ") + kAliasesEntry + " but no " + + kSafeTensorsEntry); + } + const std::vector aliases = zip_->read(kAliasesEntry); + aliases_ = parse_aliases(ByteSpan(aliases), *tensors_); + } +} + +std::optional Package::constant_info( + const std::string& key) const { + if (!tensors_) { + return std::nullopt; + } + const auto alias = aliases_.find(key); + const std::string& owner = alias == aliases_.end() ? key : alias->second; + + const TensorEntry* entry = tensors_->find(owner); + if (entry == nullptr) { + return std::nullopt; + } + + ConstantInfo out; + out.package_id = id_; + out.dtype = entry->dtype; + out.sizes = &entry->sizes; + out.nbytes = entry->nbytes; + out.owner = owner; + return out; +} + +std::optional Package::acquire_constant( + const std::string& key) const { + const std::optional info = constant_info(key); + if (!info) { + return std::nullopt; + } + std::vector bytes(info->nbytes); + load_constant_into(key, MutableByteSpan(bytes)); + return OwnedBytes::from_vector(std::move(bytes)); +} + +bool Package::load_constant_into( + const std::string& key, + MutableByteSpan destination) const { + const std::optional info = constant_info(key); + if (!info) { + return false; + } + if (destination.size() != info->nbytes) { + throw std::runtime_error( + "package: destination for '" + key + "' has " + + std::to_string(destination.size()) + " bytes; expected " + + std::to_string(info->nbytes)); + } + const TensorEntry* entry = tensors_->find(info->owner); + zip_->read_into( + kSafeTensorsEntry, tensor_data_offset_ + entry->offset, destination); + return true; +} + +void Package::verify_constants() const { + if (tensors_) { + zip_->verify(kSafeTensorsEntry); + } +} + +std::vector Package::keys() const { + std::vector out; + if (tensors_) { + out = tensors_->names(); + } + for (const auto& alias : aliases_) { + out.push_back(alias.first); + } + std::ranges::sort(out); + return out; +} + +const std::vector& Package::owner_keys() const { + static const std::vector kNone; + return tensors_ ? tensors_->names() : kNone; +} + +size_t Package::constant_bytes() const { + return tensors_ ? tensors_->total_bytes() : 0; +} + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/Package.h b/backends/native/runtime/deserialize/Package.h new file mode 100644 index 00000000000..77788e9e8d6 --- /dev/null +++ b/backends/native/runtime/deserialize/Package.h @@ -0,0 +1,125 @@ +// 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 + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ptn { + +// Fixed member names inside a .ptn. The package survives being renamed because +// nothing depends on the file name. +constexpr const char* kProgramEntry = "program.ptg"; +constexpr const char* kSafeTensorsEntry = "program.safetensors"; +constexpr const char* kAliasesEntry = "aliases.json"; + +// Metadata for one constant resolved out of a package. `sizes` is borrowed from +// the Package and must not outlive it. +struct ConstantInfo { + uint64_t package_id = 0; + ScalarType dtype = kFloat; + const std::vector* sizes = nullptr; + size_t nbytes = 0; + // Key that actually owns these bytes. Differs from the requested key when the + // package deduplicated two byte-identical immutable constants. + std::string owner; +}; + +// A loaded .ptn package: the serialized native Program plus the constants it +// references. +// +// Opening a file-backed package reads its directory and metadata, but leaves +// weight payloads on disk until an engine requests them. +class Package { + private: + uint64_t id_ = 0; + OwnedBytes archive_bytes_; + std::optional zip_; + std::vector program_; + // Absent when the program references no constants, in which case the package + // has no safetensors member at all. + std::optional tensors_; + size_t tensor_data_offset_ = 0; + std::unordered_map aliases_; + + Package(); + + public: + ~Package() = default; + Package(Package&&) noexcept = default; + Package& operator=(Package&& other) noexcept; + Package(const Package&) = delete; + Package& operator=(const Package&) = delete; + + // Open and parse the .ptn at `path` without loading its weight payloads. + static Package load(const std::string& path); + + // Parse a .ptn image already in hand. Takes ownership rather than copying, so + // a hundred-megabyte package is resident once. For callers that must inspect + // the bytes before deciding this is a package at all; everyone else should + // use the path overload. Throws std::runtime_error if the zip, the + // safetensors index, or the alias map is malformed, or if the required + // program member is missing. + static Package load(OwnedBytes bytes); + + // The serialized native Program flatbuffer (the program.ptg member). + ByteSpan program_bytes() const { + return ByteSpan(program_); + } + + // Zip member names present, in central-directory order. Diagnostic only. + const std::vector& member_names() const { + return zip_->names(); + } + + // Keys that own their bytes, in safetensors header order. + const std::vector& owner_keys() const; + + // Duplicate key -> owner key. + const std::unordered_map& aliases() const { + return aliases_; + } + + // Metadata for `key`, resolving an alias to its owner. nullopt when absent. + std::optional constant_info(const std::string& key) const; + + // Load one constant into a new owning buffer. nullopt when absent. + std::optional acquire_constant(const std::string& key) const; + + // Load one constant directly into an exact-sized destination. Returns false + // when absent and throws when the destination has the wrong size. + bool load_constant_into(const std::string& key, MutableByteSpan destination) + const; + + // Stream all weight bytes once to verify the zip member checksum. + void verify_constants() const; + + // Every key the package resolves, owners and aliases alike, sorted. + std::vector keys() const; + + // Total bytes across owner entries, i.e. what the constants actually cost. + size_t constant_bytes() const; + + // True if `bytes` starts with the zip local-header signature, i.e. looks like + // a package rather than a bare .ptg flatbuffer. Lets a tool accept either. + static bool looks_like_package(ByteSpan bytes); + + private: + void load_metadata(); +}; + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/SafeTensorsReader.cpp b/backends/native/runtime/deserialize/SafeTensorsReader.cpp new file mode 100644 index 00000000000..35dc7418b1c --- /dev/null +++ b/backends/native/runtime/deserialize/SafeTensorsReader.cpp @@ -0,0 +1,250 @@ +// 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. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace ptn { +namespace { + +// Reserved header member holding free-form string metadata, not a tensor. +constexpr std::string_view kMetadataKey = "__metadata__"; +struct DtypeCode { + std::string_view code; + ScalarType dtype; +}; + +// safetensors dtype codes, as written by safetensors.torch. Codes with no +// ScalarType counterpart (complex, 4-bit and 8-bit float variants) are absent +// and rejected by name, so an unsupported constant fails at load rather than +// being misread as another width. +constexpr std::array kDtypeCodes{{ + {"F64", kDouble}, + {"F32", kFloat}, + {"F16", kHalf}, + {"BF16", kBFloat16}, + {"I64", kLong}, + {"I32", kInt}, + {"I16", kShort}, + {"I8", kChar}, + {"U8", kByte}, + {"BOOL", kBool}, + {"U16", kUInt16}, + {"U32", kUInt32}, + {"U64", kUInt64}, +}}; + +ScalarType scalar_type_of(std::string_view code) { + const auto it = std::ranges::find(kDtypeCodes, code, &DtypeCode::code); + if (it == kDtypeCodes.end()) { + throw std::runtime_error( + "safetensors: unsupported dtype code: " + std::string(code)); + } + return it->dtype; +} + +const Json& required_member( + const Json& entry, + std::string_view key, + const std::string& name) { + const auto value = entry.find(key); + if (value == entry.end()) { + std::string message = "safetensors: entry '"; + message += name; + message += "' has no '"; + message += key; + message += "'"; + throw std::runtime_error(message); + } + return value.value(); +} + +std::vector read_sizes(const Json& shape, const std::string& name) { + if (!shape.is_array()) { + throw std::runtime_error( + "safetensors: entry '" + name + "' shape is not an array"); + } + std::vector sizes; + for (const Json& dim : shape) { + if (!dim.is_number_unsigned()) { + throw std::runtime_error( + "safetensors: entry '" + name + + "' has a non-negative integer dimension"); + } + const uint64_t value = dim.get(); + if (value > static_cast(INT64_MAX)) { + throw std::runtime_error( + "safetensors: entry '" + name + "' has an out-of-range dimension"); + } + sizes.push_back(static_cast(value)); + } + return sizes; +} + +// Element count of `sizes`, rejecting an overflowing product. A rank-0 shape is +// a scalar, whose element count is 1. +size_t numel_of(const std::vector& sizes, const std::string& name) { + size_t numel = 1; + for (const int64_t dim : sizes) { + if (dim < 0) { + throw std::runtime_error( + "safetensors: entry '" + name + "' has a negative dimension"); + } + const size_t d = static_cast(dim); + if (d != 0 && numel > SIZE_MAX / d) { + throw std::runtime_error( + "safetensors: entry '" + name + "' element count overflows"); + } + numel *= d; + } + return numel; +} + +} // namespace + +size_t SafeTensorsReader::header_size(ByteSpan prefix) { + static_assert( + std::endian::native == std::endian::little, + "the length prefix is little-endian; a big-endian host needs a swap"); + if (prefix.size() < kLengthPrefixSize) { + throw std::runtime_error( + "safetensors: blob is shorter than its length prefix"); + } + uint64_t size = 0; + std::memcpy(&size, prefix.data(), kLengthPrefixSize); + if (size > std::numeric_limits::max()) { + throw std::runtime_error("safetensors: header is too large"); + } + return static_cast(size); +} + +SafeTensorsReader SafeTensorsReader::open(ByteSpan blob) { + const size_t header_len = header_size(blob); + if (header_len > blob.size() - kLengthPrefixSize) { + throw std::runtime_error("safetensors: header length exceeds the blob"); + } + const ByteSpan header = + blob.subspan(kLengthPrefixSize, static_cast(header_len)); + return open_header( + header, + blob.size() - kLengthPrefixSize - static_cast(header_len)); +} + +SafeTensorsReader SafeTensorsReader::open_header( + ByteSpan header_bytes, + size_t data_size) { + const std::string_view header_text( + reinterpret_cast(header_bytes.data()), header_bytes.size()); + Json header; + try { + header = Json::parse(header_text); + } catch (const Json::exception& error) { + throw std::runtime_error( + "safetensors: invalid JSON header: " + std::string(error.what())); + } + if (!header.is_object()) { + throw std::runtime_error("safetensors: header is not a JSON object"); + } + + SafeTensorsReader out; + + for (auto member = header.begin(); member != header.end(); ++member) { + const std::string& name = member.key(); + if (name == kMetadataKey) { + if (!member.value().is_object() || !member.value().empty()) { + throw std::runtime_error( + "safetensors: __metadata__ must be an empty object"); + } + continue; + } + const Json& entry = member.value(); + if (!entry.is_object()) { + throw std::runtime_error( + "safetensors: entry '" + name + "' is not an object"); + } + + TensorEntry parsed; + const Json& dtype = required_member(entry, "dtype", name); + if (!dtype.is_string()) { + throw std::runtime_error( + "safetensors: entry '" + name + "' dtype is not a string"); + } + parsed.dtype = scalar_type_of(dtype.get_ref()); + parsed.sizes = read_sizes(required_member(entry, "shape", name), name); + + const Json& range = required_member(entry, "data_offsets", name); + if (!range.is_array() || range.size() != 2) { + throw std::runtime_error( + "safetensors: entry '" + name + "' data_offsets is not a pair"); + } + if (!range[0].is_number_unsigned() || !range[1].is_number_unsigned()) { + throw std::runtime_error( + "safetensors: entry '" + name + + "' data_offsets contains a non-negative integer"); + } + const uint64_t begin = range[0].get(); + const uint64_t end = range[1].get(); + if (begin > end || end > data_size || + end > std::numeric_limits::max()) { + throw std::runtime_error( + "safetensors: entry '" + name + + "' byte range is outside the data section"); + } + parsed.offset = static_cast(begin); + parsed.nbytes = static_cast(end - begin); + + // The payload must be exactly as large as its dtype and shape imply. + // Without this, a short entry becomes an out-of-bounds read in whatever + // consumes it, sized from the metadata rather than the bytes. + const size_t numel = numel_of(parsed.sizes, name); + const size_t element_bytes = element_size(parsed.dtype); + if (element_bytes != 0 && numel > SIZE_MAX / element_bytes) { + throw std::runtime_error( + "safetensors: entry '" + name + "' byte size overflows"); + } + const size_t expected = numel * element_bytes; + if (parsed.nbytes != expected) { + throw std::runtime_error( + "safetensors: entry '" + name + "' holds " + + std::to_string(parsed.nbytes) + + " bytes but its dtype and shape need " + std::to_string(expected)); + } + + if (!out.entries_.emplace(name, std::move(parsed)).second) { + throw std::runtime_error("safetensors: duplicate entry: " + name); + } + out.names_.push_back(name); + } + + return out; +} + +const TensorEntry* SafeTensorsReader::find(const std::string& name) const { + const auto it = entries_.find(name); + return it == entries_.end() ? nullptr : &it->second; +} + +size_t SafeTensorsReader::total_bytes() const { + return std::accumulate( + entries_.begin(), + entries_.end(), + size_t{0}, + [](size_t total, const auto& entry) { + return total + entry.second.nbytes; + }); +} + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/SafeTensorsReader.h b/backends/native/runtime/deserialize/SafeTensorsReader.h new file mode 100644 index 00000000000..d54026ef805 --- /dev/null +++ b/backends/native/runtime/deserialize/SafeTensorsReader.h @@ -0,0 +1,71 @@ +// 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 + +#include +#include +#include +#include +#include + +#include +#include + +namespace ptn { + +// One tensor's entry in a safetensors index. +struct TensorEntry { + ScalarType dtype = kFloat; + std::vector sizes; + // Byte range within the blob's data section, not the whole blob. + size_t offset = 0; + size_t nbytes = 0; +}; + +// Reader for the safetensors format: +// +// [u64 header_len][JSON header][data section] +// +// The header maps a tensor name to its dtype, shape, and byte range within the +// data section. The reserved "__metadata__" member must be empty until the +// reader exposes metadata semantics. +// +// Tensor payloads are packed with no per-tensor padding, so an entry's absolute +// alignment within the file is arbitrary. Consumers must copy bytes into any +// destination that requires stronger alignment. +class SafeTensorsReader { + private: + std::unordered_map entries_; + std::vector names_; + + public: + static constexpr size_t kLengthPrefixSize = 8; + + // Parse `blob`'s index. Throws std::runtime_error if the blob is truncated, + // the header is not a JSON object, a dtype has no ScalarType, or a byte range + // is inconsistent with its dtype and shape. + static SafeTensorsReader open(ByteSpan blob); + + // Decode the format's little-endian length prefix. + static size_t header_size(ByteSpan prefix); + + // Parse a header without loading its data section. + static SafeTensorsReader open_header(ByteSpan header, size_t data_size); + + // Entry for `name`, or nullptr when absent. + const TensorEntry* find(const std::string& name) const; + + // Tensor names, in header order, excluding "__metadata__". + const std::vector& names() const { + return names_; + } + + // Total payload bytes across all entries. + size_t total_bytes() const; +}; + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/ZipReader.cpp b/backends/native/runtime/deserialize/ZipReader.cpp new file mode 100644 index 00000000000..2efb11484cf --- /dev/null +++ b/backends/native/runtime/deserialize/ZipReader.cpp @@ -0,0 +1,253 @@ +// 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. + +#include + +#include +#include +#include +#include +#include +#include + +#include + +namespace ptn { +namespace { + +struct ZipDeleter { + void operator()(zip_t* archive) const noexcept { + zip_discard(archive); + } +}; + +struct ZipFileDeleter { + void operator()(zip_file_t* file) const noexcept { + zip_fclose(file); + } +}; + +using ZipHandle = std::unique_ptr; +using ZipFileHandle = std::unique_ptr; + +[[noreturn]] void throw_zip(zip_t* archive, const std::string& operation) { + throw std::runtime_error("zip: " + operation + ": " + zip_strerror(archive)); +} + +[[noreturn]] void throw_zip_file( + zip_file_t* file, + const std::string& operation) { + throw std::runtime_error( + "zip: " + operation + ": " + zip_file_strerror(file)); +} + +ZipHandle open_path(const std::string& path) { + int error_code = 0; + ZipHandle archive( + zip_open(path.c_str(), ZIP_RDONLY | ZIP_CHECKCONS, &error_code)); + if (archive == nullptr) { + zip_error_t error; + zip_error_init_with_code(&error, error_code); + const std::string message = + "zip: cannot open " + path + ": " + zip_error_strerror(&error); + zip_error_fini(&error); + throw std::runtime_error(message); + } + return archive; +} + +ZipHandle open_memory(ByteSpan bytes) { + zip_error_t error; + zip_error_init(&error); + zip_source_t* source = + zip_source_buffer_create(bytes.data(), bytes.size(), 0, &error); + if (source == nullptr) { + const std::string message = "zip: cannot create memory source: " + + std::string(zip_error_strerror(&error)); + zip_error_fini(&error); + throw std::runtime_error(message); + } + + ZipHandle archive( + zip_open_from_source(source, ZIP_RDONLY | ZIP_CHECKCONS, &error)); + if (archive == nullptr) { + zip_source_free(source); + const std::string message = "zip: cannot open memory source: " + + std::string(zip_error_strerror(&error)); + zip_error_fini(&error); + throw std::runtime_error(message); + } + zip_error_fini(&error); + return archive; +} + +} // namespace + +struct ZipReader::Impl { + explicit Impl(ZipHandle handle) : archive(std::move(handle)) {} + + ZipHandle archive; +}; + +ZipReader::ZipReader(std::unique_ptr impl) : impl_(std::move(impl)) { + const zip_int64_t count = zip_get_num_entries(impl_->archive.get(), 0); + if (count < 0) { + throw_zip(impl_->archive.get(), "cannot enumerate members"); + } + + names_.reserve(static_cast(count)); + for (zip_uint64_t index = 0; index < static_cast(count); + ++index) { + zip_stat_t stat; + zip_stat_init(&stat); + if (zip_stat_index(impl_->archive.get(), index, ZIP_FL_UNCHANGED, &stat) != + 0) { + throw_zip(impl_->archive.get(), "cannot stat member"); + } + constexpr zip_uint64_t kRequired = ZIP_STAT_NAME | ZIP_STAT_SIZE | + ZIP_STAT_COMP_METHOD | ZIP_STAT_ENCRYPTION_METHOD; + if ((stat.valid & kRequired) != kRequired || stat.name == nullptr) { + throw std::runtime_error("zip: member metadata is incomplete"); + } + const std::string name(stat.name); + if (stat.comp_method != ZIP_CM_STORE) { + throw std::runtime_error("zip: member is compressed: " + name); + } + if (stat.encryption_method != ZIP_EM_NONE) { + throw std::runtime_error("zip: member is encrypted: " + name); + } + if (stat.size > std::numeric_limits::max()) { + throw std::runtime_error("zip: member is too large: " + name); + } + if (!entries_.emplace(name, Entry{index, static_cast(stat.size)}) + .second) { + throw std::runtime_error("zip: duplicate member name: " + name); + } + names_.push_back(name); + } +} + +ZipReader::~ZipReader() = default; +ZipReader::ZipReader(ZipReader&&) noexcept = default; +ZipReader& ZipReader::operator=(ZipReader&&) noexcept = default; + +ZipReader ZipReader::open(const std::string& path) { + return ZipReader(std::make_unique(open_path(path))); +} + +ZipReader ZipReader::open(ByteSpan archive) { + return ZipReader(std::make_unique(open_memory(archive))); +} + +std::optional ZipReader::member_size(std::string_view name) const { + const Entry* entry = find_entry(name); + return entry == nullptr ? std::nullopt : std::optional(entry->size); +} + +std::vector ZipReader::read(std::string_view name) const { + const Entry* entry = find_entry(name); + if (entry == nullptr) { + throw std::runtime_error("zip: no member named " + std::string(name)); + } + std::vector bytes(entry->size); + read_entry_into(name, *entry, 0, MutableByteSpan(bytes)); + return bytes; +} + +void ZipReader::read_into( + std::string_view name, + size_t offset, + MutableByteSpan destination) const { + const Entry* entry = find_entry(name); + if (entry == nullptr) { + throw std::runtime_error("zip: no member named " + std::string(name)); + } + read_entry_into(name, *entry, offset, destination); +} + +const ZipReader::Entry* ZipReader::find_entry(std::string_view name) const { + const auto entry = entries_.find(name); + return entry == entries_.end() ? nullptr : &entry->second; +} + +void ZipReader::read_entry_into( + std::string_view name, + const Entry& entry, + size_t offset, + MutableByteSpan destination) const { + if (offset > entry.size || destination.size() > entry.size - offset) { + throw std::runtime_error( + "zip: read range is outside member " + std::string(name)); + } + if (destination.empty()) { + return; + } + + ZipFileHandle file( + zip_fopen_index(impl_->archive.get(), entry.index, ZIP_FL_UNCHANGED)); + if (file == nullptr) { + throw_zip(impl_->archive.get(), "cannot open member " + std::string(name)); + } + if (zip_fseek(file.get(), static_cast(offset), SEEK_SET) != 0) { + throw_zip_file(file.get(), "cannot seek member " + std::string(name)); + } + + size_t written = 0; + while (written < destination.size()) { + const zip_uint64_t request = + static_cast(destination.size() - written); + const zip_int64_t count = + zip_fread(file.get(), destination.data() + written, request); + if (count < 0) { + throw_zip_file(file.get(), "cannot read member " + std::string(name)); + } + if (count == 0) { + throw std::runtime_error( + "zip: unexpected end of member " + std::string(name)); + } + written += static_cast(count); + } +} + +void ZipReader::verify(std::string_view name) const { + const Entry* entry = find_entry(name); + if (entry == nullptr) { + throw std::runtime_error("zip: no member named " + std::string(name)); + } + + ZipFileHandle file( + zip_fopen_index(impl_->archive.get(), entry->index, ZIP_FL_UNCHANGED)); + if (file == nullptr) { + throw_zip(impl_->archive.get(), "cannot open member " + std::string(name)); + } + + std::array buffer{}; + size_t read = 0; + while (read < entry->size) { + const size_t request = std::min(buffer.size(), entry->size - read); + const zip_int64_t count = zip_fread(file.get(), buffer.data(), request); + if (count < 0) { + throw_zip_file(file.get(), "cannot verify member " + std::string(name)); + } + if (count == 0) { + throw std::runtime_error( + "zip: unexpected end of member " + std::string(name)); + } + read += static_cast(count); + } + + uint8_t extra = 0; + const zip_int64_t count = zip_fread(file.get(), &extra, 1); + if (count < 0) { + throw_zip_file(file.get(), "cannot verify member " + std::string(name)); + } + if (count != 0) { + throw std::runtime_error( + "zip: member is larger than its metadata: " + std::string(name)); + } +} + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/ZipReader.h b/backends/native/runtime/deserialize/ZipReader.h new file mode 100644 index 00000000000..0abc2c012b7 --- /dev/null +++ b/backends/native/runtime/deserialize/ZipReader.h @@ -0,0 +1,87 @@ +// 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 + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace ptn { + +// Read-only access to stored members of a zip archive. +class ZipReader { + private: + struct Entry { + uint64_t index = 0; + size_t size = 0; + }; + + struct StringHash { + using is_transparent = void; + + size_t operator()(std::string_view value) const noexcept { + return std::hash{}(value); + } + }; + + struct Impl; + + std::unique_ptr impl_; + std::unordered_map> entries_; + std::vector names_; + + explicit ZipReader(std::unique_ptr impl); + const Entry* find_entry(std::string_view name) const; + void read_entry_into( + std::string_view name, + const Entry& entry, + size_t offset, + MutableByteSpan destination) const; + + public: + ~ZipReader(); + ZipReader(ZipReader&&) noexcept; + ZipReader& operator=(ZipReader&&) noexcept; + ZipReader(const ZipReader&) = delete; + ZipReader& operator=(const ZipReader&) = delete; + + // Opens an archive without loading its member payloads. + static ZipReader open(const std::string& path); + + // Opens an archive over caller-owned memory. `archive` must outlive this + // reader and every read made through it. + static ZipReader open(ByteSpan archive); + + // Member size, or nullopt when the member is absent. + std::optional member_size(std::string_view name) const; + + // Copies one complete member. + std::vector read(std::string_view name) const; + + // Copies a member range directly into caller-owned storage. + void read_into( + std::string_view name, + size_t offset, + MutableByteSpan destination) const; + + // Reads a complete member and verifies its checksum without retaining it. + void verify(std::string_view name) const; + + const std::vector& names() const { + return names_; + } +}; + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/targets.bzl b/backends/native/runtime/deserialize/targets.bzl new file mode 100644 index 00000000000..80bb83872ab --- /dev/null +++ b/backends/native/runtime/deserialize/targets.bzl @@ -0,0 +1,68 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + # Borrowed byte-range view shared by the package readers (a std::span alias, + # named so the borrow contract has somewhere to live). + runtime.cxx_library( + name = "byte_span", + srcs = [], + exported_headers = ["ByteSpan.h"], + visibility = ["//executorch/backends/native/..."], + ) + + # Owning byte buffer behind a package: heap read or read-only mmap. + runtime.cxx_library( + name = "owned_bytes", + srcs = ["OwnedBytes.cpp"], + exported_headers = ["OwnedBytes.h"], + exported_deps = [":byte_span"], + visibility = ["//executorch/backends/native/..."], + ) + + # JSON representation used by package metadata readers. + runtime.cxx_library( + name = "json", + srcs = [], + exported_headers = ["Json.h"], + exported_external_deps = ["nlohmann_json"], + visibility = ["//executorch/backends/native/..."], + ) + + # Read-only reader for stored (uncompressed) zip archives, which is what a .ptn + # package is. + runtime.cxx_library( + name = "zip_reader", + srcs = ["ZipReader.cpp"], + exported_headers = ["ZipReader.h"], + exported_deps = [":byte_span"], + deps = ["fbsource//third-party/libzip:zip"], + visibility = ["//executorch/backends/native/..."], + ) + + # safetensors index reader. + runtime.cxx_library( + name = "safetensors_reader", + srcs = ["SafeTensorsReader.cpp"], + exported_headers = ["SafeTensorsReader.h"], + exported_deps = [ + ":byte_span", + "//executorch/backends/native/runtime/graph:scalar_type", + ], + deps = [":json"], + visibility = ["//executorch/backends/native/..."], + ) + # The .ptn package: program flatbuffer plus its constants. + runtime.cxx_library( + name = "package", + srcs = ["Package.cpp"], + exported_headers = ["Package.h"], + exported_deps = [ + ":byte_span", + ":owned_bytes", + ":safetensors_reader", + ":zip_reader", + "//executorch/backends/native/runtime/graph:scalar_type", + ], + deps = [":json"], + visibility = ["PUBLIC"], + ) diff --git a/backends/native/runtime/engine/BUCK b/backends/native/runtime/engine/BUCK new file mode 100644 index 00000000000..0ab35888218 --- /dev/null +++ b/backends/native/runtime/engine/BUCK @@ -0,0 +1,11 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +# Any targets that should be shared between fbcode and xplat must be defined in +# targets.bzl. This file can contain cell-only targets. + +non_fbcode_target(_kind = define_common_targets) + +fbcode_target(_kind = define_common_targets) diff --git a/backends/native/runtime/engine/Engine.cpp b/backends/native/runtime/engine/Engine.cpp new file mode 100644 index 00000000000..53385469fc5 --- /dev/null +++ b/backends/native/runtime/engine/Engine.cpp @@ -0,0 +1,17 @@ +// 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. + +#include + +namespace ptn { + +// Out of line so each vtable is emitted here rather than in every translation +// unit that includes the header. +EngineExecutable::~EngineExecutable() = default; + +EngineContext::~EngineContext() = default; + +} // namespace ptn diff --git a/backends/native/runtime/engine/Engine.h b/backends/native/runtime/engine/Engine.h new file mode 100644 index 00000000000..7fc4843cd27 --- /dev/null +++ b/backends/native/runtime/engine/Engine.h @@ -0,0 +1,111 @@ +// 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 + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace ptn { + +// One dependency-closed region of a Method, lowered onto a backend and ready to +// run: whatever the backend needed to compile it, plus the staging it reads +// inputs from and writes outputs to. +// +// Under full delegation -- the only mode today -- that region is the whole +// Method, and the executable's inputs and outputs are the Method's. Under +// runtime partitioning one Method yields several executables interleaved with +// other backends, and the inputs and outputs are region boundaries instead. +// +// Obtained from EngineContext::compile, never constructed directly. Not +// thread-safe and not re-entrant: one executable runs one call at a time. +// Concurrent inference means several executables (or, once the working-set +// split lands, several working sets over one compiled program). +class EngineExecutable { + protected: + EngineExecutable() = default; + + public: + EngineExecutable(const EngineExecutable&) = delete; + EngineExecutable& operator=(const EngineExecutable&) = delete; + EngineExecutable(EngineExecutable&&) = delete; + EngineExecutable& operator=(EngineExecutable&&) = delete; + virtual ~EngineExecutable(); + + // Counts and shapes of the compiled Method's user inputs / outputs, in graph + // order. Sizes are static upper bounds, so a dynamic dim reports its maximum. + virtual size_t num_inputs() const = 0; + virtual size_t num_outputs() const = 0; + virtual std::vector input_sizes(size_t i) const = 0; + virtual std::vector output_sizes(size_t i) const = 0; + virtual ScalarType input_dtype(size_t i) const = 0; + virtual ScalarType output_dtype(size_t i) const = 0; + + // Copy `numel` elements from host `data` into input i, converting from + // `src_dtype` to the input's dtype when they differ. `numel` must equal the + // input's element count. + virtual void + set_input(size_t i, const void* data, size_t numel, ScalarType src_dtype) = 0; + + // Run the compiled Method. Blocks until every output is readable, so a + // get_output right after it needs no further synchronization. + virtual void execute() = 0; + + // Copy output i back into host `data`, converting to `dst_dtype` from the + // output's dtype when they differ. Only meaningful after an execute(). + virtual void + get_output(size_t i, void* data, size_t numel, ScalarType dst_dtype) = 0; +}; + +// A compute backend, at process scope: the device context and kernel registry +// that every Method run on that device shares. One per device, constructed +// through the backend's own factory (e.g. make_vulkan_engine()) since selecting +// a backend is the caller's decision, not this interface's. +// +// Must outlive every executable it compiled. +class EngineContext { + protected: + EngineContext() = default; + + public: + EngineContext(const EngineContext&) = delete; + EngineContext& operator=(const EngineContext&) = delete; + EngineContext(EngineContext&&) = delete; + EngineContext& operator=(EngineContext&&) = delete; + virtual ~EngineContext(); + + // Backend identity ("vulkan"), and the device it selected ("SwiftShader + // Device"). Diagnostics only; nothing dispatches on either. + virtual const std::string& name() const = 0; + virtual const std::string& device_name() const = 0; + + // Lower `method` onto this backend and prepack the constants it binds, + // fetched from `package` by data_key. `method` must outlive the returned + // executable; `package` is needed only for this call. + // + // Compiles the method whole, which is full delegation -- the only mode today. + // Runtime partitioning narrows the unit to a region of a method and yields + // several executables per method; that arrives as an added entry point, not a + // change to this one. + // + // Throws std::runtime_error when the backend cannot run the method: an + // unsupported op or dtype, a binding whose constant the package does not + // hold, a constant whose byte count contradicts its TensorMeta, an unbounded + // dynamic dim, or a higher-order-op subgraph. A backend is free to reject + // anything else it cannot lower; there is no partial success. + virtual std::unique_ptr compile( + const Method& method, + const Package& package) = 0; +}; + +} // namespace ptn diff --git a/backends/native/runtime/engine/targets.bzl b/backends/native/runtime/engine/targets.bzl new file mode 100644 index 00000000000..f94d3f83d7e --- /dev/null +++ b/backends/native/runtime/engine/targets.bzl @@ -0,0 +1,19 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + # The runtime <-> compute-backend boundary: the abstract EngineContext + # (process-wide device state) and EngineExecutable (one lowered, ready-to-run + # Method). Pure std; each backend implements both in its own package. + runtime.cxx_library( + name = "engine", + srcs = ["Engine.cpp"], + exported_headers = [ + "Engine.h", + ], + exported_deps = [ + "//executorch/backends/native/runtime:method", + "//executorch/backends/native/runtime/deserialize:package", + "//executorch/backends/native/runtime/graph:scalar_type", + ], + visibility = ["//executorch/backends/native/..."], + ) diff --git a/backends/native/runtime/targets.bzl b/backends/native/runtime/targets.bzl index 865f2c0242f..6cc4555fb9c 100644 --- a/backends/native/runtime/targets.bzl +++ b/backends/native/runtime/targets.bzl @@ -46,11 +46,16 @@ def define_common_targets(): runtime.cxx_library( name = "runtime", srcs = [ + "Deserialize.cpp", "Program.cpp", ], exported_headers = [ "Program.h", ], + exported_deps = [ + # Program.h publicly exposes Method (build_method), so the IR is exported. + ":method", + ], deps = [ ":native_graph_schema", ], diff --git a/backends/native/test/runtime/BUCK b/backends/native/test/runtime/BUCK new file mode 100644 index 00000000000..36909de98fe --- /dev/null +++ b/backends/native/test/runtime/BUCK @@ -0,0 +1,8 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +non_fbcode_target(_kind = define_common_targets) + +fbcode_target(_kind = define_common_targets) diff --git a/backends/native/test/runtime/deserialize/BUCK b/backends/native/test/runtime/deserialize/BUCK new file mode 100644 index 00000000000..36909de98fe --- /dev/null +++ b/backends/native/test/runtime/deserialize/BUCK @@ -0,0 +1,8 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +non_fbcode_target(_kind = define_common_targets) + +fbcode_target(_kind = define_common_targets) diff --git a/backends/native/test/runtime/deserialize/targets.bzl b/backends/native/test/runtime/deserialize/targets.bzl new file mode 100644 index 00000000000..8a4c80f5c9d --- /dev/null +++ b/backends/native/test/runtime/deserialize/targets.bzl @@ -0,0 +1,36 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + runtime.cxx_test( + name = "owned_bytes_test", + srcs = ["test_owned_bytes.cpp"], + deps = [ + "//executorch/backends/native/runtime/deserialize:owned_bytes", + ], + ) + + runtime.cxx_test( + name = "package_test", + srcs = ["test_package.cpp"], + deps = [ + "//executorch/backends/native/runtime/deserialize:package", + "fbsource//third-party/libzip:zip", + ], + ) + + runtime.cxx_test( + name = "safetensors_reader_test", + srcs = ["test_safetensors_reader.cpp"], + deps = [ + "//executorch/backends/native/runtime/deserialize:safetensors_reader", + ], + ) + + runtime.cxx_test( + name = "zip_reader_test", + srcs = ["test_zip_reader.cpp"], + deps = [ + "//executorch/backends/native/runtime/deserialize:zip_reader", + "fbsource//third-party/libzip:zip", + ], + ) diff --git a/backends/native/test/runtime/deserialize/test_owned_bytes.cpp b/backends/native/test/runtime/deserialize/test_owned_bytes.cpp new file mode 100644 index 00000000000..43d97f765eb --- /dev/null +++ b/backends/native/test/runtime/deserialize/test_owned_bytes.cpp @@ -0,0 +1,131 @@ +// 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. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace ptn { +namespace { + +class OwnedBytesTest : public ::testing::Test { + private: + std::vector paths_; + + protected: + std::string temp_path(std::string_view suffix) { + const auto* info = ::testing::UnitTest::GetInstance()->current_test_info(); + const std::filesystem::path path = std::filesystem::temp_directory_path() / + (std::string("owned_bytes_") + info->name() + std::string(suffix)); + std::error_code error; + std::filesystem::remove(path, error); + paths_.push_back(path); + return path.string(); + } + + void write_file(const std::string& path, std::string_view contents) { + std::ofstream file(path, std::ios::binary | std::ios::trunc); + ASSERT_TRUE(file); + file.write(contents.data(), static_cast(contents.size())); + ASSERT_TRUE(file); + } + + void TearDown() override { + for (const std::filesystem::path& path : paths_) { + std::error_code error; + std::filesystem::remove(path, error); + } + } +}; + +static_assert(!std::is_copy_constructible_v); +static_assert(!std::is_copy_assignable_v); +static_assert(std::is_nothrow_move_constructible_v); +static_assert(std::is_nothrow_move_assignable_v); + +// cppcheck-suppress-begin syntaxError +TEST_F(OwnedBytesTest, TakesOwnershipOfVector) { + std::vector source{1, 2, 3}; + const uint8_t* data = source.data(); + + OwnedBytes bytes = OwnedBytes::from_vector(std::move(source)); + const ByteSpan span = bytes.span(); + OwnedBytes moved = std::move(bytes); + + EXPECT_FALSE(moved.is_mapped()); + EXPECT_EQ(span.data(), data); + EXPECT_EQ(moved.span().data(), data); + EXPECT_EQ( + std::vector(span.begin(), span.end()), + (std::vector{1, 2, 3})); +} + +TEST_F(OwnedBytesTest, ReadsFileIntoHeap) { + const std::string path = temp_path("_heap.bin"); + ASSERT_NO_FATAL_FAILURE(write_file(path, "abc")); + + OwnedBytes bytes = OwnedBytes::from_file(path, false); + const ByteSpan span = bytes.span(); + const OwnedBytes moved = std::move(bytes); + + EXPECT_FALSE(moved.is_mapped()); + EXPECT_TRUE(bytes.span().empty()); + EXPECT_EQ( + std::vector(span.begin(), span.end()), + (std::vector{'a', 'b', 'c'})); +} + +TEST_F(OwnedBytesTest, MapsFile) { + const std::string path = temp_path("_mapped.bin"); + ASSERT_NO_FATAL_FAILURE(write_file(path, "abc")); + +#if defined(_WIN32) + EXPECT_THROW(OwnedBytes::from_file(path), std::runtime_error); +#else + OwnedBytes bytes = OwnedBytes::from_file(path); + const ByteSpan span = bytes.span(); + const OwnedBytes moved = std::move(bytes); + EXPECT_TRUE(moved.is_mapped()); + EXPECT_TRUE(bytes.span().empty()); + EXPECT_EQ( + std::vector(span.begin(), span.end()), + (std::vector{'a', 'b', 'c'})); +#endif +} + +TEST_F(OwnedBytesTest, EmptyFileUsesHeapStorage) { + const std::string path = temp_path("_empty.bin"); + ASSERT_NO_FATAL_FAILURE(write_file(path, "")); + + const OwnedBytes bytes = OwnedBytes::from_file(path); + + EXPECT_FALSE(bytes.is_mapped()); + EXPECT_TRUE(bytes.span().empty()); +} + +TEST_F(OwnedBytesTest, RejectsInvalidPaths) { + EXPECT_THROW( + OwnedBytes::from_file(temp_path("_missing.bin"), false), + std::runtime_error); + EXPECT_THROW( + OwnedBytes::from_file( + std::filesystem::temp_directory_path().string(), false), + std::runtime_error); +} +// cppcheck-suppress-end syntaxError + +} // namespace +} // namespace ptn diff --git a/backends/native/test/runtime/deserialize/test_package.cpp b/backends/native/test/runtime/deserialize/test_package.cpp new file mode 100644 index 00000000000..0d215409e5d --- /dev/null +++ b/backends/native/test/runtime/deserialize/test_package.cpp @@ -0,0 +1,159 @@ +// 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. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace ptn { +namespace { + +std::vector make_safetensors( + std::string_view header, + std::string_view data) { + const uint64_t header_size = header.size(); + std::vector bytes(sizeof(header_size) + header.size() + data.size()); + std::memcpy(bytes.data(), &header_size, sizeof(header_size)); + std::memcpy(bytes.data() + sizeof(header_size), header.data(), header.size()); + std::memcpy( + bytes.data() + sizeof(header_size) + header.size(), + data.data(), + data.size()); + return bytes; +} + +ByteSpan as_bytes(std::string_view value) { + return ByteSpan(reinterpret_cast(value.data()), value.size()); +} + +struct ZipDiscard { + void operator()(zip_t* archive) const noexcept { + zip_discard(archive); + } +}; + +class TempPackage { + private: + const std::string program_ = "program"; + const std::string aliases_ = R"({"tied_weight":"weight"})"; + std::filesystem::path path_; + + public: + TempPackage() { + path_ = std::filesystem::temp_directory_path() / + ("ptn_package_" + std::to_string(reinterpret_cast(this)) + + ".ptn"); + int error = 0; + std::unique_ptr archive( + zip_open(path_.string().c_str(), ZIP_CREATE | ZIP_TRUNCATE, &error)); + if (archive == nullptr) { + throw std::runtime_error("failed to create test package"); + } + const std::vector tensors = make_safetensors( + R"({"weight":{"dtype":"U8","shape":[4],"data_offsets":[0,4]},"bias":{"dtype":"I16","shape":[1],"data_offsets":[4,6]}})", + "dataxy"); + add(archive.get(), kProgramEntry, as_bytes(program_)); + add(archive.get(), kSafeTensorsEntry, ByteSpan(tensors)); + add(archive.get(), kAliasesEntry, as_bytes(aliases_)); + if (zip_close(archive.get()) != 0) { + throw std::runtime_error("failed to close test package"); + } + archive.release(); + } + + ~TempPackage() { + std::error_code error; + std::filesystem::remove(path_, error); + } + + std::string path() const { + return path_.string(); + } + + private: + static void add(zip_t* archive, const char* name, ByteSpan bytes) { + zip_source_t* source = + zip_source_buffer(archive, bytes.data(), bytes.size(), 0); + if (source == nullptr) { + throw std::runtime_error("failed to create test package source"); + } + const zip_int64_t index = + zip_file_add(archive, name, source, ZIP_FL_ENC_UTF_8); + if (index < 0) { + zip_source_free(source); + throw std::runtime_error("failed to add test package member"); + } + if (zip_set_file_compression(archive, index, ZIP_CM_STORE, 0) != 0) { + throw std::runtime_error("failed to store test package member"); + } + } +}; + +// cppcheck-suppress-begin syntaxError +TEST(PackageTest, LoadsConstantsOnDemand) { + const TempPackage file; + const Package package = Package::load(file.path()); + + EXPECT_EQ(package.owner_keys(), (std::vector{"weight", "bias"})); + EXPECT_EQ(package.constant_bytes(), 6); + + const std::optional info = package.constant_info("tied_weight"); + ASSERT_TRUE(info); + EXPECT_NE(info->package_id, 0); + EXPECT_EQ(info->dtype, kByte); + EXPECT_EQ(*info->sizes, (std::vector{4})); + EXPECT_EQ(info->nbytes, 4); + EXPECT_EQ(info->owner, "weight"); + + std::array destination{}; + EXPECT_TRUE(package.load_constant_into("weight", destination)); + EXPECT_EQ(destination, (std::array{'d', 'a', 't', 'a'})); + + const std::optional acquired = + package.acquire_constant("tied_weight"); + ASSERT_TRUE(acquired); + EXPECT_TRUE(std::ranges::equal(acquired->span(), destination)); + EXPECT_NO_THROW(package.verify_constants()); +} + +TEST(PackageTest, ReportsMissingConstantsAndWrongDestinations) { + const TempPackage file; + const Package package = Package::load(file.path()); + + EXPECT_EQ(package.constant_info("missing"), std::nullopt); + EXPECT_EQ(package.acquire_constant("missing"), std::nullopt); + std::array destination{}; + EXPECT_FALSE(package.load_constant_into("missing", destination)); + EXPECT_THROW( + package.load_constant_into("weight", destination), std::runtime_error); +} + +TEST(PackageTest, SupportsCallerOwnedArchiveBytes) { + const TempPackage file; + Package package = Package::load(file.path()); + package = Package::load(OwnedBytes::from_file(file.path(), false)); + + const std::optional weight = package.acquire_constant("weight"); + ASSERT_TRUE(weight); + EXPECT_TRUE(std::ranges::equal( + weight->span(), (std::array{'d', 'a', 't', 'a'}))); +} +// cppcheck-suppress-end syntaxError + +} // namespace +} // namespace ptn diff --git a/backends/native/test/runtime/deserialize/test_safetensors_reader.cpp b/backends/native/test/runtime/deserialize/test_safetensors_reader.cpp new file mode 100644 index 00000000000..e404017b40e --- /dev/null +++ b/backends/native/test/runtime/deserialize/test_safetensors_reader.cpp @@ -0,0 +1,93 @@ +// 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. + +#include + +#include +#include +#include +#include +#include + +#include + +namespace ptn { +namespace { + +std::vector make_safetensors( + std::string_view header, + std::string_view data) { + const uint64_t header_size = header.size(); + std::vector bytes(sizeof(header_size) + header.size() + data.size()); + std::memcpy(bytes.data(), &header_size, sizeof(header_size)); + std::memcpy(bytes.data() + sizeof(header_size), header.data(), header.size()); + std::memcpy( + bytes.data() + sizeof(header_size) + header.size(), + data.data(), + data.size()); + return bytes; +} + +// cppcheck-suppress-begin syntaxError +TEST(SafeTensorsReaderTest, ReadsIndexAndPayloadsInHeaderOrder) { + const std::vector bytes = make_safetensors( + R"({"second":{"dtype":"U8","shape":[2],"data_offsets":[4,6]},"__metadata__":{},"first":{"dtype":"F32","shape":[1],"data_offsets":[0,4]}})", + "abcdXY"); + + const SafeTensorsReader reader = SafeTensorsReader::open(bytes); + + EXPECT_EQ(reader.names(), (std::vector{"second", "first"})); + ASSERT_NE(reader.find("first"), nullptr); + EXPECT_EQ(reader.find("first")->dtype, kFloat); + EXPECT_EQ(reader.find("first")->sizes, (std::vector{1})); + EXPECT_EQ(reader.find("first")->offset, 0); + EXPECT_EQ(reader.find("first")->nbytes, 4); + EXPECT_EQ(reader.total_bytes(), 6); + EXPECT_EQ(reader.find("missing"), nullptr); +} + +TEST(SafeTensorsReaderTest, RejectsNonEmptyMetadata) { + EXPECT_THROW( + SafeTensorsReader::open( + make_safetensors(R"({"__metadata__":{"source":"test"}})", "")), + std::runtime_error); +} + +TEST(SafeTensorsReaderTest, ReadsHeaderWithoutPayload) { + const std::string header = + R"({"x":{"dtype":"U8","shape":[2],"data_offsets":[0,2]}})"; + const SafeTensorsReader reader = SafeTensorsReader::open_header( + ByteSpan(reinterpret_cast(header.data()), header.size()), + 2); + + ASSERT_NE(reader.find("x"), nullptr); + EXPECT_EQ(reader.find("x")->nbytes, 2); +} + +TEST(SafeTensorsReaderTest, RejectsInvalidMetadata) { + EXPECT_THROW( + SafeTensorsReader::open(make_safetensors("[]", "")), std::runtime_error); + EXPECT_THROW( + SafeTensorsReader::open(make_safetensors( + R"({"x":{"dtype":"U8","shape":[1.0],"data_offsets":[0,1]}})", "x")), + std::runtime_error); + EXPECT_THROW( + SafeTensorsReader::open(make_safetensors( + R"({"x":{"dtype":"U8","shape":[1],"data_offsets":[0,2]}})", "x")), + std::runtime_error); +} + +TEST(SafeTensorsReaderTest, RejectsByteSizeOverflow) { + EXPECT_THROW( + SafeTensorsReader::open(make_safetensors( + R"({"x":{"dtype":"F64","shape":[2305843009213693952],"data_offsets":[0,0]}})", + "")), + std::runtime_error); +} +// cppcheck-suppress-end syntaxError + +} // namespace +} // namespace ptn diff --git a/backends/native/test/runtime/deserialize/test_zip_reader.cpp b/backends/native/test/runtime/deserialize/test_zip_reader.cpp new file mode 100644 index 00000000000..62cdaf4706c --- /dev/null +++ b/backends/native/test/runtime/deserialize/test_zip_reader.cpp @@ -0,0 +1,114 @@ +// 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. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace ptn { +namespace { + +struct ZipDiscard { + void operator()(zip_t* archive) const noexcept { + zip_discard(archive); + } +}; + +class TempZip { + public: + TempZip() { + path_ = std::filesystem::temp_directory_path() / + ("ptn_zip_reader_" + std::to_string(reinterpret_cast(this)) + + ".zip"); + int error = 0; + std::unique_ptr archive( + zip_open(path_.string().c_str(), ZIP_CREATE | ZIP_TRUNCATE, &error)); + if (archive == nullptr) { + throw std::runtime_error("failed to create test zip"); + } + add(archive.get(), "program.ptg", "program"); + add(archive.get(), "program.safetensors", "0123456789"); + if (zip_close(archive.get()) != 0) { + throw std::runtime_error("failed to close test zip"); + } + archive.release(); + } + + ~TempZip() { + std::error_code error; + std::filesystem::remove(path_, error); + } + + const std::string path() const { + return path_.string(); + } + + private: + static void add(zip_t* archive, const char* name, std::string_view bytes) { + zip_source_t* source = + zip_source_buffer(archive, bytes.data(), bytes.size(), 0); + if (source == nullptr) { + throw std::runtime_error("failed to create test zip source"); + } + const zip_int64_t index = + zip_file_add(archive, name, source, ZIP_FL_ENC_UTF_8); + if (index < 0) { + zip_source_free(source); + throw std::runtime_error("failed to add test zip member"); + } + if (zip_set_file_compression(archive, index, ZIP_CM_STORE, 0) != 0) { + throw std::runtime_error("failed to store test zip member"); + } + } + + std::filesystem::path path_; +}; + +// cppcheck-suppress-begin syntaxError +TEST(ZipReaderTest, ReadsStoredMemberRanges) { + const TempZip file; + ZipReader zip = ZipReader::open(file.path()); + + EXPECT_EQ( + zip.names(), + (std::vector{"program.ptg", "program.safetensors"})); + EXPECT_EQ(zip.member_size("program.safetensors"), 10); + EXPECT_EQ(zip.member_size("missing"), std::nullopt); + + std::array bytes{}; + zip.read_into("program.safetensors", 3, MutableByteSpan(bytes)); + EXPECT_EQ(bytes, (std::array{'3', '4', '5', '6'})); + EXPECT_EQ( + zip.read("program.ptg"), + (std::vector{'p', 'r', 'o', 'g', 'r', 'a', 'm'})); + EXPECT_NO_THROW(zip.verify("program.safetensors")); +} + +TEST(ZipReaderTest, RejectsInvalidRanges) { + const TempZip file; + ZipReader zip = ZipReader::open(file.path()); + + std::array bytes{}; + EXPECT_THROW( + zip.read_into("program.safetensors", 8, MutableByteSpan(bytes)), + std::runtime_error); + EXPECT_THROW( + zip.read_into("missing", 0, MutableByteSpan(bytes)), std::runtime_error); +} +// cppcheck-suppress-end syntaxError + +} // namespace +} // namespace ptn diff --git a/backends/native/test/runtime/targets.bzl b/backends/native/test/runtime/targets.bzl new file mode 100644 index 00000000000..97f9110f7d9 --- /dev/null +++ b/backends/native/test/runtime/targets.bzl @@ -0,0 +1,11 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + runtime.cxx_test( + name = "program_test", + srcs = ["test_program_deserialize.cpp"], + deps = [ + "//executorch/backends/native/runtime:native_graph_schema", + "//executorch/backends/native/runtime:runtime", + ], + ) diff --git a/backends/native/test/runtime/test_program_deserialize.cpp b/backends/native/test/runtime/test_program_deserialize.cpp new file mode 100644 index 00000000000..ec3cde4f10c --- /dev/null +++ b/backends/native/test/runtime/test_program_deserialize.cpp @@ -0,0 +1,299 @@ +// 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. + +#include + +#include +#include +#include + +#include +#include + +#include + +namespace ptn { +namespace { + +namespace fbs = ::native_backend; + +flatbuffers::Offset< + flatbuffers::Vector>> +create_strings( + flatbuffers::FlatBufferBuilder& builder, + const std::vector& values) { + std::vector> strings; + strings.reserve(values.size()); + for (const std::string& value : values) { + strings.push_back(builder.CreateString(value)); + } + return builder.CreateVector(strings); +} + +flatbuffers::Offset create_graph( + flatbuffers::FlatBufferBuilder& builder, + const std::vector>& nodes = {}, + const std::vector& inputs = {}, + const std::vector& outputs = {}, + const std::vector>& tensor_values = + {}) { + return fbs::CreateGraph( + builder, + builder.CreateVector(nodes), + inputs.empty() ? 0 : create_strings(builder, inputs), + outputs.empty() ? 0 : create_strings(builder, outputs), + tensor_values.empty() ? 0 : builder.CreateVector(tensor_values)); +} + +flatbuffers::Offset create_method( + flatbuffers::FlatBufferBuilder& builder, + const std::string& name, + flatbuffers::Offset graph, + const std::vector>& output_specs = {}, + const std::vector>& constants = {}, + const std::vector>& + mutable_buffers = {}) { + return fbs::CreateMethod( + builder, + builder.CreateString(name), + graph, + constants.empty() ? 0 : builder.CreateVector(constants), + output_specs.empty() ? 0 : builder.CreateVector(output_specs), + mutable_buffers.empty() ? 0 : builder.CreateVector(mutable_buffers)); +} + +std::vector finish_program( + flatbuffers::FlatBufferBuilder& builder, + const std::vector>& methods) { + const auto program = fbs::CreateProgram( + builder, builder.CreateString("1"), builder.CreateVector(methods)); + fbs::FinishProgramBuffer(builder, program); + return { + builder.GetBufferPointer(), + builder.GetBufferPointer() + builder.GetSize()}; +} + +Program load_program(const std::vector& bytes) { + return Program::load(bytes.data(), bytes.size()); +} + +// cppcheck-suppress-begin syntaxError +TEST(ProgramTest, LoadRejectsEmptyMethodName) { + flatbuffers::FlatBufferBuilder builder; + const auto graph = create_graph(builder); + const auto bytes = + finish_program(builder, {create_method(builder, "", graph)}); + + EXPECT_THROW(load_program(bytes), std::runtime_error); +} + +TEST(ProgramTest, LoadRejectsDuplicateMethodNames) { + flatbuffers::FlatBufferBuilder builder; + const auto graph = create_graph(builder); + const auto first = create_method(builder, "forward", graph); + const auto second = create_method(builder, "forward", graph); + const auto bytes = finish_program(builder, {first, second}); + + EXPECT_THROW(load_program(bytes), std::runtime_error); +} + +TEST(ProgramTest, GetMethodRejectsMismatchedOutputSpecs) { + flatbuffers::FlatBufferBuilder builder; + const auto graph = create_graph(builder, {}, {}, {"output"}); + const auto first = fbs::CreateOutputSpecDirect(builder, "output"); + const auto second = fbs::CreateOutputSpecDirect(builder, "extra"); + const auto method = create_method(builder, "forward", graph, {first, second}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + + EXPECT_THROW(program.get_method("forward"), std::runtime_error); +} + +TEST(ProgramTest, GetMethodPreservesAliasWhenTargetIsCreatedOnDemand) { + flatbuffers::FlatBufferBuilder builder; + const auto output = fbs::CreateOutputDirect(builder, "view", "input"); + const std::vector> outputs = {output}; + const auto node = fbs::CreateNodeDirect( + builder, + "view", + fbs::OpKind::CALL_FUNCTION, + "aten.view", + nullptr, + &outputs); + const auto graph = create_graph(builder, {node}, {"input"}, {"view"}); + const auto bytes = + finish_program(builder, {create_method(builder, "forward", graph)}); + const Program program = load_program(bytes); + + const Graph& loaded = program.get_method("forward").graph; + ASSERT_EQ(loaded.input_ids.size(), 1); + ASSERT_EQ(loaded.output_ids.size(), 1); + EXPECT_EQ(loaded.value(loaded.output_ids[0]).alias_id, loaded.input_ids[0]); +} + +TEST(ProgramTest, GetMethodRejectsSelfAlias) { + flatbuffers::FlatBufferBuilder builder; + const auto output = fbs::CreateOutputDirect(builder, "view", "view"); + const std::vector> outputs = {output}; + const auto node = fbs::CreateNodeDirect( + builder, + "view", + fbs::OpKind::CALL_FUNCTION, + "aten.view", + nullptr, + &outputs); + const auto graph = create_graph(builder, {node}, {}, {"view"}); + const auto bytes = + finish_program(builder, {create_method(builder, "forward", graph)}); + const Program program = load_program(bytes); + + EXPECT_THROW(program.get_method("forward"), std::runtime_error); +} + +TEST(ProgramTest, GetMethodRejectsDynamicTensorExtent) { + flatbuffers::FlatBufferBuilder builder; + const std::vector> sizes = { + fbs::CreateDim(builder, 2, 16)}; + const auto meta = + fbs::CreateTensorMetaDirect(builder, fbs::ScalarType::FLOAT, &sizes); + const auto tensor = fbs::CreateTensorValueDirect(builder, "input", meta); + const auto graph = create_graph(builder, {}, {"input"}, {}, {tensor}); + const auto bytes = + finish_program(builder, {create_method(builder, "forward", graph)}); + const Program program = load_program(bytes); + + EXPECT_THROW(program.get_method("forward"), std::runtime_error); +} + +TEST(ProgramTest, GetMethodRejectsUnknownEnumValues) { + { + flatbuffers::FlatBufferBuilder builder; + const auto node = fbs::CreateNodeDirect( + builder, "node", static_cast(127), "unknown"); + const auto graph = create_graph(builder, {node}); + const auto bytes = + finish_program(builder, {create_method(builder, "forward", graph)}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } + { + flatbuffers::FlatBufferBuilder builder; + const auto output = fbs::CreateOutput( + builder, + builder.CreateString("output"), + 0, + static_cast(127)); + const std::vector> outputs = {output}; + const auto node = fbs::CreateNodeDirect( + builder, + "node", + fbs::OpKind::CALL_FUNCTION, + "unknown", + nullptr, + &outputs); + const auto graph = create_graph(builder, {node}, {}, {"output"}); + const auto bytes = + finish_program(builder, {create_method(builder, "forward", graph)}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } + { + flatbuffers::FlatBufferBuilder builder; + const auto meta = fbs::CreateTensorMeta(builder); + const auto constant = fbs::CreateNamedTensorRefDirect( + builder, "input", "weight", meta, static_cast(127)); + const auto graph = create_graph(builder, {}, {"input"}); + const auto method = + create_method(builder, "forward", graph, {}, {constant}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } + { + flatbuffers::FlatBufferBuilder builder; + const auto argument = + fbs::CreateArgument(builder, static_cast(127), 0); + const auto named_argument = + fbs::CreateNamedArgumentDirect(builder, "input", argument); + const std::vector> inputs = { + named_argument}; + const auto node = fbs::CreateNodeDirect( + builder, "node", fbs::OpKind::CALL_FUNCTION, "unknown", &inputs); + const auto graph = create_graph(builder, {node}); + const auto bytes = + finish_program(builder, {create_method(builder, "forward", graph)}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } + { + flatbuffers::FlatBufferBuilder builder; + const auto output_spec = fbs::CreateOutputSpec( + builder, + builder.CreateString("output"), + static_cast(127)); + const auto graph = create_graph(builder, {}, {}, {"output"}); + const auto method = create_method(builder, "forward", graph, {output_spec}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } +} + +TEST(ProgramTest, GetMethodRejectsUnresolvedDataBinding) { + flatbuffers::FlatBufferBuilder builder; + const auto meta = fbs::CreateTensorMeta(builder); + const auto constant = fbs::CreateNamedTensorRefDirect( + builder, "missing", "weight", meta, fbs::InputKind::PARAMETER); + const auto graph = create_graph(builder); + const auto method = create_method(builder, "forward", graph, {}, {constant}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + + EXPECT_THROW(program.get_method("forward"), std::runtime_error); +} + +TEST(ProgramTest, GetMethodRejectsUnresolvedMutationTarget) { + { + flatbuffers::FlatBufferBuilder builder; + const auto output_spec = fbs::CreateOutputSpecDirect( + builder, "output", fbs::OutputKind::BUFFER_MUTATION, "missing"); + const auto graph = create_graph(builder, {}, {}, {"output"}); + const auto method = create_method(builder, "forward", graph, {output_spec}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } + { + flatbuffers::FlatBufferBuilder builder; + const auto output_spec = fbs::CreateOutputSpecDirect( + builder, "output", fbs::OutputKind::USER_INPUT_MUTATION, "missing"); + const auto graph = create_graph(builder, {}, {}, {"output"}); + const auto method = create_method(builder, "forward", graph, {output_spec}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } +} + +TEST(ProgramTest, GetMethodRejectsDuplicateDataBinding) { + flatbuffers::FlatBufferBuilder builder; + const auto meta = fbs::CreateTensorMeta(builder); + const auto constant = fbs::CreateNamedTensorRefDirect( + builder, "state", "state", meta, fbs::InputKind::BUFFER); + const auto mutable_buffer = + fbs::CreateMutableBufferSpecDirect(builder, "state", "state"); + const auto graph = create_graph(builder, {}, {"state"}); + const auto method = create_method( + builder, "forward", graph, {}, {constant}, {mutable_buffer}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + + EXPECT_THROW(program.get_method("forward"), std::runtime_error); +} +// cppcheck-suppress-end syntaxError + +} // namespace +} // namespace ptn diff --git a/backends/nxp/BUCK b/backends/nxp/BUCK index 6dec42f04d7..81a1e04cfee 100644 --- a/backends/nxp/BUCK +++ b/backends/nxp/BUCK @@ -68,7 +68,6 @@ fbcode_target(_kind = runtime.python_library, "fbsource//third-party/pypi/neutron_converter:neutron_converter", "//caffe2:torch", "//executorch/exir:lib", - "//executorch/backends/nxp/tests:ops_aliases", ], ) diff --git a/backends/nxp/README.md b/backends/nxp/README.md index 4188dd8f810..05204fd7c78 100644 --- a/backends/nxp/README.md +++ b/backends/nxp/README.md @@ -33,9 +33,9 @@ The eIQ Neutron NPU Backend should be considered as prototype quality at this mo improvements. NXP and the ExecuTorch community is actively developing this codebase. ## Neutron Backend implementation and SW architecture -Neutron Backend uses the eIQ Neutron Converter as ML compiler to compile the delegated subgraph to Neutron microcode. -The Neutron Converter accepts the ML model in LiteRT format, for the **eIQ Neutron N3** class therefore the Neutron Backend -uses the LiteRT flatbuffers format as IR between the ExecuTorch and Neutron Converter ML compiler. +Neutron Backend uses the eIQ Neutron Compiler as ML compiler to compile the delegated subgraph to Neutron microcode. +The Neutron Compiler accepts the ML model in LiteRT format, for the **eIQ Neutron N3** class therefore the Neutron Backend +uses the LiteRT flatbuffers format as IR between the ExecuTorch and Neutron Compiler ML compiler. ## Layout * `backend/ir/` - TFLite/LiteRT based IR to represent the Edge Subgraph, taken from onnx2tflite code base and extended to diff --git a/backends/nxp/aten_passes/add_batch_size_for_3d_input_pool_2d_ops.py b/backends/nxp/aten_passes/add_batch_size_for_3d_input_pool_2d_ops.py index 7435b3b6969..cbd000befbb 100644 --- a/backends/nxp/aten_passes/add_batch_size_for_3d_input_pool_2d_ops.py +++ b/backends/nxp/aten_passes/add_batch_size_for_3d_input_pool_2d_ops.py @@ -13,7 +13,7 @@ class AddBatchSizeFor3DInputPool2DOps(PassBase): """Adds batch size dimension for aten.adaptive_avg_pool2d.default, aten.avg_pool2d.default - and aten.max_pool2d.default ops with 3D input, as the Neutron Converter is unable to convert these ops with 3D input. + and aten.max_pool2d.default ops with 3D input, as the Neutron Compiler is unable to compile these ops with 3D input. │ ┌──────▼──────┐ diff --git a/backends/nxp/backend/edge_helper.py b/backends/nxp/backend/edge_helper.py index 408b90e264d..4708b6dfdd0 100644 --- a/backends/nxp/backend/edge_helper.py +++ b/backends/nxp/backend/edge_helper.py @@ -8,7 +8,7 @@ import torch -from executorch.backends.nxp.tests.ops_aliases import ( +from executorch.backends.nxp.backend.ops_aliases import ( AddTensor, Amax, Amin, diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py index bd812cccb76..d4029ebebce 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py @@ -6,6 +6,7 @@ import executorch.backends.nxp.backend.ir.lib.tflite.Padding as tflPadding import torch +from executorch.backends.nxp.backend.edge_helper import input_rank from executorch.backends.nxp.backend.ir.converter.conversion import common from executorch.backends.nxp.backend.ir.converter.node_converter import ( CustomDelegationOptions, @@ -45,6 +46,10 @@ def _is_supported_in_IR( parameters_mapping: dict[str, Parameter], custom_delegation_options: CustomDelegationOptions, ) -> bool: + # The input must be 4D. + if input_rank(node, 0) != 4: + return False + input_size = node.args[0].meta["val"].shape output_size = node.args[1] diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py index b3157ab4c4b..5791edcfac9 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py @@ -5,6 +5,8 @@ import numpy as np import torch + +from executorch.backends.nxp.backend.edge_helper import input_rank from executorch.backends.nxp.backend.ir.converter.conversion import ( aten_translator, common, @@ -36,6 +38,10 @@ def _is_supported_in_IR( parameters_mapping: dict[str, Parameter], custom_delegation_options: CustomDelegationOptions, ) -> bool: + # The input must be 4D. + if input_rank(node, 0) != 4: + return False + n_args = len(node.args) padding = node.args[3] if n_args >= 4 else [0, 0] 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 c8d24ea34a6..507de4283c1 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 @@ -7,7 +7,7 @@ import numpy as np import torch -from executorch.backends.nxp.backend.edge_helper import try_get_arg +from executorch.backends.nxp.backend.edge_helper import input_rank, try_get_arg from executorch.backends.nxp.backend.ir.converter.conversion import ( aten_translator, common, @@ -42,6 +42,10 @@ def _is_supported_in_IR( parameters_mapping: dict[str, Parameter], custom_delegation_options: CustomDelegationOptions, ) -> bool: + # The input must be 4D. + if input_rank(node, 0) != 4: + return False + kernel_size, stride, padding, dilation, ceil_mode = ( MaxPool2DWithIndicesConverter._get_node_args(node) ) diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py index 4d03e5e97b7..b9c5a11ca35 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py @@ -41,7 +41,7 @@ def supports_partitioning_result( is_alone_in_partition = cls.is_node_alone_in_partition(node, partition_list) if is_alone_in_partition and keepdim and all(input_shape[d] == 1 for d in dim): - # The operator is a no-op, so the Neutron Converter will skip it. If it's the only node in the + # The operator is a no-op, so the Neutron Compiler will skip it. If it's the only node in the # partition, the graph would end up empty. return False diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/permute_copy_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/permute_copy_converter.py index 3e4908c2211..6f0eb6ad757 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/permute_copy_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/permute_copy_converter.py @@ -400,7 +400,7 @@ def supports_partitioning_result( is_alone_in_partition = cls.is_node_alone_in_partition(node, partition_list) if has_static_input and is_alone_in_partition: # Transpose with a static input is a no-op on Neutron. If it was the only operator in the partition, - # Neutron Converter would produce and empty graph, so delegation is prohibited. + # Neutron Compiler would produce and empty graph, so delegation is prohibited. return False return True diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_bilinear2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_bilinear2d_converter.py index 2f0126e9aae..544638f5b5b 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_bilinear2d_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_bilinear2d_converter.py @@ -39,7 +39,7 @@ def supports_partitioning_result( is_alone_in_partition = cls.is_node_alone_in_partition(node, partition_list) if is_alone_in_partition and input_shape == output_shape: - # The operator is a no-op, so the Neutron Converter will skip it. If it's the only node in the + # The operator is a no-op, so the Neutron Compiler will skip it. If it's the only node in the # partition, the graph would end up empty. return False diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py index a3c8db14f51..cdf8ad7d46e 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py @@ -41,7 +41,7 @@ def supports_partitioning_result( is_alone_in_partition = cls.is_node_alone_in_partition(node, partition_list) if is_alone_in_partition and h_scale == w_scale == 1: - # The operator is a no-op, so the Neutron Converter will skip it. If it's the only node in the + # The operator is a no-op, so the Neutron Compiler will skip it. If it's the only node in the # partition, the graph would end up empty. return False diff --git a/backends/nxp/backend/neutron_converter_manager.py b/backends/nxp/backend/neutron_compiler_manager.py similarity index 68% rename from backends/nxp/backend/neutron_converter_manager.py rename to backends/nxp/backend/neutron_compiler_manager.py index 92b4e25a5de..f7ca416cb0c 100644 --- a/backends/nxp/backend/neutron_converter_manager.py +++ b/backends/nxp/backend/neutron_compiler_manager.py @@ -6,19 +6,35 @@ import logging import multiprocessing import os +import warnings try: - from eiq_neutron_sdk import neutron_converter, neutron_library_utils + from eiq_neutron_sdk import neutron_compiler, neutron_library_utils + + _USING_NEUTRON_COMPILER = True except ImportError: - raise RuntimeError( - "eIQ Neutron SDK not found. To install it, run 'examples/nxp/setup.sh'." - ) + try: + from eiq_neutron_sdk import ( + neutron_converter as neutron_compiler, + neutron_library_utils, + ) + + _USING_NEUTRON_COMPILER = False + warnings.warn( + "The support for eIQ Neutron SDK <= 3.2.2 will be removed in future releases.", + DeprecationWarning, + stacklevel=2, + ) + except ImportError: + raise RuntimeError( + "eIQ Neutron SDK not found. To install it, run 'examples/nxp/setup.sh'." + ) def _build_compilation_context(compilation_opts): """Build a CompilationContext from a plain dict of options.""" - cctx = neutron_converter.CompilationContext() - cctx.targetOpts = neutron_converter.getNeutronTarget(compilation_opts["target"]) + cctx = neutron_compiler.CompilationContext() + cctx.targetOpts = neutron_compiler.getNeutronTarget(compilation_opts["target"]) cctx.compilationOpts.minNumOpsPerGraph = compilation_opts["minNumOpsPerGraph"] cctx.compilationOpts.excludeGraphPasses = compilation_opts["excludeGraphPasses"] cctx.compilationOpts.fetchConstantsToSRAM = compilation_opts["fetchConstantsToSRAM"] @@ -37,19 +53,22 @@ def _build_compilation_context(compilation_opts): return cctx -def convert_unsafe(tflite_model, compilation_opts, queue): +def compile_unsafe(tflite_model, compilation_opts, queue): """ - Run neutron_converter on given tflite_model with the provided compilation options. + Run neutron_compiler on given tflite_model with the provided compilation options. This routine is supposed to run in a separate process. - If properly finished, the output queue contains the converted model, - otherwise the neutron_converter exits and the output queue is empty. + If properly finished, the output queue contains the compiled model, + otherwise the neutron_compiler exits and the output queue is empty. """ cctx = _build_compilation_context(compilation_opts) - model_converted = neutron_converter.convertModel(list(tflite_model), cctx) - queue.put(model_converted) + if _USING_NEUTRON_COMPILER: + model_compiled = neutron_compiler.compileModel(list(tflite_model), cctx) + else: + model_compiled = neutron_compiler.convertModel(list(tflite_model), cctx) + queue.put(model_compiled) -class NeutronConverterManager: +class NeutronCompilerManager: """ Manager for conversion of TFLite model in flatbuffers format into TFLite model that contains NeutronGraph nodes. @@ -69,8 +88,8 @@ def _rename_partition_kernel_selection_file(delegation_tag): except OSError: logging.error("Failed to rename partition kernel selection file.") - def get_converter(self): - return neutron_converter + def get_compiler(self): + return neutron_compiler def get_library_utils(self): return neutron_library_utils @@ -84,7 +103,7 @@ def verify_target(self, target: str): f"Target `{target}` is not a valid target. Must be one of `{valid_targets}`." ) - def convert( + def compile( self, tflite_model: bytes, target: str, @@ -93,9 +112,9 @@ def convert( use_profiling: bool = False, ) -> bytes: """ - Call Neutron Converter. + Call Neutron Compiler. - :param tflite_model: A generic TFLite model to be converted. + :param tflite_model: A generic TFLite model to be compiled. :param target: The target platform. :param delegation_tag: The delegation tag of model partition. :param fetch_constants_to_sram: Add microcode that fetches weights from external memory. @@ -104,7 +123,7 @@ def convert( :return: TFLite model with Neutron microcode as bytes. """ - # Neutron converter crashes if we provide invalid target -> verify. + # Neutron compiler crashes if we provide invalid target -> verify. self.verify_target(target) compilation_opts = { @@ -124,7 +143,7 @@ def convert( queue = multiprocessing.Manager().Queue() process = multiprocessing.Process( - target=convert_unsafe, + target=compile_unsafe, args=(tflite_model, compilation_opts, queue), ) process.start() @@ -132,20 +151,23 @@ def convert( if queue.empty(): # signals the unsafe task did not run till the end raise RuntimeError( - f"Neutron converter module terminated unexpectedly with exit code {process.exitcode}" + f"Neutron compiler module terminated unexpectedly with exit code {process.exitcode}" ) - model_converted = queue.get() + model_compiled = queue.get() process.close() except (EOFError, OSError, TypeError) as e: # Multiprocessing failed (likely due to environment restrictions) # Fall back to direct execution logging.warning( - f"Multiprocessing not available ({e}), running neutron converter directly" + f"Multiprocessing not available ({e}), running neutron compiler directly" ) cctx = _build_compilation_context(compilation_opts) - model_converted = neutron_converter.convertModel(list(tflite_model), cctx) + if _USING_NEUTRON_COMPILER: + model_compiled = neutron_compiler.compileModel(list(tflite_model), cctx) + else: + model_compiled = neutron_compiler.convertModel(list(tflite_model), cctx) if self.dump_kernel_selection_code: self._rename_partition_kernel_selection_file(delegation_tag) - return bytes(model_converted) + return bytes(model_compiled) diff --git a/backends/nxp/backend/neutron_map.py b/backends/nxp/backend/neutron_map.py index da497565726..a5becafc4f0 100644 --- a/backends/nxp/backend/neutron_map.py +++ b/backends/nxp/backend/neutron_map.py @@ -91,15 +91,15 @@ def get_tensors_name(tensors: str) -> list[str]: class NeutronMap: - """Mapping between Neutron, TFLite, and Edge operators based on the Neutron converter log. + """Mapping between Neutron, TFLite, and Edge operators based on the Neutron compiler log. - Parses the Neutron converter log to extract information about TFLite nodes and Neutron subgraphs. + Parses the Neutron compiler log to extract information about TFLite nodes and Neutron subgraphs. Maps TFLite operators to corresponding Neutron operators. Maps Edge operators to Neutron operators via the Edge-to-TFLite mapping. Attributes: - tflite_nodes (list[Node]): TFLite node information extracted from the converter log. - neutron_subgraphs (list[SubgraphInfo]): Neutron subgraph information extracted from the converter log. + tflite_nodes (list[Node]): TFLite node information extracted from the compiler log. + neutron_subgraphs (list[SubgraphInfo]): Neutron subgraph information extracted from the compiler log. neutron_graphs (list[int]): Indices of final Neutron graphs derived from neutron_subgraphs. edge_to_tflite_map (dict[int, tuple[int, ...]]): Mapping from Edge operators to TFLite operators. edge_to_neutron_map (dict[int, tuple[int, ...]]): Mapping from Edge operators to Neutron operators. @@ -118,12 +118,12 @@ class NeutronMap: tflite_to_neutron_map: dict[int, tuple[int, ...]] def __init__( - self, neutron_converter_log: str, edge_to_tflite_map: dict[int, tuple[int, ...]] + self, neutron_compiler_log: str, edge_to_tflite_map: dict[int, tuple[int, ...]] ) -> None: - """Initialize neutron map from neutron converter log. + """Initialize neutron map from neutron compiler log. - :param neutron_converter_log: neutron converter log obtained during model conversion. It should contain - original tflite graph and neutron graph dump. To add these dumps to converter log the dumpAfterImport and + :param neutron_compiler_log: neutron compiler log obtained during model compilation. It should contain + original tflite graph and neutron graph dump. To add these dumps to compiler log the dumpAfterImport and dumpAfterGenerate flags have to be set to "console". """ super().__init__() @@ -134,12 +134,12 @@ def __init__( self.tflite_to_neutron_map = {} self.edge_to_neutron_map = {} self.neutron_kernels_num = 0 - self._split_profiling_log(neutron_converter_log) + self._split_profiling_log(neutron_compiler_log) def _split_profiling_log(self, log: str) -> None: """Process profiling log to split it into original TFLite and converted Neutron nodes. - :param log: Neutron converter log obtained during model conversion, containing the original + :param log: Neutron compiler log obtained during model compilation, containing the original TFLite graph and Neutron graph dump. :return: None. Sets class attributes tflite_nodes and neutron_subgraphs with node information. """ @@ -175,7 +175,7 @@ def _split_profiling_log(self, log: str) -> None: def _get_neutron_subgraphs(self, graph_dump: str) -> list[SubgraphInfo]: """Parse Neutron graph dump and extract subgraph information. - :param graph_dump: String containing the Neutron graph dump from the converter log. + :param graph_dump: String containing the Neutron graph dump from the compiler log. :return: List of SubgraphInfo objects containing subgraph metadata and operator nodes. """ diff --git a/backends/nxp/backend/neutron_target_spec.py b/backends/nxp/backend/neutron_target_spec.py index 5a75caf9a75..a51d437f060 100644 --- a/backends/nxp/backend/neutron_target_spec.py +++ b/backends/nxp/backend/neutron_target_spec.py @@ -8,8 +8,8 @@ from enum import Enum import torch -from executorch.backends.nxp.backend.neutron_converter_manager import ( - NeutronConverterManager, +from executorch.backends.nxp.backend.neutron_compiler_manager import ( + NeutronCompilerManager, ) from executorch.exir.dialects._ops import ops as exir_ops from torch.fx import Node @@ -98,10 +98,10 @@ class NeutronTargetSpec: def __init__(self, target: str): - converter_manager = NeutronConverterManager() - converter_manager.verify_target(target) - neutron_converter = converter_manager.get_converter() - self.neutron_target = neutron_converter.getNeutronTarget(target) + compiler_manager = NeutronCompilerManager() + compiler_manager.verify_target(target) + neutron_compiler = compiler_manager.get_compiler() + self.neutron_target = neutron_compiler.getNeutronTarget(target) if self.is_subsystem(): raise ValueError( diff --git a/backends/nxp/backend/node_format_inference.py b/backends/nxp/backend/node_format_inference.py index 689de41f3e4..64595eb411f 100644 --- a/backends/nxp/backend/node_format_inference.py +++ b/backends/nxp/backend/node_format_inference.py @@ -15,7 +15,7 @@ try_get_arg, ) from executorch.backends.nxp.backend.edge_program_converter import functions_converters -from executorch.backends.nxp.tests.ops_aliases import ( +from executorch.backends.nxp.backend.ops_aliases import ( AdaptiveAvgPool2D, Amax, Amin, diff --git a/backends/nxp/tests/ops_aliases.py b/backends/nxp/backend/ops_aliases.py similarity index 77% rename from backends/nxp/tests/ops_aliases.py rename to backends/nxp/backend/ops_aliases.py index 5ccc1d67de0..fea772eb8ee 100644 --- a/backends/nxp/tests/ops_aliases.py +++ b/backends/nxp/backend/ops_aliases.py @@ -3,8 +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. -# This file defines ops aliases for shorter and more readable test description. List is sorted alphabetically. -# When finding a missing alias, add it at the correct place. +# This file defines ops aliases for shorter and more readable descriptions of edge operators. +# List is sorted alphabetically. When finding a missing alias, add it at the correct place. import operator @@ -17,6 +17,7 @@ AddTensor = exir_ops.edge.aten.add.Tensor Amax = exir_ops.edge.aten.amax.default Amin = exir_ops.edge.aten.amin.default +AsStridedCopy = exir_ops.edge.aten.as_strided_copy.default AvgPool2D = exir_ops.edge.aten.avg_pool2d.default BMM = exir_ops.edge.aten.bmm.default Cat = exir_ops.edge.aten.cat.default @@ -27,6 +28,9 @@ 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 +DequantizePerTensorTensor = ( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor +) ExecutorchDelegateCall = torch.ops.higher_order.executorch_call_delegate Exp = exir_ops.edge.aten.exp.default GetItem = operator.getitem @@ -35,6 +39,7 @@ HardTanh = exir_ops.edge.aten.hardtanh.default HardTanh_ = exir_ops.edge.aten.hardtanh_.default LeakyRelu = exir_ops.edge.aten.leaky_relu.default +Linear = exir_ops.edge.aten.linear.default Log = exir_ops.edge.aten.log.default MM = exir_ops.edge.aten.mm.default Maximum = exir_ops.edge.aten.maximum.default @@ -43,12 +48,17 @@ MeanDim = exir_ops.edge.aten.mean.dim Minimum = exir_ops.edge.aten.minimum.default MulTensor = exir_ops.edge.aten.mul.Tensor +NativebatchNormLegitNoStats = exir_ops.edge.aten._native_batch_norm_legit.no_stats +NativebatchNormLegitNoTraining = ( + exir_ops.edge.aten._native_batch_norm_legit_no_training.default +) Neg = exir_ops.edge.aten.neg.default Pad = exir_ops.edge.aten.pad.default PermuteCopy = exir_ops.edge.aten.permute_copy.default Prelu = exir_ops.edge.aten.prelu.default QuantizePerChannel = exir_ops.edge.quantized_decomposed.quantize_per_channel.default QuantizePerTensor = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default +QuantizePerTensorTensor = exir_ops.edge.quantized_decomposed.quantize_per_tensor.tensor Relu = exir_ops.edge.aten.relu.default Rsqrt = exir_ops.edge.aten.rsqrt.default Sigmoid = exir_ops.edge.aten.sigmoid.default @@ -56,6 +66,9 @@ SliceCopy = exir_ops.edge.aten.slice_copy.Tensor Softmax = exir_ops.edge.aten._softmax.default Squeeze = exir_ops.edge.aten.squeeze.default +SqueezeCopy = exir_ops.edge.aten.squeeze_copy.default +SqueezeCopyDim = exir_ops.edge.aten.squeeze_copy.dim +SqueezeCopyDims = exir_ops.edge.aten.squeeze_copy.dims SqueezeDim = exir_ops.edge.aten.squeeze.dim SqueezeDims = exir_ops.edge.aten.squeeze.dims SubTensor = exir_ops.edge.aten.sub.Tensor @@ -63,6 +76,7 @@ Tanh = exir_ops.edge.aten.tanh.default Tanh_ = exir_ops.edge.aten.tanh_.default Unsqueeze = exir_ops.edge.aten.unsqueeze.default +UnsqueezeCopy = exir_ops.edge.aten.unsqueeze_copy.default UpsampleBilinear2D = exir_ops.edge.aten.upsample_bilinear2d.vec UpsampleNearest2D = exir_ops.edge.aten.upsample_nearest2d.vec ViewCopy = exir_ops.edge.aten.view_copy.default diff --git a/backends/nxp/edge_passes/convert_reshaping_nodes_to_view.py b/backends/nxp/edge_passes/convert_reshaping_nodes_to_view.py index a82c823f593..ad296b165b2 100644 --- a/backends/nxp/edge_passes/convert_reshaping_nodes_to_view.py +++ b/backends/nxp/edge_passes/convert_reshaping_nodes_to_view.py @@ -6,8 +6,14 @@ import torch +from executorch.backends.nxp.backend.ops_aliases import ( + SqueezeCopy, + SqueezeCopyDim, + SqueezeCopyDims, + UnsqueezeCopy, + ViewCopy, +) from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass -from executorch.exir.dialects._ops import ops as exir_ops from torch._subclasses import FakeTensor, FakeTensorMode from torch.fx import GraphModule, Node from torch.fx.passes.infra.pass_base import PassResult @@ -43,25 +49,20 @@ class ConvertReshapingNodesToViewPass(NeutronEdgePass): @staticmethod def _is_squeeze(node_: Node) -> bool: return node_.op == "call_function" and ( - node_.target == exir_ops.edge.aten.squeeze_copy.dim - or node_.target == exir_ops.edge.aten.squeeze_copy.dims - or node_.target == exir_ops.edge.aten.squeeze_copy.default + node_.target == SqueezeCopyDim + or node_.target == SqueezeCopyDims + or node_.target == SqueezeCopy ) @staticmethod def _is_unsqueeze(node_: Node) -> bool: - return ( - node_.op == "call_function" - and node_.target == exir_ops.edge.aten.unsqueeze_copy.default - ) + return node_.op == "call_function" and node_.target == UnsqueezeCopy def _create_view_copy_node(self, *view_args) -> Node: - view_target = exir_ops.edge.aten.view_copy.default + view_target = ViewCopy view_node = self.graph_module.graph.call_function(view_target, view_args) - view_node.meta["source_fn_stack"] = [ - (view_node.name, exir_ops.edge.aten.view_copy.default) - ] + view_node.meta["source_fn_stack"] = [(view_node.name, ViewCopy)] x_val = view_args[0].meta["val"] with FakeTensorMode() as mode: diff --git a/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py b/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py index dafea5d259b..dc9cc28a25e 100644 --- a/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py +++ b/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py @@ -3,36 +3,35 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import operator - import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AdaptiveAvgPool2D, + AddMM, + AvgPool2D, + Clone, + CloneDimOrder, + Convolution, + DequantizePerTensor, + GetItem, + HardTanh, + MaxPool2DWithIndices, + MM, + PermuteCopy, + QuantizePerTensor, + Relu, + Sigmoid, + SqueezeCopyDims, + Tanh, + UnsqueezeCopy, + ViewCopy, +) + from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass from executorch.backends.nxp.neutron_partitioner import QDQClusterRecognizer -from executorch.backends.nxp.tests.ops_aliases import PermuteCopy - -# noinspection PyProtectedMember -from executorch.exir.dialects._ops import ops as exir_ops from torch.fx import Node from torch.fx.passes.infra.pass_base import PassResult -# Operator aliases for better readability. -AddMM = exir_ops.edge.aten.addmm.default -AvgPool2D = exir_ops.edge.aten.avg_pool2d.default -MaxPool2D = exir_ops.edge.aten.max_pool2d_with_indices.default -Conv = exir_ops.edge.aten.convolution.default -Clone = exir_ops.edge.aten.clone.default -CloneDimOrder = exir_ops.edge.dim_order_ops._clone_dim_order.default -Getitem = operator.getitem -HardTanh = exir_ops.edge.aten.hardtanh.default -MM = exir_ops.edge.aten.mm.default -Relu = exir_ops.edge.aten.relu.default -Sigmoid = exir_ops.edge.aten.sigmoid.default -SqueezeCopy = exir_ops.edge.aten.squeeze_copy.dims -Tanh = exir_ops.edge.aten.tanh.default -UnsqueezeCopy = exir_ops.edge.aten.unsqueeze_copy.default -ViewCopy = exir_ops.edge.aten.view_copy.default - def insert_qdq_pair_after_node( graph: torch.fx.Graph, anchor: torch.fx.Node, q_params: tuple @@ -41,7 +40,7 @@ def insert_qdq_pair_after_node( with graph.inserting_after(anchor): quantize_op = graph.create_node( op="call_function", - target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + target=QuantizePerTensor, args=(), # Will be added later. ) quantize_op.meta = anchor.meta @@ -50,7 +49,7 @@ def insert_qdq_pair_after_node( with graph.inserting_after(quantize_op): dequantize_op = graph.create_node( op="call_function", - target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + target=DequantizePerTensor, args=(quantize_op,) + q_params, ) dequantize_op.meta = quantize_op.meta @@ -65,8 +64,7 @@ def _is_dequantize(node_: Node) -> bool: return ( hasattr(node_, "op") and node_.op == "call_function" - and node_.target - == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default + and node_.target == DequantizePerTensor ) @@ -74,8 +72,7 @@ def _is_quantize(node_: Node) -> bool: return ( hasattr(node_, "op") and node_.op == "call_function" - and node_.target - == exir_ops.edge.quantized_decomposed.quantize_per_tensor.default + and node_.target == QuantizePerTensor ) @@ -117,7 +114,7 @@ class MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass): PermuteCopy, ], ViewCopy: [Clone, CloneDimOrder], - Conv: [ + Convolution: [ ViewCopy, # For 1D conv ], # AvgPool1D is represented in edge as Unsqueeze -> AvgPool2D -> Squeeze. The reshaping nodes must be moved out @@ -126,9 +123,15 @@ class MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass): ViewCopy, UnsqueezeCopy, ], - # MaxPool1D is represented in edge as Unsqueeze -> MaxPool2D -> Getitem -> Squeeze. The reshaping nodes must be moved out + # MaxPool1D is represented in edge as Unsqueeze -> MaxPool2D -> GetItem -> Squeeze. The reshaping nodes must be moved out # of the cluster. Instead of [Un]squeeze, ViewCopy can be used as well. - MaxPool2D: [ + MaxPool2DWithIndices: [ + ViewCopy, + UnsqueezeCopy, + ], + # AdaptiveAvgPool1D is represented in edge as Unsqueeze -> AdaptiveAvgPool2D -> Squeeze. The reshaping nodes + # must be moved out of the cluster. Instead of [Un]squeeze, ViewCopy can be used as well. + AdaptiveAvgPool2D: [ ViewCopy, UnsqueezeCopy, ], @@ -222,7 +225,7 @@ class MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass): Sigmoid, Tanh, ], - Conv: [ + Convolution: [ HardTanh, Relu, Sigmoid, @@ -234,13 +237,19 @@ class MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass): # of the cluster. Instead of [Un]squeeze, ViewCopy can be used as well. AvgPool2D: [ ViewCopy, - SqueezeCopy, + SqueezeCopyDims, ], - # MaxPool1D is represented in edge as Unsqueeze -> MaxPool2D -> Getitem -> Squeeze. The reshaping nodes must be moved out + # MaxPool1D is represented in edge as Unsqueeze -> MaxPool2D -> GetItem -> Squeeze. The reshaping nodes must be moved out # of the cluster. Instead of [Un]squeeze, ViewCopy can be used as well. - Getitem: [ + GetItem: [ + ViewCopy, + SqueezeCopyDims, + ], + # AdaptiveAvgPool1D is represented in edge as Unsqueeze -> AdaptiveAvgPool2D -> Squeeze. The reshaping nodes + # must be moved out of the cluster. Instead of [Un]squeeze, ViewCopy can be used as well. + AdaptiveAvgPool2D: [ ViewCopy, - SqueezeCopy, + SqueezeCopyDims, ], } @@ -278,7 +287,7 @@ def run(self, graph_module: torch.fx.GraphModule) -> PassResult: # satisfy the requirements of the `QDQClusterRecognizer`. actual_main_cluster_node = ( main_cluster_node - if main_cluster_node.target != Getitem + if main_cluster_node.target != GetItem else main_cluster_node.args[0] ) cluster = QDQClusterRecognizer().get_qdq_cluster(actual_main_cluster_node) diff --git a/backends/nxp/edge_passes/neutron_edge_pass_manager.py b/backends/nxp/edge_passes/neutron_edge_pass_manager.py index 3a7fc3ffbfa..e95f8d18144 100644 --- a/backends/nxp/edge_passes/neutron_edge_pass_manager.py +++ b/backends/nxp/edge_passes/neutron_edge_pass_manager.py @@ -17,7 +17,7 @@ from executorch.backends.nxp.edge_passes.remove_as_strided_copy_nodes import ( RemoveUselessAsStridedCopyNodes, ) -from torch.fx.passes.infra.pass_manager import PassManager +from executorch.exir.pass_manager import PassManager class NeutronEdgePassManager(PassManager): diff --git a/backends/nxp/edge_passes/remove_additional_quantize_dequantize_nodes_pass.py b/backends/nxp/edge_passes/remove_additional_quantize_dequantize_nodes_pass.py index 4edcc0b0e97..549b77f6eca 100644 --- a/backends/nxp/edge_passes/remove_additional_quantize_dequantize_nodes_pass.py +++ b/backends/nxp/edge_passes/remove_additional_quantize_dequantize_nodes_pass.py @@ -7,9 +7,18 @@ import torch from executorch.backends.nxp.backend.edge_helper import get_quantization_parameters_for +from executorch.backends.nxp.backend.ops_aliases import ( + Cat, + DequantizePerChannel, + DequantizePerTensor, + DequantizePerTensorTensor, + PermuteCopy, + QuantizePerChannel, + QuantizePerTensor, + QuantizePerTensorTensor, +) from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass from executorch.backends.nxp.neutron_partitioner import QDQClusterRecognizer -from executorch.exir.dialects._ops import ops as exir_ops from torch.fx.passes.infra.pass_base import PassResult @@ -36,15 +45,15 @@ class RemoveAdditionalQDQClustersPass(NeutronEdgePass): """ qdq_per_channel_nodes = ( - exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, - exir_ops.edge.quantized_decomposed.quantize_per_channel.default, + DequantizePerChannel, + QuantizePerChannel, ) qdq_per_tensor_nodes = ( - exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, - exir_ops.edge.quantized_decomposed.quantize_per_tensor.tensor, - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor, + QuantizePerTensor, + QuantizePerTensorTensor, + DequantizePerTensor, + DequantizePerTensorTensor, ) def run(self, graph_module: torch.fx.GraphModule) -> PassResult: @@ -55,8 +64,8 @@ def run(self, graph_module: torch.fx.GraphModule) -> PassResult: for cluster in qdq_clusterer.cluster_map.values(): # For now, enable only permute_copy and cat. if cluster.compute_node.target not in [ - exir_ops.edge.aten.permute_copy.default, - exir_ops.edge.aten.cat.default, + PermuteCopy, + Cat, ]: continue diff --git a/backends/nxp/edge_passes/remove_as_strided_copy_nodes.py b/backends/nxp/edge_passes/remove_as_strided_copy_nodes.py index 1ed9a9f607c..d2257a453ac 100644 --- a/backends/nxp/edge_passes/remove_as_strided_copy_nodes.py +++ b/backends/nxp/edge_passes/remove_as_strided_copy_nodes.py @@ -4,8 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from executorch.backends.nxp.backend.ops_aliases import AsStridedCopy, MeanDim from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass -from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.passes import dead_code_elimination_pass from torch.fx import GraphModule from torch.fx.passes.infra.pass_base import PassResult @@ -18,12 +18,12 @@ def __init__(self): def gen_pattern_as_strided_copy(self, graph_module: GraphModule): # Unedited method taken from `backends/samsung/_passes/remove_useless_ops.py`. for node in list(graph_module.graph.nodes): # noqa: C416 - if node.target != exir_ops.edge.aten.mean.dim: + if node.target != MeanDim: continue if len(node.users) != 1: continue successor = list(node.users.keys())[0] - if successor.target != exir_ops.edge.aten.as_strided_copy.default: + if successor.target != AsStridedCopy: continue is_pattern = True count = 0 diff --git a/backends/nxp/edge_passes/remove_io_quant_ops_pass.py b/backends/nxp/edge_passes/remove_io_quant_ops_pass.py index a87eac7360c..bb2ae14c397 100644 --- a/backends/nxp/edge_passes/remove_io_quant_ops_pass.py +++ b/backends/nxp/edge_passes/remove_io_quant_ops_pass.py @@ -5,8 +5,11 @@ import torch +from executorch.backends.nxp.backend.ops_aliases import ( + DequantizePerTensor, + QuantizePerTensor, +) from executorch.exir import EdgeProgramManager -from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass from executorch.exir.passes.quantize_io_pass import QuantizeInputs, QuantizeOutputs from torch.fx.passes.infra.pass_base import PassResult @@ -37,10 +40,7 @@ def _get_quantizable_input_indices(self): raise ValueError(f"Input {input_index} has more than one users") quantize = next(iter(target_placeholder.users)) - if ( - quantize.target - != exir_ops.edge.quantized_decomposed.quantize_per_tensor.default - ): + if quantize.target != QuantizePerTensor: continue inputs_to_quantization.append(input_index) @@ -59,10 +59,7 @@ def _get_quantizable_output_indices(self): user_outputs = list(outputs[0].args[0]) for output_index, user_output in enumerate(user_outputs): - if ( - user_output.target - != exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default - ): + if user_output.target != DequantizePerTensor: continue outputs_to_quantization.append(output_index) diff --git a/backends/nxp/nxp_backend.py b/backends/nxp/nxp_backend.py index 2f4bb07316f..0d4324d57c2 100644 --- a/backends/nxp/nxp_backend.py +++ b/backends/nxp/nxp_backend.py @@ -25,8 +25,8 @@ EdgeProgramToIRConverter, ) from executorch.backends.nxp.backend.ir.conversion_config import ConversionConfig -from executorch.backends.nxp.backend.neutron_converter_manager import ( - NeutronConverterManager, +from executorch.backends.nxp.backend.neutron_compiler_manager import ( + NeutronCompilerManager, ) from executorch.backends.nxp.backend.neutron_map import NeutronMap @@ -45,6 +45,14 @@ torch.ops.aten.prelu.default, ] +# Aten operators that must be preserved (not decomposed) during lowering to the edge dialect, because the Neutron +# backend can handle them natively. +default_preserve_ops = [ + torch.ops.aten.hardswish.default, + torch.ops.aten.pad.default, + torch.ops.aten.prelu.default, +] + class NeutronCompileSpecBuilder: config: NeutronTargetSpec @@ -86,10 +94,10 @@ def neutron_compile_spec( :param use_neutron_for_format_conversion: If True, the EdgeProgramToIRConverter will insert `Transpose` ops to ensure that the IO matches the executorch partition, which will be delegated to Neutron. - :param fetch_constants_to_sram: If True, the Neutron Converter will insert microinstructions to prefetch weights + :param fetch_constants_to_sram: If True, the Neutron Compiler will insert microinstructions to prefetch weights from FLASH to SRAM. This should be used when the whole model does not fit into SRAM. - :param dump_kernel_selection_code: Whether Neutron converter dumps kernel selection code. - :param use_profiling: If true Neutron Converter will enable profiling for neutron delegated model + :param dump_kernel_selection_code: Whether Neutron Compiler dumps kernel selection code. + :param use_profiling: If true Neutron Compiler will enable profiling for neutron delegated model :return: self for method chaining """ @@ -282,9 +290,9 @@ def preprocess( # noqa C901 ) with capture_fd_output() as tmp: - neutron_model = NeutronConverterManager( + neutron_model = NeutronCompilerManager( dump_kernel_selection_code - ).convert( + ).compile( tflite_model, target, delegation_tag, diff --git a/backends/nxp/quantizer/neutron_quantizer.py b/backends/nxp/quantizer/neutron_quantizer.py index 1f1d26ce229..7f186b5f769 100644 --- a/backends/nxp/quantizer/neutron_quantizer.py +++ b/backends/nxp/quantizer/neutron_quantizer.py @@ -13,7 +13,8 @@ from executorch.backends.nxp.quantizer.patterns import ( AbsPattern, ActivationsConcatClusterPattern, - AdaptiveAvgPoolPattern, + AdaptiveAvgPool1DPattern, + AdaptiveAvgPool2DPattern, AddmmPattern, AddTensorPattern, AmaxPattern, @@ -267,7 +268,8 @@ def __init__(self, neutron_target_spec: NeutronTargetSpec, is_qat: bool = False) super().__init__( [ OpQuantizer(AbsPattern(is_qat=is_qat), static_qconfig), - OpQuantizer(AdaptiveAvgPoolPattern(is_qat=is_qat), static_qconfig), + OpQuantizer(AdaptiveAvgPool1DPattern(is_qat=is_qat), static_qconfig), + OpQuantizer(AdaptiveAvgPool2DPattern(is_qat=is_qat), static_qconfig), OpQuantizer(AddTensorPattern(is_qat=is_qat), static_qconfig), OpQuantizer(AddmmPattern(self, is_qat=is_qat), static_fc_qconfig), OpQuantizer(AmaxPattern(is_qat=is_qat), static_qconfig), diff --git a/backends/nxp/quantizer/patterns.py b/backends/nxp/quantizer/patterns.py index 3bbd2bed54e..4ac4b777cba 100644 --- a/backends/nxp/quantizer/patterns.py +++ b/backends/nxp/quantizer/patterns.py @@ -279,7 +279,16 @@ def partition_types(self): return [torch.ops.aten.abs.default] -class AdaptiveAvgPoolPattern(SharedSpecPattern): +class AdaptiveAvgPool1DPattern(SharedSpecPattern): + """ + Quantizer for AdaptiveAvgPool1D operator. + """ + + def partition_types(self): + return [torch.ops.aten.adaptive_avg_pool1d.default] + + +class AdaptiveAvgPool2DPattern(SharedSpecPattern): """ Quantizer for AdaptiveAvgPool2D operator. """ diff --git a/backends/nxp/recipes/__init__.py b/backends/nxp/recipes/__init__.py new file mode 100644 index 00000000000..f768d4fa88d --- /dev/null +++ b/backends/nxp/recipes/__init__.py @@ -0,0 +1,18 @@ +# Copyright 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. + +from executorch.export import recipe_registry + +from .nxp_recipe_provider import NeutronRecipeConfig, NXPRecipeProvider +from .nxp_recipe_types import NXPRecipeType + +# Auto-register NXP recipe provider +recipe_registry.register_backend_recipe_provider(NXPRecipeProvider()) + +__all__ = [ + "NeutronRecipeConfig", + "NXPRecipeProvider", + "NXPRecipeType", +] diff --git a/backends/nxp/recipes/nxp_recipe_provider.py b/backends/nxp/recipes/nxp_recipe_provider.py new file mode 100644 index 00000000000..6592f1b0dd9 --- /dev/null +++ b/backends/nxp/recipes/nxp_recipe_provider.py @@ -0,0 +1,420 @@ +# Copyright 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 logging +from copy import deepcopy +from dataclasses import dataclass +from functools import partial +from typing import Any, Callable, cast, Iterable, Optional, Sequence + +import torch + +from executorch.backends.nxp.aten_passes.fuse_batch_norm_with_linear_pass import ( + FuseBatchNormWithLinearPass, +) +from executorch.backends.nxp.aten_passes.simulated_linear_bn_fusion_passes import ( + AddSimulatedLinearBatchNormFusionQATPass, + RemoveSimulatedLinearBatchNormFusionQATPass, +) +from executorch.backends.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) +from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec +from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass +from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import ( + NeutronEdgePassManager, +) +from executorch.backends.nxp.edge_passes.remove_additional_quantize_dequantize_nodes_pass import ( + RemoveAdditionalQDQClustersPass, +) +from executorch.backends.nxp.edge_passes.remove_io_quant_ops_pass import ( + RemoveIOQuantOpsPass, +) +from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner +from executorch.backends.nxp.nxp_backend import ( + core_aten_ops_exception_list, + default_preserve_ops, + generate_neutron_compile_spec, +) +from executorch.backends.nxp.quantizer.utils import ( + _replace_histogram_observers_for_integer_inputs, +) +from executorch.backends.nxp.recipes.nxp_recipe_types import NXP_BACKEND, NXPRecipeType +from executorch.backends.nxp.tests.executorch_pipeline import ( + get_default_quantizer, + handle_kernel_selection, + ModelInputSpec, + to_model_input_spec, +) +from executorch.backends.transforms.quantize_fused_convbn_bias_pass import ( + QuantizeFusedConvBnBiasAtenPass, +) +from executorch.exir import ( + EdgeCompileConfig, + EdgeProgramManager, + ExecutorchBackendConfig, + ExportedProgram, +) + +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.backend.partitioner import Partitioner +from executorch.export import ( + BackendRecipeProvider, + ExportRecipe, + LoweringRecipe, + QuantizationRecipe, + RecipeType, +) +from torchao.quantization.pt2e.quantizer import Quantizer + + +class NeutronEdgePassManagerWrapper: + def __init__(self, passes: list[NeutronEdgePass] | None = None): + self.neutron_edge_pass_manager = NeutronEdgePassManager(passes) + + def __call__( + self, method_name: str, exported_program: ExportedProgram + ) -> NeutronEdgePassManager: + return self.neutron_edge_pass_manager + + +NEUTRON_RECIPE_CONFIG_KEY = "neutron_recipe_config" + + +@dataclass +class NeutronRecipeConfig: + """Configuration shared by all NXP recipe types. + + Parameters that vary the *type* of export (delegate vs no-delegate, PTQ vs QAT) + are expressed by choosing a different NXPRecipeType rather than by flags here. + + Attributes: + input_spec: Model input description. Accepts a single shape tuple, a list of + shape tuples (one per input), or a list of ModelInputSpec objects. + target: Neutron hardware target string. Default: "imxrt700". + operators_not_to_delegate: Optional list of op names excluded from NPU delegation. + For example ["aten::convolution"]. + intermediates_dir: Optional directory to dump intermediate compilation artifacts. + get_quantizer_fn: Optional factory that returns a custom Quantizer. When None, + the default NeutronQuantizer is used. + custom_delegation_options: Optional fine-grained control over which ops are + delegated. Default: CustomDelegationOptions(). + remove_quant_io_ops: If True, remove quantize/dequantize ops at the IO boundary + (useful for integer-IO deployments). + use_quant_state_dict: If False, the post-quantization parameter values are not + passed to NeutronPartitioner. + use_neutron_for_format_conversion: Whether Neutron handles data-format conversion. + fetch_constants_to_sram: Place constant tensors in SRAM on the target. + dump_kernel_selection_code: Generate kernel-selection files after compilation. + use_profiling: Enable Neutron execution profiling. IMPORTANT: To also generate an + ETRecord, pass generate_etrecord=True to export() separately. + train_fn: Training function required for QAT recipe types (INT8_QAT_NEUTRON and + INT8_QAT_NO_DELEGATE). Receives the prepared GraphModule and must + perform the training loop. Ignored for PTQ recipe types. + """ + + input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]] + target: str = "imxrt700" + operators_not_to_delegate: list[str] | None = None + intermediates_dir: str | None = None + get_quantizer_fn: Callable[[], Quantizer] | None = None + custom_delegation_options: CustomDelegationOptions | None = None + remove_quant_io_ops: bool = False + use_quant_state_dict: bool = True + use_neutron_for_format_conversion: bool = True + fetch_constants_to_sram: bool = False + dump_kernel_selection_code: bool = False + use_profiling: bool = False + train_fn: Callable[["torch.fx.GraphModule"], None] | None = None + + +class NXPRecipeProvider(BackendRecipeProvider): + + @property + def backend_name(self) -> str: + return NXP_BACKEND + + def get_supported_recipes(self) -> Sequence[RecipeType]: + return list(NXPRecipeType) + + def create_recipe( + self, recipe_type: RecipeType, **kwargs: Any + ) -> Optional[ExportRecipe]: + if recipe_type not in self.get_supported_recipes(): + logging.warning(f"NXP backend: Recipe `{recipe_type}` is not valid.") + return None + + original_rc = kwargs.get(NEUTRON_RECIPE_CONFIG_KEY) + if original_rc is None: + raise KeyError( + f"NXP backend: create_recipe() requires `{NEUTRON_RECIPE_CONFIG_KEY}=`." + ) + if not isinstance(original_rc, NeutronRecipeConfig): + raise TypeError( + f"NXP backend: `{NEUTRON_RECIPE_CONFIG_KEY}` must be a NeutronRecipeConfig, " + f"got {type(original_rc).__name__}." + ) + + rc = cast(NeutronRecipeConfig, deepcopy(original_rc)) + if rc.custom_delegation_options is None: + rc.custom_delegation_options = CustomDelegationOptions() + + rc.input_spec = to_model_input_spec(rc.input_spec) + + match recipe_type: + case NXPRecipeType.INT8_PTQ_NEUTRON: + return self._build_recipe(recipe_type, rc, is_qat=False, delegate=True) + case NXPRecipeType.INT8_PTQ_NO_DELEGATE: + return self._build_recipe(recipe_type, rc, is_qat=False, delegate=False) + case NXPRecipeType.INT8_QAT_NEUTRON: + return self._build_recipe(recipe_type, rc, is_qat=True, delegate=True) + case NXPRecipeType.INT8_QAT_NO_DELEGATE: + return self._build_recipe(recipe_type, rc, is_qat=True, delegate=False) + case _: + raise NotImplementedError( + f"NXP backend: Recipe `{recipe_type}` is not supported." + ) + + def _build_recipe( + self, + recipe_type: NXPRecipeType, + rc: NeutronRecipeConfig, + *, + is_qat: bool, + delegate: bool, + ) -> ExportRecipe: + if is_qat and rc.train_fn is None: + raise ValueError( + f"NXP backend: Recipe `{recipe_type}` requires `train_fn` to be set in " + f"NeutronRecipeConfig. Provide a callable that trains the prepared model." + ) + + neutron_target_spec = NeutronTargetSpec(rc.target) + + if rc.get_quantizer_fn is None: + rc.get_quantizer_fn = partial( + get_default_quantizer, neutron_target_spec, is_qat + ) + + quantization_recipe = _build_quantization_recipe(rc, is_qat) + compile_spec = generate_neutron_compile_spec( + rc.target, + intermediates_dir=rc.intermediates_dir, + operators_not_to_delegate=rc.operators_not_to_delegate, + use_neutron_for_format_conversion=rc.use_neutron_for_format_conversion, + fetch_constants_to_sram=rc.fetch_constants_to_sram, + dump_kernel_selection_code=rc.dump_kernel_selection_code, + use_profiling=rc.use_profiling, + ) + lowering_recipe = _build_lowering_recipe( + compile_spec, neutron_target_spec, rc, delegate=delegate + ) + + return ExportRecipe( + name=recipe_type.value, + quantization_recipe=quantization_recipe, + lowering_recipe=lowering_recipe, + executorch_backend_config=ExecutorchBackendConfig( + extract_delegate_segments=False + ), + ) + + +# --------------------------------------------------------------------------- +# Pass wrappers +# --------------------------------------------------------------------------- +# ExirPassBase subclasses return a PassResult with a .graph_module attribute, +# but QuantizeStage._apply_passes expects callable(GraphModule) -> GraphModule. +# These thin wrappers bridge the two conventions. + + +def _wrap_exir_pass(pass_cls, *args, **kwargs): + """Return a callable(GraphModule) -> GraphModule wrapping an ExirPass instance.""" + _pass_instance = pass_cls(*args, **kwargs) + + def _wrapped(m): + return _pass_instance(m).graph_module + + _wrapped.__qualname__ = f"_wrap_exir_pass({pass_cls.__name__})" + return _wrapped + + +def _histogram_observer_fix_pass(m): + """Callable(GraphModule) -> GraphModule that replaces HistogramObserver for integer inputs.""" + _replace_histogram_observers_for_integer_inputs(m) + return m + + +# --------------------------------------------------------------------------- +# Module-level builder helpers +# --------------------------------------------------------------------------- + + +def _build_quantization_recipe( + rc: NeutronRecipeConfig, is_qat: bool +) -> QuantizationRecipe: + """Build the QuantizationRecipe for PTQ or QAT. + + PTQ uses the standard QuantizeStage flow (prepare_pt2e -> calibrate -> convert_pt2e). + QAT uses the QAT flow (prepare_qat_pt2e -> BN-fusion passes -> train_fn -> convert_pt2e). + + The NXP-specific passes are injected via the QuantizationRecipe hook lists so + that QuantizeStage executes them in the correct order. + """ + _quantizer = rc.get_quantizer_fn() + + # post_prepare_passes: always fix HistogramObserver for non-float inputs. + # For QAT, also insert the simulated linear-BN fusion before training so + # fake-quantize nodes see fused weights during the training loop. + post_prepare: list[Callable] = [] + if is_qat: + post_prepare.append(_wrap_exir_pass(AddSimulatedLinearBatchNormFusionQATPass)) + post_prepare.append(_histogram_observer_fix_pass) + + if is_qat: + # pre_convert_passes: tear down the simulated fusion and fold BN into + # the linear weights before convert_pt2e. + pre_convert: list[Callable] = [ + _wrap_exir_pass(RemoveSimulatedLinearBatchNormFusionQATPass), + _wrap_exir_pass(FuseBatchNormWithLinearPass), + ] + + # post_convert_passes: fix up quantization parameters for fused conv+BN + # bias nodes after convert_pt2e has inserted the quantize/dequantize ops. + post_convert: list[Callable] = [ + _wrap_exir_pass( + QuantizeFusedConvBnBiasAtenPass, + default_zero_bias=False, + symmetric_quant=True, + ) + ] + + return QuantizationRecipe( + quantizers=[_quantizer], + is_qat=True, + train_fn=rc.train_fn, + post_prepare_passes=post_prepare, + pre_convert_passes=pre_convert, + post_convert_passes=post_convert, + ) + else: + return QuantizationRecipe( + quantizers=[_quantizer], + post_prepare_passes=post_prepare, + ) + + +def _build_lowering_recipe( + compile_spec: list[CompileSpec], + neutron_target_spec: NeutronTargetSpec, + rc: NeutronRecipeConfig, + *, + delegate: bool, +) -> LoweringRecipe: + """Build the LoweringRecipe, optionally including NPU delegation.""" + partitioners = _build_partitioners(compile_spec, neutron_target_spec, rc, delegate) + pre_partitioning_callback = _build_pre_partitioning_callback(rc) + edge_manager_transform_passes = _build_edge_manager_transform_passes(rc) + + # The edge pass manager must be wrapped: EdgeTransformAndLowerStage calls + # edge_transform_passes with (method_name, ep) and expects a PassManager back. + return LoweringRecipe( + partitioners=partitioners, + edge_transform_passes=[NeutronEdgePassManagerWrapper()], + edge_compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _core_aten_ops_exception_list=core_aten_ops_exception_list, + ), + pre_partitioning_callback=pre_partitioning_callback, + edge_manager_transform_passes=edge_manager_transform_passes, + ) + + +def _build_partitioners( + compile_spec: list[CompileSpec], + neutron_target_spec: NeutronTargetSpec, + rc: NeutronRecipeConfig, + delegate: bool, +) -> list: + """Create the NeutronPartitioner list. Empty when delegate=False.""" + if not delegate: + return [] + return [ + NeutronPartitioner( + compile_spec, + neutron_target_spec, + rc.custom_delegation_options, + preserve_ops=default_preserve_ops, + ) + ] + + +def _build_pre_partitioning_callback(rc: NeutronRecipeConfig): + """Return a callback that assigns the post-quantization state_dict to NeutronPartitioner. + + NeutronPartitioner requires static parameter data. Since the partitioner is instantiated + during recipe creation (before model data is available), assignment is deferred to a + callback invoked just before partitioning. + """ + _use_quant_state_dict = rc.use_quant_state_dict + + def _callback( + _partitioners: list[Partitioner] | None, + programs: dict[str, ExportedProgram], + ) -> None: + if not _partitioners: + return + + if _use_quant_state_dict: + post_quant_state_dict: dict | None = {} + for _, program in programs.items(): + post_quant_state_dict.update(program.state_dict) + else: + post_quant_state_dict = None + + for _partitioner in _partitioners: + if isinstance(_partitioner, NeutronPartitioner): + _partitioner.post_quantization_state_dict = post_quant_state_dict + + return _callback + + +def _build_edge_manager_transform_passes(rc: NeutronRecipeConfig) -> list: + """Build edge_manager_transform_passes for the post-partitioning graph cleanup. + + These run in EdgeProgramManagerTransformStage, after to_edge_transform_and_lower: + - RemoveIOQuantOpsPass (optional, when remove_quant_io_ops=True) + - RemoveAdditionalQDQClustersPass (always applied) + - handle_kernel_selection side-effect (optional, when dump_kernel_selection_code=True) + + Each callable receives EdgeProgramManager and returns passes for epm.transform(), + or an empty list when no graph transformation is needed (side-effect only). + """ + passes = [] + + if rc.remove_quant_io_ops: + + def _remove_io_quant_ops(epm: EdgeProgramManager) -> list: + return [RemoveIOQuantOpsPass(edge_program_manager=epm)] + + passes.append(_remove_io_quant_ops) + + def _remove_additional_qdq_clusters( + epm: EdgeProgramManager, + ) -> NeutronEdgePassManager: + return NeutronEdgePassManager([RemoveAdditionalQDQClustersPass()]) + + passes.append(_remove_additional_qdq_clusters) + + if rc.dump_kernel_selection_code: + + def _handle_kernel_selection_side_effect(_epm: EdgeProgramManager) -> list: + # Side-effect only: write kernel-selection files. No graph transform needed. + handle_kernel_selection() + return [] + + passes.append(_handle_kernel_selection_side_effect) + + return passes diff --git a/backends/nxp/recipes/nxp_recipe_types.py b/backends/nxp/recipes/nxp_recipe_types.py new file mode 100644 index 00000000000..b816a064c67 --- /dev/null +++ b/backends/nxp/recipes/nxp_recipe_types.py @@ -0,0 +1,40 @@ +# Copyright 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. + +from executorch.export import RecipeType + + +NXP_BACKEND: str = "nxp" + + +class NXPRecipeType(RecipeType): + """NXP-specific recipe types for Neutron NPU export. + + Choose the recipe that matches your intended export configuration: + - INT8_PTQ_NEUTRON: standard post-training quantization, delegates to Neutron NPU. + - INT8_PTQ_NO_DELEGATE: PTQ without NPU delegation (useful for debugging or CPU-only deployment). + - INT8_QAT_NEUTRON: quantization-aware training, delegates to Neutron NPU. + - INT8_QAT_NO_DELEGATE: QAT without NPU delegation (useful for accuracy evaluation). + """ + + # INT8 static PTQ (weights + activations). Calibration dataset required. + # Applicable operators are delegated to the Neutron NPU. + INT8_PTQ_NEUTRON = "nxp_int8_ptq_neutron" + + # INT8 PTQ without NPU delegation. Produces a quantized graph that runs on CPU. + # Useful for accuracy evaluation or debugging before enabling delegation. + INT8_PTQ_NO_DELEGATE = "nxp_int8_ptq_no_delegate" + + # INT8 QAT (weights + activations). A train_fn must be provided in NeutronRecipeConfig. + # Applicable operators are delegated to the Neutron NPU. + INT8_QAT_NEUTRON = "nxp_int8_qat_neutron" + + # INT8 QAT without NPU delegation. Produces a quantized graph that runs on CPU. + # Useful for accuracy evaluation or debugging before enabling delegation. + INT8_QAT_NO_DELEGATE = "nxp_int8_qat_no_delegate" + + @classmethod + def get_backend_name(cls) -> str: + return NXP_BACKEND diff --git a/backends/nxp/run_unittests.sh b/backends/nxp/run_unittests.sh index 66e51c39a1d..d515d12ae1c 100755 --- a/backends/nxp/run_unittests.sh +++ b/backends/nxp/run_unittests.sh @@ -10,6 +10,10 @@ EXECUTORCH_DIR=$(dirname $(dirname $SCRIPT_DIR)) cd $EXECUTORCH_DIR +# Cap pytest-xdist's workers to the container's CPU quota. Applies to +# `-n logical` as well, despite the variable's name. +source .ci/scripts/pytest-parallelism.sh + # '-c /dev/null' is used to ignore root level pytest.ini. pytest -c /dev/null -n "logical" backends/nxp/tests/ diff --git a/backends/nxp/runtime/NeutronBackend.cpp b/backends/nxp/runtime/NeutronBackend.cpp index 4c75f76712a..1effb21a536 100644 --- a/backends/nxp/runtime/NeutronBackend.cpp +++ b/backends/nxp/runtime/NeutronBackend.cpp @@ -442,8 +442,20 @@ class NeutronBackend final : public PyTorchBackendInterface { auto arg = args[cfg->inputMap[i]]->toTensor(); auto dim_order = arg.dim_order().data(); - if (cfg->inputTranspositionFlags[i] && - multipleChannelsPresent(arg.sizes())) { + if (cfg->inputTranspositionFlags[i]) { + if (!multipleChannelsPresent(arg.sizes())) { + // The input has only 1 channel, so NCHW and NHWC data is equivalent + // and no transposition is needed. + if (!is_channels_last_dim_order(dim_order, arg.dim()) && + !is_contiguous_dim_order(dim_order, arg.dim())) { + ET_LOG(Error, "Input %d uses unsupported dim-order.", i); + print_dim_order(dim_order, arg.dim()); + return Error::InvalidProgram; + } + + cfg->dcfg.inputs[i] = arg.const_data_ptr(); + continue; + } // The input must be transposed. if (arg.sizes().size() < 3) { ET_LOG(Error, "Unable to transpose 1D and 2D input to channel last"); @@ -496,10 +508,21 @@ class NeutronBackend final : public PyTorchBackendInterface { auto arg = args[cfg->numInputArgs + cfg->outputMap[i]]->toTensor(); auto dim_order = arg.dim_order().data(); - if (cfg->outputTranspositionFlags[i] && - multipleChannelsPresent(arg.sizes())) { - // The output will have to be transposed. + if (cfg->outputTranspositionFlags[i]) { + if (!multipleChannelsPresent(arg.sizes())) { + // The output has only 1 channel, so NCHW and NHWC data is equivalent + // and no transposition is needed. + if (!is_channels_last_dim_order(dim_order, arg.dim()) && + !is_contiguous_dim_order(dim_order, arg.dim())) { + ET_LOG(Error, "Output %d uses unsupported dim-order.", i); + print_dim_order(dim_order, arg.dim()); + return Error::InvalidProgram; + } + cfg->dcfg.outputs[i] = arg.mutable_data_ptr(); + continue; + } + // The output will have to be transposed. if (is_channels_last_dim_order(dim_order, arg.dim())) { // The tensor will already be correctly permuted. No transposition // needed. @@ -636,6 +659,14 @@ class NeutronBackend final : public PyTorchBackendInterface { index++; } } + // The neutronGetSdkVersion() function is available starting with Neutron + // Software 3.2.1. The code below is not backward compatible with earlier + // Neutron Software versions. + NeutronSdkVersion neutron_sdk_version = neutronGetSdkVersion(); + uint16_t neutron_sdk_version_uint16 = + static_cast(neutron_sdk_version.major << 8) | + static_cast(neutron_sdk_version.minor << 4) | + static_cast(neutron_sdk_version.patch); event_tracer_log_profiling_delegate( tracer, nullptr, @@ -643,9 +674,8 @@ class NeutronBackend final : public PyTorchBackendInterface { neutron_events[events_num - 1].startEvent.time, neutron_events[events_num - 1].stopEvent.time + stop_ticks - start_ticks, - static_cast( - &neutron_events[events_num - 1].startEvent.functionCode), - sizeof(uint8_t)); + static_cast(&neutron_sdk_version_uint16), + sizeof(uint16_t)); } #endif diff --git a/backends/nxp/runtime/NeutronDriver.h b/backends/nxp/runtime/NeutronDriver.h index 5c47bd74eab..0d29f8b172b 100644 --- a/backends/nxp/runtime/NeutronDriver.h +++ b/backends/nxp/runtime/NeutronDriver.h @@ -42,7 +42,7 @@ typedef void* NeutronModelHandle; typedef struct { /// Neutron microcode buffer address. - /// The Neutron microcode is generated by the Neutron converter tool. + /// The Neutron microcode is generated by the Neutron compiler tool. /// The microcode buffer, 16 bytes aligned, is allocated and initialized by /// the application or ML framework. The microcode buffer is passed by /// reference to the Neutron firmware. The microcode buffer is specific for a @@ -50,7 +50,7 @@ typedef struct { const void* microcode; /// Neutron weights buffer address. - /// The Neutron weights is generated by the Neutron converter tool. + /// The Neutron weights is generated by the Neutron compiler tool. /// The weights buffer, 16 bytes aligned, is allocated and initialized by the /// application or ML framework. The weights buffer address is passed by /// reference to the Neutron-firmware. The weights buffer is specific for a @@ -58,7 +58,7 @@ typedef struct { const void* weights; /// Neutron kernels buffer address. - /// The Neutron kernels are generated by the Neutron converter tool. + /// The Neutron kernels are generated by the Neutron compiler tool. /// The kernels buffer, 16 bytes aligned, is allocated and initialized by the /// application or ML framework. The kernels buffer address is passed by /// reference to the Neutron-firmware. The kernels buffer is specific for a @@ -124,6 +124,16 @@ typedef struct { void (*wait)(uint32_t channel); } NeutronConfig; +/// This structure contains semantic version of the Neutron SDK +/// (major.minor.patch) and the SHA version (string and uint32_t). +typedef struct { + uint32_t major; + uint32_t minor; + uint32_t patch; + const char* hashString; + uint32_t hashUint32; +} NeutronSdkVersion; + /* Invalid handle, returned by neutronModelPrepare() if an error occurred. */ #define NEUTRON_INVALID_HANDLE NULL @@ -224,6 +234,9 @@ NeutronError neutronSetConfig(NeutronConfig* config); /// - Used to get NeutronContext size. size_t neutronGetModelContextSize(); +/// - Used to get Neutron SDK version. +NeutronSdkVersion neutronGetSdkVersion(); + /// - Allocates size bytes and returns a pointer to the allocated memory. /// The returned pointer address will be a multiple of the alignment. /// Returns NULL on failure. diff --git a/backends/nxp/tests/BUCK b/backends/nxp/tests/BUCK index 7879e1e3db5..3d991e1eef0 100644 --- a/backends/nxp/tests/BUCK +++ b/backends/nxp/tests/BUCK @@ -4,17 +4,6 @@ load("@fbcode_macros//build_defs:python_pytest.bzl", "python_pytest") oncall("executorch") -fbcode_target(_kind = runtime.python_library, - name = "ops_aliases", - srcs = [ - "ops_aliases.py", - ], - deps = [ - "//caffe2:torch", - "//executorch/exir:lib", - ], -) - fbcode_target(_kind = runtime.python_library, name = "models", srcs = [ @@ -159,7 +148,7 @@ fbcode_target(_kind = python_pytest, fbcode_target(_kind = python_pytest, name = "test_neutron_converter_manager", srcs = [ - "generic_tests/test_neutron_converter_manager.py", + "generic_tests/test_neutron_compiler_manager.py", ], deps = [ "//executorch/backends/nxp:neutron_sdk", diff --git a/backends/nxp/tests/cortex_m_benchmarking.py b/backends/nxp/tests/cortex_m_benchmarking.py new file mode 100644 index 00000000000..7859e9e2c3b --- /dev/null +++ b/backends/nxp/tests/cortex_m_benchmarking.py @@ -0,0 +1,79 @@ +# Copyright 2025-2026 Arm Limited and/or its affiliates. +# Copyright 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. + +from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig +from executorch.backends.cortex_m.test.tester import CortexMQuantize, CortexMTester +from executorch.backends.nxp.tests.utils import ( + process_input_sample, + process_output_sample, + read_prepared_samples, + store_results, +) +from executorch.backends.test.harness.stages import StageType + + +class CortexMNXPBenchmarkTester(CortexMTester): + def __init__( + self, + module, + example_inputs, + target_config: CortexMTargetConfig | None = None, + timeout: int = 120, + ): + target_config = target_config or CortexMTargetConfig( + cpu=CortexM.M33 + ) # set default to M33 for NXP boards + super().__init__(module, example_inputs, target_config, timeout) + + def run_benchmark( + self, + calibration_samples, + input_spec, + output_spec, + testing_dataset_dir, + cpu_results_dir, + npu_results_dir, + ): + quantization_stage = CortexMQuantize(calibration_samples=calibration_samples) + + self.quantize(quantization_stage) + self.export() + self.to_edge() + self.run_passes() + self.to_executorch() + self.serialize() + self.run_program( + input_spec, + output_spec, + testing_dataset_dir, + cpu_results_dir, + npu_results_dir, + ) + + return self.stages[StageType.SERIALIZE].executorch_program_manager + + def run_program( + self, + input_spec, + output_spec, + testing_dataset_dir, + cpu_results_dir, + npu_results_dir, + ): + all_outputs = [] + + for input_samples in read_prepared_samples(testing_dataset_dir, input_spec): + current_input_samples = process_input_sample(input_spec, input_samples) + + # Run the model. + output = self.stages[StageType.SERIALIZE].run_artifact( + *current_input_samples + ) + current_outputs = process_output_sample(output, output_spec) + all_outputs.append(current_outputs) + + # Store all the results. + store_results(all_outputs, cpu_results_dir, npu_results_dir) diff --git a/backends/nxp/tests/executorch_pipeline.py b/backends/nxp/tests/executorch_pipeline.py index 964b8de159c..0567f62b101 100644 --- a/backends/nxp/tests/executorch_pipeline.py +++ b/backends/nxp/tests/executorch_pipeline.py @@ -33,6 +33,7 @@ from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner from executorch.backends.nxp.nxp_backend import ( core_aten_ops_exception_list, + default_preserve_ops, generate_neutron_compile_spec, ) from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer @@ -92,7 +93,7 @@ def get_random_calibration_inputs( ] -def _get_default_quantizer(target_spec: NeutronTargetSpec, use_qat: bool) -> Quantizer: +def get_default_quantizer(target_spec: NeutronTargetSpec, use_qat: bool) -> Quantizer: return NeutronQuantizer(target_spec, is_qat=use_qat) @@ -154,7 +155,7 @@ def _nested( return _nested -def _get_example_input( +def get_example_input( input_spec: tuple[ModelInputSpec, ...], ) -> tuple[torch.Tensor, ...]: example_input = [] @@ -195,12 +196,10 @@ def to_quantized_edge_program( ) -> EdgeProgramManager: _neutron_target_spec = NeutronTargetSpec(target) if get_quantizer_fn is None: - get_quantizer_fn = partial( - _get_default_quantizer, _neutron_target_spec, use_qat - ) + get_quantizer_fn = partial(get_default_quantizer, _neutron_target_spec, use_qat) input_spec = to_model_input_spec(input_spec) calibration_inputs = get_calibration_inputs_fn(input_spec) - example_input = _get_example_input(input_spec) + example_input = get_example_input(input_spec) # Make sure the model is in the evaluation mode. model.eval() @@ -215,12 +214,6 @@ def to_quantized_edge_program( train_fn=train_fn, ) - # List of operators to not decompose during the lowering. - preserve_ops = [ - torch.ops.aten.prelu.default, - torch.ops.aten.pad.default, - torch.ops.aten.hardswish.default, - ] compile_spec = generate_neutron_compile_spec( target, intermediates_dir=intermediates_dir, @@ -240,7 +233,7 @@ def to_quantized_edge_program( _neutron_target_spec, custom_delegation_options, post_quant_state_dict, - preserve_ops=preserve_ops, + preserve_ops=default_preserve_ops, ) ] else: @@ -318,7 +311,7 @@ def to_edge_program( model: nn.Module, input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]], ) -> EdgeProgramManager: - example_input = _get_example_input(to_model_input_spec(input_spec)) + example_input = get_example_input(to_model_input_spec(input_spec)) # Make sure the model is in the evaluation mode. model.eval() diff --git a/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py b/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py index b426c34c260..a3eae135144 100644 --- a/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py +++ b/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py @@ -15,6 +15,13 @@ from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import ( NeutronAtenPassManager, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AdaptiveAvgPool2D, + AvgPool2D, + GetItem, + MaxPool2DWithIndices, + ViewCopy, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import neutron_target_spec from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops @@ -22,19 +29,12 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( AdaptiveAvgPool2dModule, AvgPool2dModule, MaxPool2dModule, ) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AdaptiveAvgPool2D, - AvgPool2D, - GetItem, - MaxPool2DWithIndices, - ViewCopy, -) @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/generic_tests/test_aot_example.py b/backends/nxp/tests/generic_tests/test_aot_example.py index 1f8dc410917..b75d3605d34 100644 --- a/backends/nxp/tests/generic_tests/test_aot_example.py +++ b/backends/nxp/tests/generic_tests/test_aot_example.py @@ -2,6 +2,7 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + import contextlib import os import subprocess @@ -186,6 +187,9 @@ def test_aot_example__mlperf_tiny_ic(): """Test that the MLPerf Tiny image classification model (ResNet-8) can be lowered to Neutron backend via `aot_neutron_compile.py` and all ops are delegated.""" + # Number of random samples to generate, must be divisible by number of classes + num_random_samples = 60 + # Run the compilation script as a module (like run_aot_example.sh does). # The calibration data of this model is generated randomly, so no dataset download is needed. cmd = [ @@ -199,6 +203,8 @@ def test_aot_example__mlperf_tiny_ic(): "--target", "imxrt700", "--use_random_dataset", + "--num_random_samples", + str(num_random_samples), ] # Output file will be created in executorch_root @@ -215,7 +221,10 @@ def test_aot_example__mlperf_tiny_ic(): def test_aot_example__mlperf_tiny_ic__profiling(): """Test that the MLPerf Tiny image classification model (ResNet-8) can be lowered to Neutron backend via - `aot_neutron_compile.py` and all ops are delegated.""" + `aot_neutron_compile.py` and profiling works as intended.""" + + # Number of random samples to generate, must be divisible by number of classes + num_random_samples = 60 # Run the compilation script as a module (like run_aot_example.sh does) # Channels-last is buggy, so channels-first is used instead @@ -232,6 +241,8 @@ def test_aot_example__mlperf_tiny_ic__profiling(): "--remove-quant-io-ops", "--use_profiling", # Generate profilable model and create ETRecord "--use_random_dataset", # Avoid downloading the dataset. + "--num_random_samples", + str(num_random_samples), ] # Output files will be created in executorch_root. @@ -249,3 +260,77 @@ def test_aot_example__mlperf_tiny_ic__profiling(): with _cleanup_generated_files(pte_file, etrecord_file): result = _run_compile(cmd) _assert_profiling(result, pte_file, etrecord_file) + + +def test_aot_example__mlperf_tiny_kws(): + """Test that the MLPerf Tiny keyword spotting model (DS-CNN) can be lowered to Neutron backend via + `aot_neutron_compile.py` and all ops are delegated.""" + + # Number of random samples to generate, must be divisible by number of classes + num_random_samples = 60 + + # Run the compilation script as a module (like run_aot_example.sh does). + cmd = [ + sys.executable, + "-m", + "examples.nxp.aot_neutron_compile", + "--model_name", + "mlperf_tiny_keyword_spotting", + "--delegate", + "--quantize", + "--target", + "imxrt700", + "--use_random_dataset", + "--num_random_samples", + str(num_random_samples), + ] + + # Output file will be created in executorch_root + pte_file = Path( + os.path.join(EXECUTORCH_ROOT, "mlperf_tiny_keyword_spotting_nxp_delegate.pte") + ) + + with _cleanup_generated_files(pte_file): + result = _run_compile(cmd) + _assert_delegation(result, pte_file) + + +def test_aot_example__mlperf_tiny_kws__profiling(): + """Test that the MLPerf Tiny keyword spotting model (DS-CNN) can be lowered to Neutron backend via + `aot_neutron_compile.py` and profiling works as intended.""" + + # Number of random samples to generate, must be divisible by number of classes + num_random_samples = 60 + + # Run the compilation script as a module (like run_aot_example.sh does) + cmd = [ + sys.executable, + "-m", + "examples.nxp.aot_neutron_compile", + "--model_name", + "mlperf_tiny_keyword_spotting", + "--delegate", + "--quantize", + "--target", + "imxrt700", + "--remove-quant-io-ops", + "--use_profiling", # Generate profilable model and create ETRecord + "--use_random_dataset", + "--num_random_samples", + str(num_random_samples), + ] + + pte_file = Path( + os.path.join( + EXECUTORCH_ROOT, "mlperf_tiny_keyword_spotting_nxp_delegate_profile.pte" + ) + ) + etrecord_file = Path( + os.path.join( + EXECUTORCH_ROOT, "etrecord", "mlperf_tiny_keyword_spotting_etrecord.bin" + ) + ) + + with _cleanup_generated_files(pte_file, etrecord_file): + result = _run_compile(cmd) + _assert_profiling(result, pte_file, etrecord_file) diff --git a/backends/nxp/tests/generic_tests/test_batch_norm_fusion.py b/backends/nxp/tests/generic_tests/test_batch_norm_fusion.py index 64cd7ce83f0..85c02e0d53d 100644 --- a/backends/nxp/tests/generic_tests/test_batch_norm_fusion.py +++ b/backends/nxp/tests/generic_tests/test_batch_norm_fusion.py @@ -6,6 +6,8 @@ from copy import deepcopy import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import ( @@ -27,7 +29,7 @@ graph_contains_any_of_ops, OverrideTargetSupportCheck, ) -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.simple_models import ( ConvBatchNormModule, LinearBatchNormModule, ) diff --git a/backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py b/backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py index 1b1aaed897e..0f931075927 100644 --- a/backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py +++ b/backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py @@ -4,6 +4,8 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch @@ -13,12 +15,15 @@ from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters import ( ViewCopyConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + Cat, + ExecutorchDelegateCall, + SubTensor, + ViewCopy, +) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops -from executorch.exir.dialects._ops import ops as exir_ops - -# noinspection PyProtectedMember -ExecutorchDelegateCall = torch.ops.higher_order.executorch_call_delegate class SingleViewCopyModule(torch.nn.Module): @@ -70,7 +75,7 @@ def test_single_view_copy_partition(): ep = to_quantized_edge_program(module, input_shape).exported_program() # Make sure the `view_copy` was not delegated. - assert graph_contains_any_of_ops(ep.graph, [exir_ops.edge.aten.view_copy.default]) + assert graph_contains_any_of_ops(ep.graph, [ViewCopy]) assert not graph_contains_any_of_ops(ep.graph, [ExecutorchDelegateCall]) @@ -114,18 +119,18 @@ def test_noop_partitions__concatenate_one_tensor_and_add_zeros(): assert graph_contains_any_of_ops( ep.graph, [ - exir_ops.edge.aten.cat.default, - exir_ops.edge.aten.add.Tensor, + Cat, + AddTensor, ], ) @pytest.mark.xfail( strict=True, - reason="Neutron Converter currently supports these 2 noops in sequence.", + reason="Neutron Compiler currently supports these 2 noops in sequence.", ) def test_noop_partitions__concatenate_one_tensor_and_add_zeros__forced_delegation(): - # When the noop `Concatenate` and noop `Add` are in sequence, Neutron Converter supports them. This edge case is + # When the noop `Concatenate` and noop `Add` are in sequence, Neutron Compiler supports them. This edge case is # not reflected in our logic. But as this edge case is extremely rare (and even if it ever happened in a real # model, the consequences would be minimal), fixing it is not a priority. @@ -158,8 +163,8 @@ def test_noop_partitions__add_sub(): assert graph_contains_any_of_ops( ep.graph, [ - exir_ops.edge.aten.add.Tensor, - exir_ops.edge.aten.sub.Tensor, + AddTensor, + SubTensor, ], ) diff --git a/backends/nxp/tests/test_convert_1d_conv_to_2d.py b/backends/nxp/tests/generic_tests/test_convert_1d_conv_to_2d.py similarity index 98% rename from backends/nxp/tests/test_convert_1d_conv_to_2d.py rename to backends/nxp/tests/generic_tests/test_convert_1d_conv_to_2d.py index b7db6fbc46e..d4a07d86fab 100644 --- a/backends/nxp/tests/test_convert_1d_conv_to_2d.py +++ b/backends/nxp/tests/generic_tests/test_convert_1d_conv_to_2d.py @@ -4,6 +4,8 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import ( @@ -14,6 +16,10 @@ from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + Convolution, + ExecutorchDelegateCall, +) from executorch.backends.nxp.tests.executorch_pipeline import ( neutron_target_spec, to_quantized_edge_program, @@ -22,8 +28,10 @@ convert_run_compare, graph_contains_any_of_ops, ) -from executorch.backends.nxp.tests.models import Conv1dModule, ConvTranspose1dModule -from executorch.exir.dialects._ops import ops as exir_ops +from executorch.backends.nxp.tests.simple_models import ( + Conv1dModule, + ConvTranspose1dModule, +) from torch import nn from torch.export import ExportedProgram @@ -46,9 +54,6 @@ def reseed_model_per_test_run(): 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 - @pytest.mark.parametrize( "input_shape, kernel_size, stride, padding, dilation, groups, bias", @@ -301,7 +306,7 @@ def test_convert_conv_1d_to_conv2d_full_pipeline( # Make sure `edge.aten.convolution.default` is in the model. assert graph_contains_any_of_ops( exported_program.graph, - [EdgeConvolution], + [Convolution], ) example_input = (np.random.random(input_shape).astype(np.float32) * 50).astype( @@ -378,7 +383,7 @@ def test_convert_conv_1d_to_conv2d_transp_full_pipeline( # Make sure `edge.aten.convolution.default` is in the model. assert graph_contains_any_of_ops( exported_program.graph, - [EdgeConvolution], + [Convolution], ) example_input = (np.random.random(input_shape).astype(np.float32) * 50).astype( diff --git a/backends/nxp/tests/generic_tests/test_convert_div_to_mul.py b/backends/nxp/tests/generic_tests/test_convert_div_to_mul.py index 3415b79a39d..89879226c3a 100644 --- a/backends/nxp/tests/generic_tests/test_convert_div_to_mul.py +++ b/backends/nxp/tests/generic_tests/test_convert_div_to_mul.py @@ -13,16 +13,16 @@ ConvertDivToMulPass, NeutronAtenPassManager, ) +from executorch.backends.nxp.backend.ops_aliases import MulTensor from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import neutron_target_spec from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( NonstaticDivLinearModel, StaticDivLinearModel, ) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import MulTensor @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py b/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py index b3d68eeeb56..be214b77568 100644 --- a/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py +++ b/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py @@ -18,20 +18,20 @@ from executorch.backends.nxp.backend.edge_helper import ( try_get_tensor_constant_from_node, ) +from executorch.backends.nxp.backend.ops_aliases import AddTensor, MulTensor, SubTensor from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import neutron_target_spec from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import ( - AddScalarModule, - MulScalarModule, - SubScalarModule, -) from executorch.backends.nxp.tests.nsys_testing import ( AllCloseOutputComparator, lower_run_compare, ) -from executorch.backends.nxp.tests.ops_aliases import AddTensor, MulTensor, SubTensor +from executorch.backends.nxp.tests.simple_models import ( + AddScalarModule, + MulScalarModule, + SubScalarModule, +) @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/generic_tests/test_debug_results.py b/backends/nxp/tests/generic_tests/test_debug_results.py index ccf5c13b501..387acd69b2c 100644 --- a/backends/nxp/tests/generic_tests/test_debug_results.py +++ b/backends/nxp/tests/generic_tests/test_debug_results.py @@ -7,17 +7,19 @@ import os import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec from executorch.backends.nxp.tests.graph_verifier import BaseGraphVerifier -from executorch.backends.nxp.tests.models import AddTensorModule, AvgPool2dModule from executorch.backends.nxp.tests.nsys_testing import ( get_test_name, lower_run_compare, OUTPUTS_DIR, ) +from executorch.backends.nxp.tests.simple_models import AddTensorModule, AvgPool2dModule @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py b/backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py index de4e684405c..cddf2feff34 100644 --- a/backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py +++ b/backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py @@ -4,6 +4,8 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch @@ -12,16 +14,16 @@ NeutronAtenPassManager, SplitGRUBasedOnNumLayers, ) +from executorch.backends.nxp.backend.ops_aliases import SliceCopy from executorch.backends.nxp.tests.executorch_pipeline import neutron_target_spec from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( GRUModel, SplitWithSections, SplitWithSize, ) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import SliceCopy @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/test_fold_redundant_qdq.py b/backends/nxp/tests/generic_tests/test_fold_redundant_qdq.py similarity index 97% rename from backends/nxp/tests/test_fold_redundant_qdq.py rename to backends/nxp/tests/generic_tests/test_fold_redundant_qdq.py index e79d29b1166..7da9986ad7b 100644 --- a/backends/nxp/tests/test_fold_redundant_qdq.py +++ b/backends/nxp/tests/generic_tests/test_fold_redundant_qdq.py @@ -6,13 +6,12 @@ import torch +from executorch.backends.nxp.backend.ops_aliases import ExecutorchDelegateCall from executorch.backends.nxp.edge_passes.fold_redundant_qdq_pass import ( FoldRedundantDequantizeQuantizePass, ) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -ExecutorchDelegateCall = torch.ops.higher_order.executorch_call_delegate - class ConvDropoutConvModule(torch.nn.Module): """Two conv clusters separated by an eval-mode dropout (an identity). diff --git a/backends/nxp/tests/test_fuse_batch_norm_single_user.py b/backends/nxp/tests/generic_tests/test_fuse_batch_norm_single_user.py similarity index 100% rename from backends/nxp/tests/test_fuse_batch_norm_single_user.py rename to backends/nxp/tests/generic_tests/test_fuse_batch_norm_single_user.py diff --git a/backends/nxp/tests/generic_tests/test_gru_splitting.py b/backends/nxp/tests/generic_tests/test_gru_splitting.py index 297f9677fb2..03f3f4a2947 100644 --- a/backends/nxp/tests/generic_tests/test_gru_splitting.py +++ b/backends/nxp/tests/generic_tests/test_gru_splitting.py @@ -4,6 +4,8 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch diff --git a/backends/nxp/tests/generic_tests/test_integration.py b/backends/nxp/tests/generic_tests/test_integration.py index 9916cba5bdd..869aed432f8 100644 --- a/backends/nxp/tests/generic_tests/test_integration.py +++ b/backends/nxp/tests/generic_tests/test_integration.py @@ -5,14 +5,14 @@ import executorch.extension.pybindings.portable_lib import executorch.kernels.quantized # noqa F401 +from executorch.backends.nxp.backend.ops_aliases import AddMM, Convolution from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops -from executorch.backends.nxp.tests.ops_aliases import AddMM, Convolution from executorch.backends.nxp.tests.use_qat import * # noqa F401 from executorch.backends.nxp.tests.executorch_pipeline import ( to_quantized_executorch_program, ) -from executorch.backends.nxp.tests.models import ConvFCSoftmaxModule +from executorch.backends.nxp.tests.simple_models import ConvFCSoftmaxModule from executorch.devtools.backend_debug import get_delegation_info from executorch.examples.nxp.experimental.cifar_net.cifar_net import CifarNet diff --git a/backends/nxp/tests/generic_tests/test_kernel_selection.py b/backends/nxp/tests/generic_tests/test_kernel_selection.py index 732913c0485..2dd498881c7 100644 --- a/backends/nxp/tests/generic_tests/test_kernel_selection.py +++ b/backends/nxp/tests/generic_tests/test_kernel_selection.py @@ -7,11 +7,13 @@ import eiq_neutron_sdk import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.simple_models import ( AdaptiveAvgPool2dConvModule, Conv2dReLUMaxPoolModule, ) diff --git a/backends/nxp/tests/generic_tests/test_move_activation_before_concatenation.py b/backends/nxp/tests/generic_tests/test_move_activation_before_concatenation.py index 6aa07dbba8d..64a7a5e162d 100644 --- a/backends/nxp/tests/generic_tests/test_move_activation_before_concatenation.py +++ b/backends/nxp/tests/generic_tests/test_move_activation_before_concatenation.py @@ -19,6 +19,15 @@ from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddMM, + Cat, + Convolution, + HardTanh, + Relu, + Sigmoid, + Tanh, +) from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer from executorch.backends.nxp.quantizer.utils import calibrate_and_quantize from executorch.backends.nxp.tests.executorch_pipeline import ( @@ -33,21 +42,20 @@ ToChannelFirstPreprocess, ToChannelLastPreprocess, ) -from executorch.backends.nxp.tests.models import get_activation -from executorch.exir.dialects._ops import ops as exir_ops +from executorch.backends.nxp.tests.simple_models import get_activation from parameterized import parameterized from torch import nn from torch.export import ExportedProgram from torch.fx import GraphModule concat_cluster_ops = [ - exir_ops.edge.aten.addmm.default, - exir_ops.edge.aten.convolution.default, - exir_ops.edge.aten.hardtanh.default, - exir_ops.edge.aten.relu.default, - exir_ops.edge.aten.sigmoid.default, - exir_ops.edge.aten.tanh.default, - exir_ops.edge.aten.cat.default, + AddMM, + Convolution, + HardTanh, + Relu, + Sigmoid, + Tanh, + Cat, ] diff --git a/backends/nxp/tests/generic_tests/test_neutron_backend.py b/backends/nxp/tests/generic_tests/test_neutron_backend.py index 867b585ef64..f2eaf097b14 100644 --- a/backends/nxp/tests/generic_tests/test_neutron_backend.py +++ b/backends/nxp/tests/generic_tests/test_neutron_backend.py @@ -4,7 +4,10 @@ # LICENSE file in the root directory of this source tree. from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.models import Conv2dModule, LinearSoftmaxModule +from executorch.backends.nxp.tests.simple_models import ( + Conv2dModule, + LinearSoftmaxModule, +) def test_neutron_backend__single_conv_model(): diff --git a/backends/nxp/tests/generic_tests/test_neutron_backend_executor.py b/backends/nxp/tests/generic_tests/test_neutron_backend_executor.py index 06a95142b1a..6dbbf177b7d 100644 --- a/backends/nxp/tests/generic_tests/test_neutron_backend_executor.py +++ b/backends/nxp/tests/generic_tests/test_neutron_backend_executor.py @@ -21,7 +21,10 @@ TFLiteExecutor, ToNHWCPreprocess, ) -from executorch.backends.nxp.tests.models import Conv2dModule, ConvFCSoftmaxModule +from executorch.backends.nxp.tests.simple_models import ( + Conv2dModule, + ConvFCSoftmaxModule, +) from torch.export import ExportedProgram diff --git a/backends/nxp/tests/generic_tests/test_neutron_converter_manager.py b/backends/nxp/tests/generic_tests/test_neutron_compiler_manager.py similarity index 79% rename from backends/nxp/tests/generic_tests/test_neutron_converter_manager.py rename to backends/nxp/tests/generic_tests/test_neutron_compiler_manager.py index 8bd3446da7a..2b3b798dd67 100644 --- a/backends/nxp/tests/generic_tests/test_neutron_converter_manager.py +++ b/backends/nxp/tests/generic_tests/test_neutron_compiler_manager.py @@ -6,34 +6,34 @@ import multiprocessing import pickle -from executorch.backends.nxp.backend.neutron_converter_manager import ( - NeutronConverterManager, +from executorch.backends.nxp.backend.neutron_compiler_manager import ( + NeutronCompilerManager, ) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.models import LinearModule +from executorch.backends.nxp.tests.simple_models import LinearModule def test_conv2d_neutron_conversion__prefetching(mocker): model = LinearModule(True) input_shape = (1, 1, 32, 32) - converter_spy = mocker.spy(NeutronConverterManager, "convert") + compiler_spy = mocker.spy(NeutronCompilerManager, "compile") _ = to_quantized_edge_program( model, input_shape, fetch_constants_to_sram=True ).exported_program() - neutron_model_prefetch = converter_spy.spy_return + neutron_model_prefetch = compiler_spy.spy_return _ = to_quantized_edge_program( model, input_shape, fetch_constants_to_sram=False ).exported_program() - neutron_model_regular = converter_spy.spy_return + neutron_model_regular = compiler_spy.spy_return assert len(neutron_model_prefetch) != len( neutron_model_regular ), "The weight prefetching flag does not make a difference!" -def test_convert_unsafe_args_are_picklable(mocker): +def test_compile_unsafe_args_are_picklable(mocker): """Verify that all args passed to `multiprocessing.Process` are picklable. The subprocess uses forkserver/spawn in some environments, which requires diff --git a/backends/nxp/tests/generic_tests/test_node_format_inference.py b/backends/nxp/tests/generic_tests/test_node_format_inference.py index 18d5f874aab..8d6ac52627d 100644 --- a/backends/nxp/tests/generic_tests/test_node_format_inference.py +++ b/backends/nxp/tests/generic_tests/test_node_format_inference.py @@ -13,18 +13,18 @@ NodeFormatInference, NXP_NODE_FORMAT, ) +from executorch.backends.nxp.backend.ops_aliases import ( + ExecutorchDelegateCall, + MaxPool2DWithIndices, +) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.simple_models import ( Conv2dModule, MaxPool2dModule, SoftmaxModule, ) -from executorch.backends.nxp.tests.ops_aliases import ( - ExecutorchDelegateCall, - MaxPool2DWithIndices, -) def test_convolution(): diff --git a/backends/nxp/tests/generic_tests/test_operator_selector.py b/backends/nxp/tests/generic_tests/test_operator_selector.py index ca301daf738..71c30c73fac 100644 --- a/backends/nxp/tests/generic_tests/test_operator_selector.py +++ b/backends/nxp/tests/generic_tests/test_operator_selector.py @@ -4,7 +4,7 @@ # LICENSE file in the root directory of this source tree.import torch from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.models import Conv2dModule +from executorch.backends.nxp.tests.simple_models import Conv2dModule def test_operator_selector_mechanism(): diff --git a/backends/nxp/tests/generic_tests/test_per_channel_conversion.py b/backends/nxp/tests/generic_tests/test_per_channel_conversion.py index af9ef08057b..25a4b28224b 100644 --- a/backends/nxp/tests/generic_tests/test_per_channel_conversion.py +++ b/backends/nxp/tests/generic_tests/test_per_channel_conversion.py @@ -12,6 +12,10 @@ from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + Convolution, + DequantizePerChannel, +) from executorch.backends.nxp.quantizer.neutron_quantizer import ( act_qspec, NeutronAtenQuantizer, @@ -29,8 +33,7 @@ ToChannelFirstPreprocess, ToChannelLastPreprocess, ) -from executorch.backends.nxp.tests.models import Conv2dModule -from executorch.exir.dialects._ops import ops as exir_ops +from executorch.backends.nxp.tests.simple_models import Conv2dModule from parameterized import parameterized from torch import fx @@ -172,16 +175,10 @@ def test_per_channel_convolution(self, _, use_qat: bool): conv_nodes = [ node for node in exported_program.graph.nodes - if node.target == exir_ops.edge.aten.convolution.default + if node.target == Convolution ] assert len(conv_nodes) == 1 conv_node = conv_nodes[0] - assert ( - conv_node.args[1].target - == exir_ops.edge.quantized_decomposed.dequantize_per_channel.default - ) - assert ( - conv_node.args[2].target - == exir_ops.edge.quantized_decomposed.dequantize_per_channel.default - ) + assert conv_node.args[1].target == DequantizePerChannel + assert conv_node.args[2].target == DequantizePerChannel diff --git a/backends/nxp/tests/generic_tests/test_profiling.py b/backends/nxp/tests/generic_tests/test_profiling.py index cd90bdd345b..09ae6c079ed 100644 --- a/backends/nxp/tests/generic_tests/test_profiling.py +++ b/backends/nxp/tests/generic_tests/test_profiling.py @@ -2,6 +2,7 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + import ast import logging import os @@ -9,19 +10,26 @@ from typing import Any, Union import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.tests.graph_verifier import BaseGraphVerifier from executorch.backends.nxp.tests.model_output_comparator import ( NumericalStatsOutputComparator, ) -from executorch.backends.nxp.tests.models import AvgPool2dModule, SoftmaxModule from executorch.backends.nxp.tests.nsys_testing import ( get_test_name, lower_run_compare, OUTPUTS_DIR, ) +from executorch.backends.nxp.tests.profiling_utils import ( + get_neutron_compiler_version, + get_neutron_driver_version, + get_neutron_kernel_kinds, +) +from executorch.backends.nxp.tests.simple_models import AvgPool2dModule, SoftmaxModule from executorch.devtools.inspector._inspector import Inspector from executorch.examples.models.mlperf_tiny import ( DeepAutoEncoder, @@ -65,12 +73,15 @@ def inspector_check(test_name: str) -> None: 5. The profiling dump event does not have associated op types. """ + # Global mapping of Neutron kernel IDs to names used by the delegate metadata parser. + kernel_kinds = {} + def parse_delegate_metadata( delegate_metadatas: list[bytes], ) -> Union[list[str], dict[str, Any]]: """Metadata parser for Neutron Backend metadata. - The parser is a callable that deserializes the data and returns neutron kernel number. + The parser deserializes delegate metadata and converts kernel IDs into human-readable kernel names when available. The deserialized data is then added back to the corresponding event in the event block for user consumption. """ @@ -81,7 +92,13 @@ def parse_delegate_metadata( if function_code == 0: metadata_list.append("Profiling dump") else: - metadata_list.append("Neutron kernel " + str(function_code)) + metadata_list.append( + kernel_kinds.get( + function_code, "Neutron kernel " + str(function_code) + ) + ) + elif len(metadata_bytes) == 2: + metadata_list.append("Profiling dump") else: metadata_list.append("Invalid metadata size") return metadata_list @@ -96,6 +113,16 @@ def parse_delegate_metadata( file_path ), f"Required profiling file does not exist: {file_path}" + # Validate driver/compiler version compatibility and load kernel names + # used to decode delegate metadata. + driver_version = get_neutron_driver_version(etdump_path) + compiler_version = get_neutron_compiler_version() + if driver_version: + assert ( + driver_version == compiler_version + ), "Driver and compiler versions do not match" + kernel_kinds = get_neutron_kernel_kinds() + # Create Inspector and parse profiling data. try: inspector = Inspector( @@ -123,18 +150,22 @@ def parse_delegate_metadata( assert numeric_events, "No numeric delegate profiling events found" - # All delegate events except the last one should describe - # individual Neutron kernels. + # All numeric delegate events except the last contain either + # resolved kernel names or fallback "Neutron kernel " metadata. for event in numeric_events[:-1]: - metadata = str(event.delegate_debug_metadatas) - - assert "Neutron kernel" in metadata, ( - f"Event {event.name}: expected 'Neutron kernel', " f"got {metadata}" - ) + metadata = event.delegate_debug_metadatas + if kernel_kinds: + assert "Neutron kernel" not in metadata, ( + f"Event {event.name}: expected kernel kind, " f"got {metadata}" + ) + else: + assert "Neutron kernel" in metadata, ( + f"Event {event.name}: expected 'Neutron kernel', " f"got {metadata}" + ) # The final numeric event should represent the profiling dump. profiling_dump_event = numeric_events[-1] - profiling_metadata = str(profiling_dump_event.delegate_debug_metadatas) + profiling_metadata = profiling_dump_event.delegate_debug_metadatas assert "Profiling dump" in profiling_metadata, ( f"Event {profiling_dump_event.name}: " @@ -142,7 +173,7 @@ def parse_delegate_metadata( ) # Profiling dump event is expected to have no associated operators. - assert profiling_dump_event.op_types == [], ( + assert not profiling_dump_event.op_types, ( f"Event {profiling_dump_event.name}: expected empty op_types, " f"got {profiling_dump_event.op_types}" ) diff --git a/backends/nxp/tests/generic_tests/test_qdq_clustering_conv.py b/backends/nxp/tests/generic_tests/test_qdq_clustering_conv.py index 6db55347452..b0a25c68dd3 100644 --- a/backends/nxp/tests/generic_tests/test_qdq_clustering_conv.py +++ b/backends/nxp/tests/generic_tests/test_qdq_clustering_conv.py @@ -4,7 +4,7 @@ # LICENSE file in the root directory of this source tree. from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.models import Conv2dModule +from executorch.backends.nxp.tests.simple_models import Conv2dModule def test_conv2d_partitioner(): diff --git a/backends/nxp/tests/generic_tests/test_quantized_input_data.py b/backends/nxp/tests/generic_tests/test_quantized_input_data.py index a9f9f3e47e6..23a39ac11de 100644 --- a/backends/nxp/tests/generic_tests/test_quantized_input_data.py +++ b/backends/nxp/tests/generic_tests/test_quantized_input_data.py @@ -5,16 +5,16 @@ import executorch.backends.nxp.tests.nsys_testing as nsys_testing import torch +from executorch.backends.nxp.backend.ops_aliases import AvgPool2D, MulTensor from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import AvgPool2dModule, MulTensorModule from executorch.backends.nxp.tests.nsys_testing import ( lower_run_compare, OUTPUTS_DIR, ReferenceModel, ) -from executorch.backends.nxp.tests.ops_aliases import AvgPool2D, MulTensor +from executorch.backends.nxp.tests.simple_models import AvgPool2dModule, MulTensorModule def test__single_quantized_inputs(mocker, request): diff --git a/backends/nxp/tests/generic_tests/test_quantizer.py b/backends/nxp/tests/generic_tests/test_quantizer.py index 6180d2fd9ae..886301fb149 100644 --- a/backends/nxp/tests/generic_tests/test_quantizer.py +++ b/backends/nxp/tests/generic_tests/test_quantizer.py @@ -9,8 +9,10 @@ from copy import deepcopy import executorch.backends.nxp.tests.executorch_pipeline as executorch_pipeline -import executorch.backends.nxp.tests.models as models +import executorch.backends.nxp.tests.simple_models as models import numpy as np + +# noinspection PyUnusedImports import pytest import torch @@ -18,6 +20,18 @@ EdgeProgramToIRConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddMM, + Convolution, + HardTanh, + MM, + NativebatchNormLegitNoStats, + NativebatchNormLegitNoTraining, + Relu, + Sigmoid, + Tanh, +) + from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer from executorch.backends.nxp.tests.executorch_pipeline import ( neutron_target_spec, @@ -31,8 +45,6 @@ ToChannelLastPreprocess, ) -from executorch.exir.dialects._ops import ops as exir_ops - requires_tflite = pytest.mark.skipif( tflite is None, reason="tensorflow/tflite not available" ) @@ -50,13 +62,13 @@ ) fuse_activation_ops = [ - exir_ops.edge.aten.addmm.default, - exir_ops.edge.aten.mm.default, - exir_ops.edge.aten.convolution.default, - exir_ops.edge.aten.hardtanh.default, - exir_ops.edge.aten.relu.default, - exir_ops.edge.aten.sigmoid.default, - exir_ops.edge.aten.tanh.default, + AddMM, + MM, + Convolution, + HardTanh, + Relu, + Sigmoid, + Tanh, ] @@ -74,8 +86,8 @@ ] batch_norm_ops = ( - exir_ops.edge.aten._native_batch_norm_legit.no_stats, - exir_ops.edge.aten._native_batch_norm_legit_no_training.default, + NativebatchNormLegitNoStats, + NativebatchNormLegitNoTraining, torch.ops.aten._native_batch_norm_legit_no_training.default, torch.ops.aten.batch_norm.default, torch.ops.aten.native_batch_norm.default, diff --git a/backends/nxp/tests/generic_tests/test_recipe_export.py b/backends/nxp/tests/generic_tests/test_recipe_export.py new file mode 100644 index 00000000000..b853a5aa560 --- /dev/null +++ b/backends/nxp/tests/generic_tests/test_recipe_export.py @@ -0,0 +1,733 @@ +# Copyright 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 pytest +import torch +import torch.nn + +from executorch.backends.nxp.aten_passes.fuse_batch_norm_with_linear_pass import ( + FuseBatchNormWithLinearPass, +) +from executorch.backends.nxp.aten_passes.simulated_linear_bn_fusion_passes import ( + AddSimulatedLinearBatchNormFusionQATPass, + RemoveSimulatedLinearBatchNormFusionQATPass, +) +from executorch.backends.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) +from executorch.backends.nxp.backend.graph_utils import batch_norm_target_ops +from executorch.backends.nxp.backend.ops_aliases import ExecutorchDelegateCall +from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import ( + NeutronEdgePassManager, +) +from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner +from executorch.backends.nxp.recipes.nxp_recipe_provider import ( + _histogram_observer_fix_pass, + NEUTRON_RECIPE_CONFIG_KEY, + NeutronRecipeConfig, + NXPRecipeProvider, +) +from executorch.backends.nxp.recipes.nxp_recipe_types import NXPRecipeType +from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec +from executorch.backends.nxp.tests.executors import ( + graph_contains_any, + graph_contains_any_of_ops, +) +from executorch.backends.nxp.tests.simple_models import ConvBatchNormModule +from executorch.backends.transforms.quantize_fused_convbn_bias_pass import ( + QuantizeFusedConvBnBiasAtenPass, +) +from executorch.export import export +from executorch.export.recipe import ExportRecipe +from torch._inductor.lowering import quantized_decomposed + + +class SimpleCNN(torch.nn.Module): + def __init__(self, channels=3): + super().__init__() + self.conv = torch.nn.Conv2d(channels, channels, kernel_size=3) + + def forward(self, x): + x = self.conv(x) + x = torch.relu(x) + x = x.reshape(1, -1) + x = x + x + return x + + +INPUT_SHAPE = (1, 3, 8, 8) + + +def _run_export( + model, rc, recipe_type=NXPRecipeType.INT8_PTQ_NEUTRON, input_shape=INPUT_SHAPE +): + example_inputs = [(torch.randn(input_shape),)] + recipe = NXPRecipeProvider().create_recipe(recipe_type, neutron_recipe_config=rc) + return export(model, example_inputs=example_inputs, export_recipe=recipe) + + +def _get_graph(sess): + return sess.get_edge_program_manager().exported_program().graph + + +def test_ptq_neutron_basic(): + """Baseline PTQ: whole model delegated, IO is quantized.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE) + sess = _run_export(model, rc) + graph = _get_graph(sess) + + assert graph_contains_any_of_ops(graph, [ExecutorchDelegateCall]) + + def is_cnn_op(n): + return any(op in n.name.lower() for op in ["conv", "relu", "view", "add"]) + + assert not graph_contains_any(graph, is_cnn_op) + + nodes = list(graph.nodes) + # Skip alloc nodes (e.g. "alloc", "alloc_1") which also have op == "call_function". + first_call = next( + n for n in nodes if n.op == "call_function" and not n.name.startswith("alloc") + ) + last_call = next(n for n in reversed(nodes) if n.op == "call_function") + assert first_call.target == quantized_decomposed.quantize_per_tensor.out + assert last_call.target == quantized_decomposed.dequantize_per_tensor.out + + +class TestInt8PTQNoDelegate: + + def test__basic(self): + """INT8_PTQ_NO_DELEGATE: model is quantized but no delegate call is present in the graph.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE) + sess = _run_export(model, rc, recipe_type=NXPRecipeType.INT8_PTQ_NO_DELEGATE) + graph = _get_graph(sess) + + assert not graph_contains_any_of_ops(graph, [ExecutorchDelegateCall]) + + def is_cnn_op(n): + return any(op in n.name.lower() for op in ["conv", "relu", "view", "add"]) + + # With no delegation, original ops should be visible in the graph. + assert graph_contains_any(graph, is_cnn_op) + + +class TestNeutronRecipeConfigFlags: + def test_operators_not_to_delegate(self): + """Ops listed in operators_not_to_delegate are not lowered to Neutron.""" + model = SimpleCNN() + rc = NeutronRecipeConfig( + INPUT_SHAPE, operators_not_to_delegate=["aten::convolution"] + ) + sess = _run_export(model, rc) + graph = _get_graph(sess) + + assert graph_contains_any_of_ops( + graph, [torch.ops.aten.convolution.out] + ) # Convolution was not delegated. + assert graph_contains_any_of_ops( + graph, [ExecutorchDelegateCall] + ) # Other operators were delegated. + + def _is_relu_add_or_view(n: torch.fx.Node) -> bool: + return any(op in n.name.lower() for op in ["relu", "add", "view"]) + + assert not graph_contains_any(graph, _is_relu_add_or_view) + + def test_remove_quant_io_ops(self): + """remove_quant_io_ops=True: no quantize op at the IO boundary.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, remove_quant_io_ops=True) + sess = _run_export(model, rc) + graph = _get_graph(sess) + nodes = list(graph.nodes) + + real_nodes = [n for n in nodes if n.op not in ("placeholder", "output")] + assert real_nodes[0].target != quantized_decomposed.quantize_per_tensor.out + assert real_nodes[-1].target != quantized_decomposed.dequantize_per_tensor.out + assert real_nodes[-1].meta["val"].dtype == torch.int8 + placeholder_nodes = [n for n in nodes if n.op == "placeholder"] + assert placeholder_nodes[0].name == "x" # Main input + assert placeholder_nodes[0].meta["val"].dtype == torch.int8 + + def test_use_quant_state_dict_false(self, mocker): + """use_quant_state_dict=False: the NeutronPartitioner used during lowering has + post_quantization_state_dict=None, confirmed by intercepting the constructor.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, use_quant_state_dict=False) + + captured = [] + original_init = NeutronPartitioner.__init__ + + def capturing_init(self_, *args, **kwargs): + original_init(self_, *args, **kwargs) + captured.append(self_) + + mocker.patch.object(NeutronPartitioner, "__init__", capturing_init) + + _run_export(model, rc) + + assert ( + len(captured) == 1 + ), "Expected exactly one NeutronPartitioner to be created." + assert captured[0].post_quantization_state_dict is None + + def test_custom_delegation_options_explicit(self, mocker): + """Explicitly provided CustomDelegationOptions are forwarded to NeutronPartitioner.""" + model = SimpleCNN() + opts = CustomDelegationOptions() + rc = NeutronRecipeConfig(INPUT_SHAPE, custom_delegation_options=opts) + + captured = [] + original_init = NeutronPartitioner.__init__ + + def capturing_init(self_, *args, **kwargs): + original_init(self_, *args, **kwargs) + captured.append(self_) + + mocker.patch.object(NeutronPartitioner, "__init__", capturing_init) + _run_export(model, rc) + + assert len(captured) == 1 + assert captured[0].custom_delegation_options == opts + + def test_intermediates_dir(self, tmp_path): + """intermediates_dir: intermediate compilation files are written to the directory.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, intermediates_dir=str(tmp_path)) + _run_export(model, rc) + assert any( + tmp_path.iterdir() + ), "No intermediate files written to intermediates_dir." + + def test_fetch_constants_to_sram_flag(self, mocker): + """fetch_constants_to_sram=True reaches the NeutronPartitioner used during export.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, fetch_constants_to_sram=True) + + captured = [] + original_init = NeutronPartitioner.__init__ + + def capturing_init(self_, *args, **kwargs): + original_init(self_, *args, **kwargs) + captured.append(self_) + + mocker.patch.object(NeutronPartitioner, "__init__", capturing_init) + _run_export(model, rc) + + assert ( + len(captured) == 1 + ), "Expected exactly one NeutronPartitioner to be created." + spec_map = {s.key: s.value.decode() for s in captured[0].delegation_spec[1]} + assert spec_map["fetch_constants_to_sram"] == "True" + + def test_use_profiling_flag(self, mocker): + """use_profiling=True reaches the NeutronPartitioner used during export.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, use_profiling=True) + + captured = [] + original_init = NeutronPartitioner.__init__ + + def capturing_init(self_, *args, **kwargs): + original_init(self_, *args, **kwargs) + captured.append(self_) + + mocker.patch.object(NeutronPartitioner, "__init__", capturing_init) + _run_export(model, rc) + + assert ( + len(captured) == 1 + ), "Expected exactly one NeutronPartitioner to be created." + spec_map = {s.key: s.value.decode() for s in captured[0].delegation_spec[1]} + assert spec_map["use_profiling"] == "True" + + def test_dump_kernel_selection_code(self, tmp_path, monkeypatch): + """dump_kernel_selection_code=True causes a kernel selection C file to be written.""" + monkeypatch.chdir(tmp_path) + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, dump_kernel_selection_code=True) + _run_export(model, rc) + assert ( + tmp_path / "_kernel_selection.c" + ).exists(), "_kernel_selection.c was not created in the working directory." + + def test_custom_quantizer_fn(self): + """get_quantizer_fn overrides the default NeutronQuantizer.""" + from executorch.backends.nxp.backend.neutron_target_spec import ( + NeutronTargetSpec, + ) + from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer + + custom_quantizer_called = [] + + def my_quantizer_fn(): + q = NeutronQuantizer(NeutronTargetSpec("imxrt700")) + custom_quantizer_called.append(True) + return q + + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, get_quantizer_fn=my_quantizer_fn) + sess = _run_export(model, rc) + assert custom_quantizer_called, "Custom quantizer factory was not called." + assert sess.get_edge_program_manager() is not None + + def test_use_neutron_for_format_conversion_false(self): + """use_neutron_for_format_conversion=False still produces a valid export.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, use_neutron_for_format_conversion=False) + sess = _run_export(model, rc) + assert sess.get_edge_program_manager() is not None + + def test_target_explicit(self): + """Specifying fake target to make sure an error is raised.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, target="FAKE") + with pytest.raises(ValueError, match="`FAKE` is not a valid target"): + _run_export(model, rc) + + +class TestInputSpecForms: + def test__single_tuple(self): + """input_spec as a plain shape tuple works.""" + model = SimpleCNN() + sess = _run_export(model, NeutronRecipeConfig((1, 3, 8, 8))) + assert sess.get_edge_program_manager() is not None + + def test__list_of_tuples(self): + """input_spec as list of shape tuples works.""" + model = SimpleCNN() + sess = _run_export(model, NeutronRecipeConfig([(1, 3, 8, 8)])) + assert sess.get_edge_program_manager() is not None + + def test__model_input_spec(self): + """input_spec as list of ModelInputSpec objects works.""" + model = SimpleCNN() + sess = _run_export(model, NeutronRecipeConfig([ModelInputSpec((1, 3, 8, 8))])) + assert sess.get_edge_program_manager() is not None + + def test__multi_input(self): + """input_spec with multiple inputs (two tensors) works.""" + + class AddModel(torch.nn.Module): + def forward(self, x, y): + return x + y + + model = AddModel() + rc = NeutronRecipeConfig([(1, 3, 8, 8), (1, 3, 8, 8)]) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + example_inputs = [(torch.randn(1, 3, 8, 8), torch.randn(1, 3, 8, 8))] + sess = export(model, example_inputs=example_inputs, export_recipe=recipe) + assert sess.get_edge_program_manager() is not None + + +class TestErrorHandling: + def test_create_recipe_missing_config_key(self): + """create_recipe without neutron_recipe_config kwarg raises KeyError.""" + with pytest.raises(KeyError, match=NEUTRON_RECIPE_CONFIG_KEY): + NXPRecipeProvider().create_recipe(NXPRecipeType.INT8_PTQ_NEUTRON) + + def test_create_recipe_invalid_recipe_type(self): + """create_recipe with an unsupported recipe type returns None with a warning.""" + from executorch.export.recipe import RecipeType + + class FakeRecipeType(RecipeType): + FAKE = "fake" + + @classmethod + def get_backend_name(cls): + return "fake_backend" + + rc = NeutronRecipeConfig(INPUT_SHAPE) + result = NXPRecipeProvider().create_recipe( + FakeRecipeType.FAKE, neutron_recipe_config=rc + ) + assert result is None + + +class TestRecipeStructureValidation: + + def test_ptq_neutron_recipe_structure(self): + """INT8_PTQ_NEUTRON recipe: correct quantizer and partitioner are set.""" + rc = NeutronRecipeConfig(INPUT_SHAPE) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + assert recipe.quantization_recipe is not None + assert len(recipe.quantization_recipe.quantizers) == 1 + assert recipe.lowering_recipe.partitioners is not None + assert len(recipe.lowering_recipe.partitioners) == 1 + + def test_ptq_neutron_recipe_name(self): + """INT8_PTQ_NEUTRON recipe has the expected name.""" + rc = NeutronRecipeConfig(INPUT_SHAPE) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + assert recipe.name == NXPRecipeType.INT8_PTQ_NEUTRON.value + + +class TestRecipeCombination: + def test__chains_pre_partitioning_callbacks(self): + """Combining two NXP recipes chains both pre_partitioning_callbacks.""" + recipe1 = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, + neutron_recipe_config=NeutronRecipeConfig(INPUT_SHAPE), + ) + recipe2 = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, + neutron_recipe_config=NeutronRecipeConfig(INPUT_SHAPE), + ) + combined = ExportRecipe.combine([recipe1, recipe2]) + assert combined.lowering_recipe.pre_partitioning_callback is not None + # Calling the combined callback should not raise. + combined.lowering_recipe.pre_partitioning_callback(None, {}) + + +class TestEdgeManagerTransformPasses: + def test__executed(self): + """edge_manager_transform_passes are called after partitioning.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + + transform_called = [] + + def tracking_pass(epm): + transform_called.append(True) + return [] + + recipe.lowering_recipe.edge_manager_transform_passes = [tracking_pass] + + example_inputs = [(torch.randn(INPUT_SHAPE),)] + export(model, example_inputs=example_inputs, export_recipe=recipe) + assert transform_called, "edge_manager_transform_passes were not executed." + + def test__qdq_pass_callable_returns_pass_manager(self, mocker): + """_remove_additional_qdq_clusters returns a bare NeutronEdgePassManager, not a + list containing one. EdgeProgramManagerTransformStage calls epm.transform(passes) + directly, so a list-of-PassManager would be silently mis-applied.""" + rc = NeutronRecipeConfig(INPUT_SHAPE) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + # remove_quant_io_ops=False (default): first callable is _remove_additional_qdq_clusters. + qdq_callable = recipe.lowering_recipe.edge_manager_transform_passes[0] + result = qdq_callable(mocker.MagicMock()) + assert isinstance(result, NeutronEdgePassManager) + + +# --------------------------------------------------------------------------- +# Helpers shared by QAT tests +# --------------------------------------------------------------------------- + + +def _noop_train_fn(model: torch.fx.GraphModule) -> None: + """A no-op train_fn used by structural/unit tests that only inspect pass shape.""" + pass + + +def _minimal_train_fn(model: torch.fx.GraphModule, shape=(1, 3, 5, 5)) -> None: + """Run a few SGD steps on random data so fake-quant observer statistics are populated. + Used by end-to-end tests. + """ + optimizer = torch.optim.SGD(model.parameters(), lr=1e-4) + for _ in range(3): + optimizer.zero_grad() + out = model(torch.randn(shape)) + loss = out.sum() + loss.backward() + optimizer.step() + + +def _run_qat_export(model, train_fn=None, recipe_type=NXPRecipeType.INT8_QAT_NEUTRON): + if train_fn is None: + train_fn = _noop_train_fn + rc = NeutronRecipeConfig(INPUT_SHAPE, train_fn=train_fn) + return _run_export(model, rc, recipe_type=recipe_type) + + +# --------------------------------------------------------------------------- +# QAT recipe: end-to-end tests +# --------------------------------------------------------------------------- + + +# Both QAT recipe types must reject a missing train_fn. +@pytest.mark.parametrize( + "recipe_type", + [NXPRecipeType.INT8_QAT_NEUTRON, NXPRecipeType.INT8_QAT_NO_DELEGATE], + ids=lambda r: r.value, +) +def test__qat_requires_train_fn(recipe_type): + """Any QAT recipe raises ValueError when train_fn is absent from NeutronRecipeConfig.""" + rc = NeutronRecipeConfig(INPUT_SHAPE) # train_fn=None (default) + with pytest.raises(ValueError, match="train_fn"): + NXPRecipeProvider().create_recipe(recipe_type, neutron_recipe_config=rc) + + +# Both _NO_DELEGATE recipe types (PTQ and QAT) must produce an empty partitioner list. +@pytest.mark.parametrize( + "recipe_type", + [NXPRecipeType.INT8_PTQ_NO_DELEGATE, NXPRecipeType.INT8_QAT_NO_DELEGATE], + ids=lambda r: r.value, +) +def test__no_delegate_recipe_has_empty_partitioners(recipe_type): + """Both NO_DELEGATE recipe types produce an empty partitioner list.""" + # train_fn is required by QAT recipes; PTQ ignores it, so always pass it. + rc = NeutronRecipeConfig(INPUT_SHAPE, train_fn=_noop_train_fn) + recipe = NXPRecipeProvider().create_recipe(recipe_type, neutron_recipe_config=rc) + assert recipe.lowering_recipe.partitioners == [] + + +class TestInt8QATNeutron: + + def test__basic(self): + """INT8_QAT_NEUTRON: full export succeeds and the graph contains a delegate call.""" + model = SimpleCNN() + sess = _run_qat_export(model) + graph = _get_graph(sess) + assert graph_contains_any_of_ops(graph, [ExecutorchDelegateCall]) + + def test__train_fn_is_called(self): + """train_fn is invoked exactly once during the QAT export pipeline.""" + model = SimpleCNN() + call_count = [] + + def counting_train_fn(m): + call_count.append(1) + + rc = NeutronRecipeConfig(INPUT_SHAPE, train_fn=counting_train_fn) + _run_export(model, rc, recipe_type=NXPRecipeType.INT8_QAT_NEUTRON) + assert ( + len(call_count) == 1 + ), f"Expected train_fn called once, got {len(call_count)}" + + def test__recipe_name(self): + """INT8_QAT_NEUTRON recipe has the expected name.""" + rc = NeutronRecipeConfig(INPUT_SHAPE, train_fn=_noop_train_fn) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_QAT_NEUTRON, neutron_recipe_config=rc + ) + assert recipe.name == NXPRecipeType.INT8_QAT_NEUTRON.value + + def test__recipe_structure(self): + """INT8_QAT_NEUTRON recipe has is_qat=True, one quantizer, one partitioner.""" + rc = NeutronRecipeConfig(INPUT_SHAPE, train_fn=_noop_train_fn) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_QAT_NEUTRON, neutron_recipe_config=rc + ) + qr = recipe.quantization_recipe + assert qr is not None + assert qr.is_qat is True + assert qr.train_fn is _noop_train_fn + assert len(qr.quantizers) == 1 + assert recipe.lowering_recipe.partitioners is not None + assert len(recipe.lowering_recipe.partitioners) == 1 + + def test__io_is_quantized_by_default(self): + """QAT export with default settings: IO boundary has quantize/dequantize ops.""" + model = SimpleCNN() + sess = _run_qat_export(model) + graph = _get_graph(sess) + nodes = list(graph.nodes) + # Skip alloc nodes (e.g. "alloc", "alloc_1") which also have op == "call_function". + first_call = next( + n + for n in nodes + if n.op == "call_function" and not n.name.startswith("alloc") + ) + last_call = next(n for n in reversed(nodes) if n.op == "call_function") + assert first_call.target == quantized_decomposed.quantize_per_tensor.out + assert last_call.target == quantized_decomposed.dequantize_per_tensor.out + + +class TestInt8QATNoDelegate: + + def test__basic(self): + """INT8_QAT_NO_DELEGATE: export succeeds without any delegate call.""" + model = SimpleCNN() + sess = _run_qat_export(model, recipe_type=NXPRecipeType.INT8_QAT_NO_DELEGATE) + graph = _get_graph(sess) + assert not graph_contains_any_of_ops(graph, [ExecutorchDelegateCall]) + + +# --------------------------------------------------------------------------- +# QAT recipe: NXP-specific pass structure tests +# --------------------------------------------------------------------------- + +# Both QAT recipe types are built from the same _build_quantization_recipe(is_qat=True) +# call, so their pass lists must be identical. The parametrization below makes this +# explicit and catches any accidental divergence. +_QAT_RECIPE_TYPES = [NXPRecipeType.INT8_QAT_NEUTRON, NXPRecipeType.INT8_QAT_NO_DELEGATE] + + +@pytest.mark.parametrize("recipe_type", _QAT_RECIPE_TYPES, ids=lambda r: r.value) +class TestQATNXPPasses: + + def _get_qat_recipe(self, recipe_type: NXPRecipeType) -> "ExportRecipe": + rc = NeutronRecipeConfig(INPUT_SHAPE, train_fn=_noop_train_fn) + return NXPRecipeProvider().create_recipe(recipe_type, neutron_recipe_config=rc) + + def test__post_prepare_passes_start_with_add_bn_fusion(self, recipe_type): + """QAT post_prepare_passes: first pass is AddSimulatedLinearBatchNormFusionQATPass wrapper.""" + recipe = self._get_qat_recipe(recipe_type) + qr = recipe.quantization_recipe + assert qr.post_prepare_passes is not None + # The first post-prepare pass must wrap AddSimulatedLinearBatchNormFusionQATPass. + # We verify by inspecting the __qualname__ set by _wrap_exir_pass. + first_pass = qr.post_prepare_passes[0] + assert ( + AddSimulatedLinearBatchNormFusionQATPass.__name__ in first_pass.__qualname__ + ) + + def test__post_prepare_passes_end_with_histogram_observer_fix(self, recipe_type): + """QAT post_prepare_passes: last pass is _histogram_observer_fix_pass.""" + recipe = self._get_qat_recipe(recipe_type) + qr = recipe.quantization_recipe + assert qr.post_prepare_passes is not None + last_pass = qr.post_prepare_passes[-1] + assert last_pass is _histogram_observer_fix_pass + + def test__pre_convert_passes_include_remove_bn_fusion_and_fold(self, recipe_type): + """QAT pre_convert_passes: contains RemoveSimulatedLinearBatchNormFusionQATPass + followed by FuseBatchNormWithLinearPass (each applied once).""" + recipe = self._get_qat_recipe(recipe_type) + qr = recipe.quantization_recipe + assert qr.pre_convert_passes is not None + assert len(qr.pre_convert_passes) == 2, ( + "Expected 2 pre_convert passes (remove + fuse), " + f"got {len(qr.pre_convert_passes)}" + ) + qualnames = [p.__qualname__ for p in qr.pre_convert_passes] + assert ( + qualnames[0] + == f"_wrap_exir_pass({RemoveSimulatedLinearBatchNormFusionQATPass.__name__})" + ) + assert ( + qualnames[1] == f"_wrap_exir_pass({FuseBatchNormWithLinearPass.__name__})" + ) + + def test__post_convert_passes_include_quant_fused_conv_bn_bias(self, recipe_type): + """QAT post_convert_passes: contains QuantizeFusedConvBnBiasAtenPass wrapper.""" + recipe = self._get_qat_recipe(recipe_type) + qr = recipe.quantization_recipe + assert qr.post_convert_passes is not None + assert len(qr.post_convert_passes) == 1 + assert f"_wrap_exir_pass({QuantizeFusedConvBnBiasAtenPass.__name__})" in ( + qr.post_convert_passes[0].__qualname__ + ) + + +# --------------------------------------------------------------------------- +# PTQ recipe: pass structure tests (complement to TestQATNXPPasses above) +# --------------------------------------------------------------------------- + + +class TestPTQNXPPasses: + """Verifies the pass structure of INT8_PTQ_NEUTRON recipes. + + These tests are separate from TestQATNXPPasses because the assertions are + PTQ-specific and independent of which QAT recipe type is being tested. + """ + + def _get_ptq_recipe(self) -> "ExportRecipe": + rc = NeutronRecipeConfig(INPUT_SHAPE) + return NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + + def test__no_pre_or_post_convert_passes(self): + """PTQ recipe does not set pre_convert_passes or post_convert_passes.""" + qr = self._get_ptq_recipe().quantization_recipe + assert qr.pre_convert_passes is None + assert qr.post_convert_passes is None + + def test__post_prepare_passes_include_histogram_observer_fix(self): + """PTQ recipe post_prepare_passes contains _histogram_observer_fix_pass.""" + qr = self._get_ptq_recipe().quantization_recipe + assert qr.post_prepare_passes is not None + assert _histogram_observer_fix_pass in qr.post_prepare_passes + + def test__post_prepare_passes_do_not_include_add_bn_fusion(self): + """PTQ recipe post_prepare_passes must NOT contain AddSimulatedLinearBatchNormFusionQATPass.""" + qr = self._get_ptq_recipe().quantization_recipe + for p in qr.post_prepare_passes or []: + assert AddSimulatedLinearBatchNormFusionQATPass.__name__ not in getattr( + p, "__qualname__", "" + ) + + +# --------------------------------------------------------------------------- +# QAT recipe: e2e test on a real model - equivalent to imperative QAT tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "bias", [True, False], ids=lambda b: "bias" if b else "no_bias" +) +class TestQATEquivalentToImperative: + """Recipe-path QAT tests that mirror the imperative-path tests in test_batch_norm_fusion.py. + + The imperative reference is test_biasless_convbn_fusion_qat (and its bias=True variant). + The recipe path must produce an equivalent result: the graph is delegated and the + BN is fully fused away by the QAT passes. + """ + + # Use (1, 3, 5, 5) so that Conv2d(kernel_size=3) produces a (1, 3, 3, 3) feature map, + # giving BatchNorm > 1 value per channel in training mode (QAT requires train mode). + _CONVBN_INPUT_SHAPE = (1, 3, 5, 5) + + def test__convbn_qat_produces_delegate_call(self, bias): + """INT8_QAT_NEUTRON on ConvBatchNormModule produces a delegate call. + Equivalent imperative test: test_biasless_convbn_fusion_qat / test_batch_norm_conv_fusing + in backends/nxp/tests/generic_tests/test_batch_norm_fusion.py.""" + model = ConvBatchNormModule( + bias=bias, + input_rank=len(self._CONVBN_INPUT_SHAPE), + num_features=self._CONVBN_INPUT_SHAPE[1], + ) + rc = NeutronRecipeConfig( + self._CONVBN_INPUT_SHAPE, + train_fn=_minimal_train_fn, + use_neutron_for_format_conversion=False, + ) + sess = _run_export( + model, + rc, + recipe_type=NXPRecipeType.INT8_QAT_NEUTRON, + input_shape=self._CONVBN_INPUT_SHAPE, + ) + graph = _get_graph(sess) + + # Same assertion as the imperative path: the model is delegated. + assert graph_contains_any_of_ops(graph, [ExecutorchDelegateCall]) + + def test__convbn_qat_bn_is_fused_away(self, bias): + """INT8_QAT_NEUTRON on ConvBatchNormModule: BN is fused away by QAT passes. + Equivalent imperative test: test_batch_norm_conv_fusing__full_pipeline__2d + in backends/nxp/tests/generic_tests/test_batch_norm_fusion.py.""" + model = ConvBatchNormModule( + bias=bias, + input_rank=len(self._CONVBN_INPUT_SHAPE), + num_features=self._CONVBN_INPUT_SHAPE[1], + ) + rc = NeutronRecipeConfig( + self._CONVBN_INPUT_SHAPE, + train_fn=_minimal_train_fn, + use_neutron_for_format_conversion=False, + ) + sess = _run_export( + model, + rc, + recipe_type=NXPRecipeType.INT8_QAT_NEUTRON, + input_shape=self._CONVBN_INPUT_SHAPE, + ) + # The edge program (before delegation) must not contain any BN ops. + edge_graph = sess.get_edge_program_manager().exported_program().graph + assert not graph_contains_any_of_ops(edge_graph, batch_norm_target_ops) diff --git a/backends/nxp/tests/generic_tests/test_removing_dead_code.py b/backends/nxp/tests/generic_tests/test_removing_dead_code.py index 8b3a979f412..7c01be7832d 100644 --- a/backends/nxp/tests/generic_tests/test_removing_dead_code.py +++ b/backends/nxp/tests/generic_tests/test_removing_dead_code.py @@ -6,6 +6,8 @@ import unittest import numpy as np + +# noinspection PyUnusedImports import pytest import torch 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 12d2f193f57..a6bd3fafdc9 100644 --- a/backends/nxp/tests/generic_tests/test_split_group_convolution.py +++ b/backends/nxp/tests/generic_tests/test_split_group_convolution.py @@ -15,6 +15,7 @@ from executorch.backends.nxp.aten_passes.split_group_convolution import ( SplitGroupConvolution, ) +from executorch.backends.nxp.backend.ops_aliases import Cat, Convolution from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner from executorch.backends.nxp.nxp_backend import generate_neutron_compile_spec from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer @@ -26,13 +27,12 @@ to_quantized_edge_program, ) from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.simple_models import ( Conv1dModule, Conv2dModule, Conv3dModule, ) from executorch.exir import EdgeCompileConfig, EdgeProgramManager -from executorch.exir.dialects._ops import ops as exir_ops from executorch.extension.export_util import export_to_edge from parameterized import parameterized from torch.fx import GraphModule @@ -129,7 +129,7 @@ def test_split_group_convolution__2d( assert nodes[-5].name == "lowered_module_0" assert not graph_contains_any_of_ops( ep.graph, - [exir_ops.edge.aten.convolution.default, exir_ops.edge.aten.cat.default], + [Convolution, Cat], ) @parameterized.expand( @@ -206,7 +206,7 @@ def test_split_group_convolution__1d( assert nodes[-5].name == "lowered_module_0" assert not graph_contains_any_of_ops( ep.graph, - [exir_ops.edge.aten.convolution.default, exir_ops.edge.aten.cat.default], + [Convolution, Cat], ) @parameterized.expand( @@ -303,5 +303,5 @@ def test_split_group_convolution__applied_by_default(self, _, is_qat: bool): assert nodes[-5].name == "lowered_module_0" assert not graph_contains_any_of_ops( ep.graph, - [exir_ops.edge.aten.convolution.default, exir_ops.edge.aten.cat.default], + [Convolution, Cat], ) diff --git a/backends/nxp/tests/graph_verifier.py b/backends/nxp/tests/graph_verifier.py index 44900b6a11b..70701bbd90f 100644 --- a/backends/nxp/tests/graph_verifier.py +++ b/backends/nxp/tests/graph_verifier.py @@ -10,17 +10,17 @@ from dataclasses import dataclass from typing import Callable, Union -from executorch.backends.nxp.neutron_partitioner import ( - NeutronPartitioner, - NXP_DELEGATION_TAG, -) -from executorch.backends.nxp.tests.ops_aliases import ( +from executorch.backends.nxp.backend.ops_aliases import ( DequantizePerChannel, DequantizePerTensor, QuantizePerChannel, QuantizePerTensor, ) +from executorch.backends.nxp.neutron_partitioner import ( + NeutronPartitioner, + NXP_DELEGATION_TAG, +) from executorch.exir.dialects.edge._ops import EdgeOpOverload from pytest_mock import MockerFixture diff --git a/backends/nxp/tests/ir/converter/node_converter/test_abs_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_abs_converter.py index d42ef4c6e7d..94a0aafc2a7 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_abs_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_abs_converter.py @@ -8,12 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Abs from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.nsys_testing import ( lower_run_compare, RandomDatasetCreator, ) -from executorch.backends.nxp.tests.ops_aliases import Abs from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py index 9646c04a3f2..75093aba530 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py @@ -8,6 +8,11 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AdaptiveAvgPool2D, + ExecutorchDelegateCall, + ViewCopy, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program @@ -16,11 +21,10 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import AdaptiveAvgPool2dModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AdaptiveAvgPool2D, - ExecutorchDelegateCall, +from executorch.backends.nxp.tests.simple_models import ( + AdaptiveAvgPool1dModule, + AdaptiveAvgPool2dModule, ) from executorch.backends.nxp.tests.use_qat import * # noqa F403 @@ -42,6 +46,11 @@ class TestAdaptiveAvgPool2D: (2, 3), id="H != W, non multiples of num_macs, batch != 1.", ), + pytest.param( + (2, 3, 10, 15), + (5, 5), + id="H != W, non multiples of num_macs, batch != 1, fixed fail.", + ), ], ) def test__basic_nsys_inference( @@ -54,9 +63,8 @@ def test__basic_nsys_inference( expected_non_delegated_ops={}, ) - output_comparator = AllCloseOutputComparator( - 7.84e-3 - ) # Accept small error due to Neutron bug (AIR-14585). + remove_quant_io_ops = True # Use quantized dataset. + output_comparator = AllCloseOutputComparator(atol=1) # Allow single bit error. lower_run_compare( model, @@ -66,32 +74,7 @@ def test__basic_nsys_inference( RandomDatasetCreator(low=-1, high=1), output_comparator=output_comparator, use_qat=use_qat, - ) - - @pytest.mark.xfail( - strict=True, - reason="Known Neutron bad compute issue. Will be fixed in Neutron SW 3.1.2.", - ) - def test__know_neutron_issue(self, mocker, request): - input_shape = (2, 3, 10, 15) - output_size = (5, 5) - model = AdaptiveAvgPool2dModule(output_size) - graph_verifier = DetailedGraphVerifier( - mocker, - expected_delegated_ops={AdaptiveAvgPool2D: 1}, - expected_non_delegated_ops={}, - ) - - # Use high tolerance so we notice when the issue is fixed. - output_comparator = AllCloseOutputComparator(7.8e-3) - - lower_run_compare( - model, - input_shape, - graph_verifier, - request, - RandomDatasetCreator(low=-1, high=1), - output_comparator=output_comparator, + remove_quant_io_ops=remove_quant_io_ops, ) def test__kernel_size_and_stride_limit(self, mocker, request): @@ -110,9 +93,8 @@ def test__kernel_size_and_stride_limit(self, mocker, request): expected_non_delegated_ops={}, ) - output_comparator = AllCloseOutputComparator( - 7.9e-3 - ) # Accept small error due to Neutron bug (AIR-14585). + remove_quant_io_ops = True # Use quantized dataset. + output_comparator = AllCloseOutputComparator(atol=1) # Allow single bit error. lower_run_compare( model, @@ -121,6 +103,7 @@ def test__kernel_size_and_stride_limit(self, mocker, request): request, RandomDatasetCreator(low=-1, high=1), output_comparator=output_comparator, + remove_quant_io_ops=remove_quant_io_ops, ) def test__kernel_size_and_stride_limit_exceeded(self): @@ -140,3 +123,31 @@ def test__kernel_size_and_stride_limit_exceeded(self): delegated_ep.graph, [ExecutorchDelegateCall] ) assert graph_contains_any_of_ops(delegated_ep.graph, [AdaptiveAvgPool2D]) + + +class TestAdaptiveAvgPool1DTo2D: + + # Just a basic test to verify that the operator gets extended to the 2D variant correctly. + def test__basic_nsys_inference(self, mocker, request, use_qat): + input_shape = (2, 4, 6) # The old flow limited the batch size to 1. + output_size = (3,) + model = AdaptiveAvgPool1dModule(output_size) + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={AdaptiveAvgPool2D: 1, ViewCopy: 2}, + expected_non_delegated_ops={}, + ) + + remove_quant_io_ops = True # Use quantized dataset. + output_comparator = AllCloseOutputComparator(atol=1) # Allow single bit error. + + lower_run_compare( + model, + input_shape, + graph_verifier, + request, + RandomDatasetCreator(low=-1, high=1), + output_comparator=output_comparator, + use_qat=use_qat, + remove_quant_io_ops=remove_quant_io_ops, + ) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py index c01d0ca818d..5c3b0dd0100 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py @@ -8,6 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import ( @@ -19,13 +25,10 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import AddTensorModule, MaxPoolAddTensorModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, +from executorch.backends.nxp.tests.simple_models import ( + AddTensorModule, + MaxPoolAddTensorModule, ) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_addmm_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_addmm_converter.py index 1db604d5b1e..713a618ea53 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_addmm_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_addmm_converter.py @@ -8,20 +8,20 @@ # noinspection PyUnusedImports import pytest import torch - -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 graph_contains_any_of_ops -from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier, Operator -from executorch.backends.nxp.tests.models import AddmmModule, LinearModule -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( +from executorch.backends.nxp.backend.ops_aliases import ( AddMM, ExecutorchDelegateCall, MM, PermuteCopy, ViewCopy, ) + +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 graph_contains_any_of_ops +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier, Operator +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import AddmmModule, LinearModule from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_amax_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_amax_converter.py index d348c12102e..892ea70018e 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_amax_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_amax_converter.py @@ -21,6 +21,13 @@ from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.transpose_options import ( Transpose, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + Amax, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, +) 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 graph_contains_any_of_ops @@ -29,13 +36,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - Amax, - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_amin_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_amin_converter.py index e2490d9c2c4..3ae13bafe40 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_amin_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_amin_converter.py @@ -21,6 +21,13 @@ from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.transpose_options import ( Transpose, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + Amin, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, +) 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 graph_contains_any_of_ops @@ -29,13 +36,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - Amin, - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 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 3db1158d637..bd28e8df6e6 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 @@ -8,17 +8,17 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AvgPool2D, + ExecutorchDelegateCall, + ViewCopy, +) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import AvgPool2dModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AvgPool2D, - ExecutorchDelegateCall, - ViewCopy, -) +from executorch.backends.nxp.tests.simple_models import AvgPool2dModule from executorch.backends.nxp.tests.use_qat import * # noqa F403 @@ -112,10 +112,10 @@ def test__stride_limit_exceeded(self): assert graph_contains_any_of_ops(delegated_ep.graph, [AvgPool2D]) -class TestAvgPool1D: +class TestAvgPool1DTo2D: # Just a basic test to verify that the operator gets extended to the 2D variant correctly. - def test__basic_nsys_inference(self, mocker, request): + def test__basic_nsys_inference(self, mocker, request, use_qat): input_shape = (2, 4, 6) # The old flow limited the batch size to 1. model = AvgPool1DModule() graph_verifier = DetailedGraphVerifier( @@ -124,4 +124,4 @@ def test__basic_nsys_inference(self, mocker, request): expected_non_delegated_ops={}, ) - lower_run_compare(model, input_shape, graph_verifier, request) + lower_run_compare(model, input_shape, graph_verifier, request, use_qat=use_qat) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_bmm_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_bmm_converter.py index c564c024623..5bb7a080b23 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_bmm_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_bmm_converter.py @@ -6,6 +6,11 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + BMM, + GetItem, + MaxPool2DWithIndices, +) from executorch.backends.nxp.edge_passes.move_auxiliary_operator_into_separate_qdq_cluster_pass import ( ViewCopy, @@ -16,12 +21,11 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( BatchMatMulMaxPoolModel, BatchMatMulModel, ) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import BMM, GetItem, MaxPool2DWithIndices from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py index b28a431e3ca..ac541ed0b77 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py @@ -8,6 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + Cat, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, +) from executorch.backends.nxp.tests.executorch_pipeline import ( ModelInputSpec, @@ -16,12 +22,6 @@ from executorch.backends.nxp.tests.executors import graph_contains_any_of_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 ( - Cat, - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 @@ -136,7 +136,7 @@ def test__different_shapes__channels_first(self, mocker, request, dim, num_input lower_run_compare(model, input_shapes, graph_verifier, request) def test__single_input__alone_in_partition__not_delegated(self): - # The operator is a noop, and there is no other op in the model. The Neutron Converter would produce an empty + # The operator is a noop, and there is no other op in the model. The Neutron Compiler would produce an empty # graph, so the `cat` is not delegated. input_shape = [ModelInputSpec((2, 3, 5))] model = CatModule(1) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_clamp_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_clamp_converter.py index b2147a0d984..04ae8757110 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_clamp_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_clamp_converter.py @@ -18,6 +18,11 @@ from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import ( BuiltinOperator as Ops, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + Clamp, + ExecutorchDelegateCall, +) from executorch.backends.nxp.tests.executorch_pipeline import ( ModelInputSpec, to_quantized_edge_program, @@ -25,11 +30,6 @@ from executorch.backends.nxp.tests.executors import graph_contains_any_of_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 ( - AddTensor, - Clamp, - ExecutorchDelegateCall, -) from executorch.backends.nxp.tests.use_qat import * # noqa: F403 F401 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_clone_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_clone_converter.py index 5ee3db6752f..aaf6a6e35d3 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_clone_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_clone_converter.py @@ -2,6 +2,7 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + import itertools import unittest @@ -16,6 +17,11 @@ PermuteCopyConverter, ) from executorch.backends.nxp.backend.node_format_inference import NodeFormatInference +from executorch.backends.nxp.backend.ops_aliases import ( + Clone, + CloneDimOrder, + PermuteCopy, +) from executorch.backends.nxp.edge_passes.move_auxiliary_operator_into_separate_qdq_cluster_pass import ( MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass, ) @@ -45,7 +51,6 @@ ToChannelLastPreprocess, ) from executorch.exir import EdgeCompileConfig -from executorch.exir.dialects._ops import ops as exir_ops from executorch.extension.export_util.utils import export_to_edge from parameterized import parameterized from torch import nn @@ -144,8 +149,8 @@ def setUpClass(cls): @staticmethod def _node_is_clone(node) -> bool: clone_ops = [ - exir_ops.edge.aten.clone.default, - exir_ops.edge.dim_order_ops._clone_dim_order.default, + Clone, + CloneDimOrder, ] def target_can_be_clone(node): @@ -215,13 +220,14 @@ def test_conv_dropout_no_quant( has_clone = graph_contains_any_of_ops( graph=edge_program.graph, ops=[ - exir_ops.edge.aten.clone.default, - exir_ops.edge.dim_order_ops._clone_dim_order.default, + Clone, + CloneDimOrder, ], ) - # Clone with inplace=True should not produce clone edge op and vice versa - assert inplace_dropout ^ has_clone + # Neither spelling leaves a clone behind on this PyTorch: the out-of-place + # one used to and no longer does. + assert not has_clone @parameterized.expand([("QAT", True), ("PTQ", False)]) def test_clone_pool_view_copy_quant( @@ -286,7 +292,7 @@ def test_clone__to_contiguous_format(self): ) # Make sure the `aten.clone` was inserted as expected. nodes = list(edge_program_manager.exported_program().graph.nodes) - assert nodes[9].target == exir_ops.edge.dim_order_ops._clone_dim_order.default + assert nodes[9].target == CloneDimOrder assert nodes[9].kwargs["dim_order"] == [0, 1, 2, 3] # Move the `clone` out of the cluster with the `view_copy`. @@ -339,9 +345,7 @@ def _unsupported_target(*_): ep = to_quantized_edge_program(model, input_shape).exported_program() nodes = list(ep.graph.nodes) - assert not graph_contains_any_of_ops( - ep.graph, [exir_ops.edge.aten.clone.default] - ) + assert not graph_contains_any_of_ops(ep.graph, [Clone]) assert nodes[3].name == "executorch_call_delegate" - assert nodes[5].target == exir_ops.edge.aten.permute_copy.default + assert nodes[5].target == PermuteCopy assert nodes[7].name == "executorch_call_delegate_1" 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 b4a64447aa6..b4de3c4984f 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 @@ -12,10 +12,10 @@ from executorch.backends.nxp.backend.ir.converter.builder.model_builder import ( ModelBuilder, ) +from executorch.backends.nxp.backend.ops_aliases import ConstantPadND, Convolution from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import PadConvModule, PadModule 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.simple_models import PadConvModule, PadModule from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py index 3d20d38bb54..b802309d8e0 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py @@ -4,22 +4,27 @@ # 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.ops_aliases import ( + Convolution, + ExecutorchDelegateCall, + ViewCopy, +) 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 graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import Conv2dModule, Conv2dTransposedModule from executorch.backends.nxp.tests.nsys_testing import ( AllCloseOutputComparator, lower_run_compare, ReferenceModel, ) -from executorch.backends.nxp.tests.ops_aliases import ( - Convolution, - ExecutorchDelegateCall, - ViewCopy, +from executorch.backends.nxp.tests.simple_models import ( + Conv2dModule, + Conv2dTransposedModule, ) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_exp_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_exp_converter.py index b304dce2c94..9f8068f67d4 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_exp_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_exp_converter.py @@ -8,10 +8,10 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Exp 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 Exp from executorch.backends.nxp.tests.use_qat import * # noqa F403 from executorch.backends.nxp.tests.dataset_creator import ( LinearRampDatasetCreator, diff --git a/backends/nxp/tests/ir/converter/node_converter/test_hardswish_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_hardswish_converter.py index deff4e12f0c..e0f626a2140 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_hardswish_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_hardswish_converter.py @@ -8,25 +8,25 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AddMM, + Convolution, + Hardswish, + PermuteCopy, + ViewCopy, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator 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 ( +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( ConvHardswishModule, HardswishModule, LinearHardswishModule, ) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddMM, - Convolution, - Hardswish, - PermuteCopy, - ViewCopy, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_hardtanh_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_hardtanh_converter.py index 66a052dba4f..dc71f5df921 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_hardtanh_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_hardtanh_converter.py @@ -18,16 +18,19 @@ from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import ( BuiltinOperator as Ops, ) +from executorch.backends.nxp.backend.ops_aliases import ( + Convolution, + ExecutorchDelegateCall, + HardTanh, +) 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 graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import Conv2dWithActivation, HardTanhModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - Convolution, - ExecutorchDelegateCall, - HardTanh, +from executorch.backends.nxp.tests.simple_models import ( + Conv2dWithActivation, + HardTanhModule, ) from executorch.backends.nxp.tests.use_qat import * # noqa F403 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 567cf85ebe5..5176a8c61db 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 @@ -8,11 +8,11 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import LeakyRelu from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator 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 LeakyRelu from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_log_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_log_converter.py index 0b7fe88cffc..ebf6a8bfb9b 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_log_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_log_converter.py @@ -8,10 +8,10 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Log 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 Log from executorch.backends.nxp.tests.use_qat import * # noqa F403 from executorch.backends.nxp.tests.dataset_creator import ( LinearRampDatasetCreator, 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 55a47146bfc..cfc4b46022a 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 @@ -9,16 +9,16 @@ import pytest import torch -from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.executors import graph_contains_any_of_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 ( +from executorch.backends.nxp.backend.ops_aliases import ( ExecutorchDelegateCall, GetItem, MaxPool2DWithIndices, ViewCopy, ) +from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program +from executorch.backends.nxp.tests.executors import graph_contains_any_of_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.use_qat import * # noqa F403 @@ -158,10 +158,10 @@ def test__padding_to_kernel_ratio_exceeded(self): to_quantized_edge_program(model, input_shape) -class TestMaxPool1D: +class TestMaxPool1DTo2D: # Just a basic test to verify that the operator gets extended to the 2D variant correctly. - def test__basic_nsys_inference__view_not_delegated(self, mocker, request): + def test__basic_nsys_inference__view_not_delegated(self, mocker, request, use_qat): input_shape = (2, 4, 6) # The old flow limited the batch size to 1. model = MaxPool1DModule() @@ -171,4 +171,4 @@ def test__basic_nsys_inference__view_not_delegated(self, mocker, request): expected_non_delegated_ops={}, ) - lower_run_compare(model, input_shape, graph_verifier, request) + lower_run_compare(model, input_shape, graph_verifier, request, use_qat=use_qat) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py index e25ed98fb3f..01df6211dfc 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py @@ -8,6 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + ExecutorchDelegateCall, + GetItem, + Maximum, + MaxPool2DWithIndices, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import ( @@ -19,13 +25,10 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import MaximumModule, MaxPoolMaximumModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - ExecutorchDelegateCall, - GetItem, - Maximum, - MaxPool2DWithIndices, +from executorch.backends.nxp.tests.simple_models import ( + MaximumModule, + MaxPoolMaximumModule, ) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py index 1674153540f..14f55de94d4 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py @@ -21,6 +21,13 @@ from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.transpose_options import ( Transpose, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, + MeanDim, +) 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 graph_contains_any_of_ops @@ -29,13 +36,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, - MeanDim, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_minimum_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_minimum_converter.py index 9dc2b8d77d7..727ff914ceb 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_minimum_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_minimum_converter.py @@ -8,6 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, + Minimum, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import ( @@ -19,13 +25,10 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import MaxPoolMinimumModule, MinimumModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, - Minimum, +from executorch.backends.nxp.tests.simple_models import ( + MaxPoolMinimumModule, + MinimumModule, ) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_mm_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_mm_converter.py index 423999dc7ec..a8997c11ef5 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_mm_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_mm_converter.py @@ -8,12 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import MM, PermuteCopy, ViewCopy from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier, Operator -from executorch.backends.nxp.tests.models import LinearModule, MmModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import MM, PermuteCopy, ViewCopy +from executorch.backends.nxp.tests.simple_models import LinearModule, MmModule from executorch.backends.nxp.tests.use_qat import * # noqa F403 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 718383284be..414cbecad85 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 @@ -8,6 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, + MulTensor, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import ( @@ -19,13 +25,10 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import MaxPoolMulTensorModule, MulTensorModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, - MulTensor, +from executorch.backends.nxp.tests.simple_models import ( + MaxPoolMulTensorModule, + MulTensorModule, ) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_neg_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_neg_converter.py index 691cb3bd2ca..182e7754013 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_neg_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_neg_converter.py @@ -8,6 +8,7 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Convolution, Neg from executorch.backends.nxp.tests.dataset_creator import ( LinearRampDatasetCreator, @@ -15,7 +16,6 @@ ) 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 Convolution, Neg from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py index 266260f9e1f..d7bbe0f21c6 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py @@ -4,16 +4,18 @@ # 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.converter.builder.model_builder import ( ModelBuilder, ) +from executorch.backends.nxp.backend.ops_aliases import Convolution, Pad from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import PadConvModule, PadModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import Convolution, Pad +from executorch.backends.nxp.tests.simple_models import PadConvModule, PadModule @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_permute_copy_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_permute_copy_converter.py index bdfd1e9da25..2317efd45d9 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_permute_copy_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_permute_copy_converter.py @@ -9,17 +9,17 @@ import pytest import torch from _pytest.mark import ParameterSet - -from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.executors import graph_contains_any_of_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 ( +from executorch.backends.nxp.backend.ops_aliases import ( ExecutorchDelegateCall, GetItem, MaxPool2DWithIndices, PermuteCopy, ) + +from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program +from executorch.backends.nxp.tests.executors import graph_contains_any_of_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.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_prelu_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_prelu_converter.py index 884e95ec20c..58a820f1a7f 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_prelu_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_prelu_converter.py @@ -11,31 +11,31 @@ from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddMM, + Convolution, + ExecutorchDelegateCall, + GtScalar, + MulTensor, + PermuteCopy, + Prelu, + ViewCopy, + WhereSelf, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops 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 ( + +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( ConvPReLUModule, LinearPReLUModule, PReLUModule, TwoPartitionPReLUModel, ) - -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddMM, - Convolution, - ExecutorchDelegateCall, - GtScalar, - MulTensor, - PermuteCopy, - Prelu, - ViewCopy, - WhereSelf, -) from torch.export import ExportedProgram from executorch.backends.nxp.tests.use_qat import * # noqa F403 from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program diff --git a/backends/nxp/tests/ir/converter/node_converter/test_relu_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_relu_converter.py index 1f274576767..d44fc1480db 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_relu_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_relu_converter.py @@ -4,16 +4,12 @@ # 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 exir_ops -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 graph_contains_any_of_ops -from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import Conv2dModule, LinearModule, ReLUModule -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( +from executorch.backends.nxp.backend.ops_aliases import ( AddMM, Convolution, DequantizePerChannel, @@ -23,6 +19,16 @@ Relu, ViewCopy, ) +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 graph_contains_any_of_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.simple_models import ( + Conv2dModule, + LinearModule, + ReLUModule, +) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_rsqrt_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_rsqrt_converter.py index 67101410d9d..916ab5ba518 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_rsqrt_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_rsqrt_converter.py @@ -8,6 +8,7 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Rsqrt from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier @@ -15,7 +16,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import Rsqrt from executorch.backends.nxp.tests.use_qat import * # noqa F403 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 5c4e4f4f007..c6ea90cb5e2 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 @@ -9,6 +9,7 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Sigmoid from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier @@ -16,7 +17,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import Sigmoid from torch import nn from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_slice_copy_tensor_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_slice_copy_tensor_converter.py index 56d0b4bbd64..00b89c859ef 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_slice_copy_tensor_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_slice_copy_tensor_converter.py @@ -8,6 +8,11 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + Convolution, + ExecutorchDelegateCall, + SliceCopy, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program @@ -16,16 +21,11 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( SliceTensorConvModule, SliceTensorModule, ) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - Convolution, - ExecutorchDelegateCall, - SliceCopy, -) @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py index 2ce0790fc98..f893cc1e394 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py @@ -4,8 +4,16 @@ # 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.ops_aliases import ( + Convolution, + ExecutorchDelegateCall, + Softmax, + ViewCopy, +) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops @@ -15,14 +23,8 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import SoftmaxModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - Convolution, - ExecutorchDelegateCall, - Softmax, - ViewCopy, -) +from executorch.backends.nxp.tests.simple_models import SoftmaxModule @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_sub_tensor_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_sub_tensor_converter.py index 1601c1e19c2..54d0b3c6cc2 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_sub_tensor_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_sub_tensor_converter.py @@ -8,6 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, + SubTensor, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import ( @@ -19,13 +25,10 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import MaxPoolSubTensorModule, SubTensorModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, - SubTensor, +from executorch.backends.nxp.tests.simple_models import ( + MaxPoolSubTensorModule, + SubTensorModule, ) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_sum_dim_int_list_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_sum_dim_int_list_converter.py index 8b28142b63a..5726776c78a 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_sum_dim_int_list_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_sum_dim_int_list_converter.py @@ -21,6 +21,13 @@ from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.transpose_options import ( Transpose, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, + SumDimIntList, +) 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 graph_contains_any_of_ops @@ -29,13 +36,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, - SumDimIntList, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_tanh_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_tanh_converter.py index 51b7ee484a7..71b6ced4ed2 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_tanh_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_tanh_converter.py @@ -7,12 +7,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Convolution, Tanh from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import Conv2dWithActivation from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import Convolution, Tanh +from executorch.backends.nxp.tests.simple_models import Conv2dWithActivation from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_upsample_bilinear2d.py b/backends/nxp/tests/ir/converter/node_converter/test_upsample_bilinear2d.py index 949f193b267..8513b87a1a3 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_upsample_bilinear2d.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_upsample_bilinear2d.py @@ -8,6 +8,11 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + ExecutorchDelegateCall, + UpsampleBilinear2D, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program @@ -17,11 +22,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - ExecutorchDelegateCall, - UpsampleBilinear2D, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_upsample_nearest2d.py b/backends/nxp/tests/ir/converter/node_converter/test_upsample_nearest2d.py index b3e28a7b2f8..dd3d7c9c147 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_upsample_nearest2d.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_upsample_nearest2d.py @@ -8,17 +8,17 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + ExecutorchDelegateCall, + UpsampleNearest2D, +) 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 graph_contains_any_of_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 ( - AddTensor, - ExecutorchDelegateCall, - UpsampleNearest2D, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_view_copy_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_view_copy_converter.py index 2a2d270e30a..009e1abc142 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_view_copy_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_view_copy_converter.py @@ -6,18 +6,12 @@ from typing import Sequence import numpy as np + +# noinspection PyUnusedImports import pytest import torch -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 graph_contains_any_of_ops -from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.model_output_comparator import ( - AllCloseOutputComparator, -) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( +from executorch.backends.nxp.backend.ops_aliases import ( AddMM, AddTensor, AvgPool2D, @@ -28,6 +22,14 @@ Relu, ViewCopy, ) +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 graph_contains_any_of_ops +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.model_output_comparator import ( + AllCloseOutputComparator, +) +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare from torch import nn from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/edge_passes/test_convert_reshaping_nodes_to_view.py b/backends/nxp/tests/ir/edge_passes/test_convert_reshaping_nodes_to_view.py index f0489b151f7..82cc3f63526 100644 --- a/backends/nxp/tests/ir/edge_passes/test_convert_reshaping_nodes_to_view.py +++ b/backends/nxp/tests/ir/edge_passes/test_convert_reshaping_nodes_to_view.py @@ -6,14 +6,17 @@ import numpy as np import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import AddTensor, ViewCopy from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import SqueezeAddModel, UnsqueezeAddModel from executorch.backends.nxp.tests.nsys_testing import ( AllCloseOutputComparator, lower_run_compare, ) -from executorch.backends.nxp.tests.ops_aliases import AddTensor, ViewCopy +from executorch.backends.nxp.tests.simple_models import ( + SqueezeAddModel, + UnsqueezeAddModel, +) @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/ir/edge_passes/test_edge_passes.py b/backends/nxp/tests/ir/edge_passes/test_edge_passes.py index 9fa56989cf0..cbd2b2ff48e 100644 --- a/backends/nxp/tests/ir/edge_passes/test_edge_passes.py +++ b/backends/nxp/tests/ir/edge_passes/test_edge_passes.py @@ -21,6 +21,12 @@ PermuteCopyConverter, ViewCopyConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + DequantizePerTensor, + PermuteCopy, + QuantizePerTensor, + ViewCopy, +) from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import ( NeutronEdgePassManager, ) @@ -42,13 +48,12 @@ EdgeProgramExecutor, OverrideTargetSupportCheck, ) -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.simple_models import ( Conv2dModule, ConvActivationModule, ConvFCFCSoftmaxModuleWithoutReshape, LinearActivationModule, ) -from executorch.exir.dialects._ops import ops as exir_ops from executorch.extension.export_util.utils import export_to_edge from parameterized import parameterized from torch.export import ExportedProgram @@ -56,10 +61,7 @@ def _is_view_copy(node_: Node) -> bool: - return ( - node_.op == "call_function" - and node_.target == exir_ops.edge.aten.view_copy.default - ) + return node_.op == "call_function" and node_.target == ViewCopy def _find_view_copy_node_indices(graph_nodes: list[Node]) -> list[int]: @@ -352,16 +354,10 @@ def test_remove_additional_quantize_dequantize_nodes_pass(self): ) nodes = list(edge_program_with_qdq_cluster.graph.nodes) assert len(nodes) == 10 - assert ( - nodes[5].target - == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default - ) - assert nodes[6].target == exir_ops.edge.aten.permute_copy.default + assert nodes[5].target == DequantizePerTensor + assert nodes[6].target == PermuteCopy assert "cluster" in nodes[6].meta - assert ( - nodes[7].target - == exir_ops.edge.quantized_decomposed.quantize_per_tensor.default - ) + assert nodes[7].target == QuantizePerTensor # Run pass for removal of additional QDQ nodes and compute in non-float types where possible edge_program_manager = edge_program_manager.transform( @@ -373,12 +369,9 @@ def test_remove_additional_quantize_dequantize_nodes_pass(self): nodes = list(edge_program_without_qdq_cluster.graph.nodes) assert len(nodes) == 8 assert nodes[4].name == "getitem" - assert nodes[5].target == exir_ops.edge.aten.permute_copy.default + assert nodes[5].target == PermuteCopy assert "cluster" not in nodes[5].meta - assert ( - nodes[6].target - == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default - ) + assert nodes[6].target == DequantizePerTensor edge_program_executor_without_qdq_cluster = EdgeProgramExecutor( edge_program_without_qdq_cluster diff --git a/backends/nxp/tests/ir/edge_passes/test_linear_bn_fusing.py b/backends/nxp/tests/ir/edge_passes/test_linear_bn_fusing.py index aadef8c7731..8a41ee7221a 100644 --- a/backends/nxp/tests/ir/edge_passes/test_linear_bn_fusing.py +++ b/backends/nxp/tests/ir/edge_passes/test_linear_bn_fusing.py @@ -3,7 +3,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import executorch.backends.nxp.tests.models as models +import executorch.backends.nxp.tests.simple_models as models import numpy as np import pytest import torch @@ -21,6 +21,7 @@ batch_norm_target_ops, is_batch_norm, ) +from executorch.backends.nxp.backend.ops_aliases import AddMM, Linear from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer from executorch.backends.nxp.tests.executorch_pipeline import ( get_random_calibration_inputs, @@ -34,7 +35,6 @@ ToChannelFirstPreprocess, ToChannelLastPreprocess, ) -from executorch.exir.dialects._ops import ops as exir_ops from torch.export import export, ExportedProgram from torchao.quantization.pt2e.prepare import _is_activation_post_process_node from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_qat_pt2e @@ -243,8 +243,8 @@ def test_linear_bn_full_qat_pipeline_conversion( assert not graph_contains_any_of_ops( graph=edge_program.graph, ops=[ - exir_ops.edge.aten.addmm.default, - exir_ops.edge.aten.linear.default, + AddMM, + Linear, ] + batch_norm_target_ops, ) @@ -303,8 +303,8 @@ def test_incompatible_linear_bn_not_fused(mocker, input_shape, linear_bias, bn_e assert graph_contains_any_of_ops( graph=edge_program.graph, ops=[ - exir_ops.edge.aten.addmm.default, - exir_ops.edge.aten.linear.default, + AddMM, + Linear, ], ) assert graph_contains_any_of_ops( diff --git a/backends/nxp/tests/ir/edge_passes/test_remove_io_quant_ops_pass.py b/backends/nxp/tests/ir/edge_passes/test_remove_io_quant_ops_pass.py index ef669897b51..b24b3e4a555 100644 --- a/backends/nxp/tests/ir/edge_passes/test_remove_io_quant_ops_pass.py +++ b/backends/nxp/tests/ir/edge_passes/test_remove_io_quant_ops_pass.py @@ -10,7 +10,7 @@ import torch from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.models import Conv2dReLUModule +from executorch.backends.nxp.tests.simple_models import Conv2dReLUModule from executorch.examples.nxp.experimental.cifar_net.cifar_net import CifarNet from executorch.exir import ExecutorchBackendConfig from executorch.exir.passes.quantize_io_pass import get_config_method_name diff --git a/backends/nxp/tests/models/__init__.py b/backends/nxp/tests/models/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backends/nxp/tests/generic_tests/test_cifarnet.py b/backends/nxp/tests/models/test_cifarnet.py similarity index 98% rename from backends/nxp/tests/generic_tests/test_cifarnet.py rename to backends/nxp/tests/models/test_cifarnet.py index 6db8ebc9a03..c1d1cd91502 100644 --- a/backends/nxp/tests/generic_tests/test_cifarnet.py +++ b/backends/nxp/tests/models/test_cifarnet.py @@ -5,6 +5,7 @@ import os.path +# noinspection PyUnusedImports import pytest import torch diff --git a/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py b/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py index 782f7c13a9a..6dabb505268 100644 --- a/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py +++ b/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py @@ -6,7 +6,11 @@ from functools import partial import numpy as np + +# noinspection PyUnusedImports +import pytest import torch + from executorch.backends.nxp.tests.dataset_creator import ( FromCalibrationDataDatasetCreator, ) @@ -23,7 +27,6 @@ ReferenceModel, ) from executorch.backends.nxp.tests.use_qat import * # noqa F403 -import pytest from executorch.examples.nxp.models.mlperf_tiny.image_classification.mlperf_tiny_image_classification import ( MLPerfTinyImageClassification, ) @@ -64,9 +67,9 @@ def test_mlperf_tiny_classification_mse_cpu_vs_npu( input_spec.dim_order = torch.channels_last quant_type_key = "QAT" if use_qat else "PTQ" - dim_order_key = "channels-last" if channels_last else "channels-first" + format_key = "channels-last" if channels_last else "channels-first" - mse = BOUNDS_MSE[quant_type_key][dim_order_key] + mse = BOUNDS_MSE[quant_type_key][format_key] comparator = NumericalStatsOutputComparator( max_mse_error=mse, use_softmax=True, is_classification_task=True ) @@ -77,15 +80,6 @@ def test_mlperf_tiny_classification_mse_cpu_vs_npu( else None ) - # This model does not work in channels-last format and QAT. See more information below. - # Github issue: https://github.com/pytorch/executorch/issues/22179 - # NXP internal issue ID: EIEX-1065 - ref_model = ( - ReferenceModel.QUANTIZED_EDGE_PYTHON - if channels_last and use_qat - else ReferenceModel.QUANTIZED_EXECUTORCH_CPP - ) - lower_run_compare( model, [input_spec], @@ -93,7 +87,7 @@ def test_mlperf_tiny_classification_mse_cpu_vs_npu( request, dataset_creator=dataset_creator, output_comparator=comparator, - reference_model=ref_model, + reference_model=ReferenceModel.QUANTIZED_EXECUTORCH_CPP, mocker=mocker, use_qat=use_qat, train_fn=train_fn, diff --git a/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py b/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py new file mode 100644 index 00000000000..ab56849e970 --- /dev/null +++ b/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py @@ -0,0 +1,140 @@ +# Copyright 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. + +from functools import partial + +import numpy as np +import pytest +import torch +from executorch.backends.nxp.tests.dataset_creator import ( + FromCalibrationDataDatasetCreator, +) +from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec +from executorch.backends.nxp.tests.graph_verifier import BaseGraphVerifier +from executorch.backends.nxp.tests.model_output_comparator import ( + ClassificationAccuracyOutputComparator, + NumericalStatsOutputComparator, +) +from executorch.backends.nxp.tests.nsys_testing import ( + lower_run_compare, + lower_run_compare_ptq_qat, + ReferenceModel, +) +from executorch.backends.nxp.tests.use_qat import * # noqa F403 +from executorch.examples.nxp.models.mlperf_tiny.keyword_spotting.mlperf_tiny_keyword_spotting import ( + MLPerfTinyKeywordSpotting, +) + +BOUNDS_MSE = { + "PTQ": { + "channels-last": np.inf, + "channels-first": 5.5e-7, + }, + "QAT": { + "channels-last": np.inf, + "channels-first": 3.3e-5, + }, +} + + +@pytest.fixture(autouse=True) +def reseed_model_per_test_run(): + torch.manual_seed(23) + np.random.seed(23) + + +@pytest.mark.parametrize( + "channels_last", + [ + False, + pytest.param( + True, + marks=pytest.mark.xfail( + reason="EIEX-1082, don't forget to readjust bounds when it start working", + strict=True, + ), + ), + ], +) +def test_mlperf_tiny_kws_mse_cpu_vs_npu(mocker, request, channels_last, use_qat): + # approx. 5 samples per class + num_samples = 60 + + kws = MLPerfTinyKeywordSpotting(num_samples=num_samples, use_random_dataset=True) + model = kws.get_eager_model() + dataset = kws.dataset + labels = kws.labels + + dataset_creator = FromCalibrationDataDatasetCreator( + dataset, num_examples=num_samples, idx_to_label=labels + ) + + input_spec = ModelInputSpec(kws.input_shape) + if channels_last: + model.to(memory_format=torch.channels_last) + input_spec.dim_order = torch.channels_last + + quant_type_key = "QAT" if use_qat else "PTQ" + format_key = "channels-last" if channels_last else "channels-first" + mse = BOUNDS_MSE[quant_type_key][format_key] + comparator = NumericalStatsOutputComparator( + max_mse_error=mse, is_classification_task=True + ) + model_verifier = BaseGraphVerifier(1, []) + train_fn = ( + partial(kws.train_model_fn, channels_last=channels_last) if use_qat else None + ) + + # This model does not work in channels-last format when running with portable kernels. + # See more information below. + # Github issue: https://github.com/pytorch/executorch/issues/22520 + # NXP internal issue ID: EIEX-1074 + ref_model = ( + ReferenceModel.QUANTIZED_EDGE_PYTHON + if channels_last + else ReferenceModel.QUANTIZED_EXECUTORCH_CPP + ) + + lower_run_compare( + model, + [input_spec], + model_verifier, + request, + dataset_creator=dataset_creator, + output_comparator=comparator, + mocker=mocker, + reference_model=ref_model, + use_qat=use_qat, + train_fn=train_fn, + ) + + +def test_mlperf_tiny_kws_ptq_qat_equivalence(request): + # approx. 5 samples per class + num_samples = 60 + + kws = MLPerfTinyKeywordSpotting(num_samples=num_samples, use_random_dataset=True) + + model = kws.get_eager_model() + dataset = kws.dataset + labels = kws.labels + + dataset_creator = FromCalibrationDataDatasetCreator( + dataset, num_examples=num_samples, idx_to_label=labels + ) + comparator = ClassificationAccuracyOutputComparator(class_dict=labels) + + input_spec = ModelInputSpec(kws.input_shape) + model_verifier = BaseGraphVerifier(1, []) + + lower_run_compare_ptq_qat( + model, + [input_spec], + model_verifier, + request, + train_fn=kws.train_model_fn, + dataset_creator=dataset_creator, + output_comparator=comparator, + ) diff --git a/backends/nxp/tests/nsys_testing.py b/backends/nxp/tests/nsys_testing.py index a8038083c37..c7364967b6b 100644 --- a/backends/nxp/tests/nsys_testing.py +++ b/backends/nxp/tests/nsys_testing.py @@ -16,14 +16,9 @@ from os import environ, mkdir from typing import Callable, Iterable -import numpy as np import torch import yaml -from executorch.backends.nxp.backend.edge_helper import is_channels_last_dim_order -from executorch.backends.nxp.backend.ir.converter.conversion import translator -from executorch.backends.nxp.backend.ir.converter.conversion.translator import ( - torch_type_to_numpy_type, -) +from executorch.backends.nxp.backend.ops_aliases import ExecutorchDelegateCall from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner from executorch.backends.nxp.tests.config_importer import test_config from executorch.backends.nxp.tests.dataset_creator import ( @@ -33,6 +28,7 @@ ) from executorch.backends.nxp.tests.executorch_pipeline import ( get_calibration_inputs_fn_from_dataset_dir, + get_example_input, ModelInputSpec, to_edge_program, to_model_input_spec, @@ -44,10 +40,15 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.ops_aliases import ExecutorchDelegateCall from executorch.backends.nxp.tests.outputs_dir_importer import outputs_dir -from executorch.backends.nxp.tests.utils import save_pte_program, store_txt_input_tensor - +from executorch.backends.nxp.tests.utils import ( + process_input_sample, + process_output_sample, + read_prepared_samples, + save_pte_program, + store_results, + store_txt_input_tensor, +) from executorch.devtools.visualization.visualization_utils import ( visualize_with_clusters, ) @@ -72,6 +73,7 @@ class ReferenceModel(Enum): # QUANTIZED_ATEN_PYTHON = 2 # Not implemented. # FLOAT_ATEN_PYTHON = 3 # Not implemented. FLOAT_PYTORCH_PYTHON = 4 + QUANTIZED_CORTEX_M = 5 def _get_dataset_cli_args(input_spec: list[ModelInputSpec], testing_dataset_dir): @@ -251,102 +253,6 @@ def _save_non_quantized_fp32_executorch_program( return non_quantized_program.exported_program() -def read_prepared_samples( - dataset_dir: str, input_spec: list[ModelInputSpec] -) -> list[tuple[np.ndarray, ...]]: - """Read numpy arrays generated by a `DatasetCreator`. - - :param dataset_dir: Directory containing the generated samples - :param input_spec: List of ModelInputSpec defining the shape and type of each input - - :return: List of tuples, where each tuple contains numpy arrays for one sample - """ - all_samples = [] - - # Multi-input: samples are in numbered subdirectories - if len(input_spec) > 1: - sample_dirs = sorted( - [ - d - for d in os.listdir(dataset_dir) - if os.path.isdir(os.path.join(dataset_dir, d)) - ] - ) - - for sample_name in sample_dirs: - sample_dir = os.path.join(dataset_dir, sample_name) - current_samples = [] - - for spec_idx, spec in enumerate(input_spec): - bin_file_path = os.path.join( - sample_dir, f"{str(spec_idx).zfill(2)}.bin" - ) - sample_vector = np.fromfile( - bin_file_path, dtype=torch_type_to_numpy_type(spec.dtype) - ).reshape(spec.shape) - current_samples.append(sample_vector) - - all_samples.append(tuple(current_samples)) - - # Single-input: binary files are directly in dataset_dir - else: - bin_files = sorted([f for f in os.listdir(dataset_dir) if f.endswith(".bin")]) - - for bin_file in bin_files: - bin_file_path = os.path.join(dataset_dir, bin_file) - sample_vector = np.fromfile( - bin_file_path, dtype=torch_type_to_numpy_type(input_spec[0].dtype) - ).reshape(input_spec[0].shape) - all_samples.append((sample_vector,)) - - return all_samples - - -def store_results( - results: list[tuple[np.ndarray, ...]], output_dir: str, reference_dir: str -): - """Store a list of output arrays in the directory structure matching the reference directory. - - :param results: List of tuples, where each tuple contains numpy arrays (outputs for one sample) - :param output_dir: Directory where results will be stored - - Directory structure created matches reference_dir: - output_dir/ - ├── sample_0/ - │ ├── 0000.bin - │ └── 0001.bin - ├── some_other_sample/ - │ ├── 0000.bin - │ └── 0001.bin - """ - os.makedirs(output_dir, exist_ok=True) - - # Get subdirectories from reference directory - sample_dirs = sorted( - [ - d - for d in os.listdir(reference_dir) - if os.path.isdir(os.path.join(reference_dir, d)) - ] - ) - - assert len(sample_dirs) == len( - results - ), f"Number of samples ({len(results)}) must match number of subdirectories in reference_dir ({len(sample_dirs)})" - - for _sample_idx, (sample_name, sample_outputs) in enumerate( - zip(sample_dirs, results) - ): - sample_dir = os.path.join(output_dir, sample_name) - os.makedirs(sample_dir, exist_ok=True) - - # Store each output tensor - for output_idx, output_array in enumerate(sample_outputs): - bin_file_name = f"{str(output_idx).zfill(4)}.bin" - bin_file_path = os.path.join(sample_dir, bin_file_name) - output_array.tofile(bin_file_path) - - def _run_python_program( model: torch.nn.Module | GraphModule, testing_dataset_dir, @@ -371,53 +277,102 @@ def _run_python_program( all_outputs = [] for input_samples in read_prepared_samples(testing_dataset_dir, input_spec): - current_input_samples = [] - for spec, sample in zip(input_spec, input_samples, strict=True): - match spec.dim_order: - case torch.contiguous_format: - # Use the data as is, just turn it into a PyTorch tensor. - sample = torch.tensor(sample) - - case torch.channels_last: - # The tensor data was stored by the DatasetCreator as channels last (NHWC), but it was now - # incorrectly parsed as contiguous/channels first (NCHW). Transpose it to channels last to preserve - # the semantics. - channels_last_shape = translator.dims_to_channels_last( - list(spec.shape) - ) - sample = np.moveaxis(sample.reshape(channels_last_shape), -1, 1) - sample = torch.tensor(sample).to(memory_format=torch.channels_last) - - case _: - raise ValueError(f"Unsupported dim_order: {spec.dim_order}") - - current_input_samples.append(sample) + current_input_samples = process_input_sample(input_spec, input_samples) # Run the model. output = model(*current_input_samples) - if isinstance(output, torch.Tensor): - output = (output,) + current_outputs = process_output_sample(output, output_spec) + all_outputs.append(current_outputs) - current_outputs = [] + # Store all the results. + store_results(all_outputs, cpu_results_dir, npu_results_dir) - for o, o_spec in zip(output, output_spec, strict=True): - dim_order = list(o_spec.dim_order()) # ExecuTorch dim order. - rank = len(o_spec.shape) - if dim_order == list(range(rank)): # Contiguous dim order. - current_outputs.append(o.detach().numpy()) - elif is_channels_last_dim_order(dim_order): # Channels last dim order. - # The NPU variant outputs channels last (NHWC). We need to convert the CPU output to match. - o = o.detach().numpy().reshape(o_spec.shape) - current_outputs.append(np.moveaxis(o, 1, -1)) +def _run_cortex_m_program( + model: torch.nn.Module | GraphModule, + test_dir, + test_name, + calibration_dataset_dir, + testing_dataset_dir, + input_spec: list[ModelInputSpec], + output_spec: list[torch.Tensor], + cpu_results_dir, + npu_results_dir, +): + """Run a model with Cortex-M backend with channels last inputs. + + :param model: Any PyTorch/ExecuTorch model runnable with Cortex-M with channels last inputs. + :param test_dir: Directory for saving test artifacts. + :param test_name: Name of the test. + :param calibration_dataset_dir: Directory containing calibration data. + :param testing_dataset_dir: Directory containing testing data. The samples have to be channels last (NHWC) for 4D tensors. + The format must match the input_spec.dim_order. + :param input_spec: List of ModelInputSpec defining the shape, type, and dimension order of each input. + :param output_spec: List of output tensor specifications. + :param cpu_results_dir: Directory where CPU results will be stored. The structure will match the existing structure + of `npu_results_dir`. + :param npu_results_dir: Directory where NPU results are already stored, to serve as reference directory structure + for `cpu_results_dir`. + """ + # Assert Cortex-M dependencies are available + assert_cortex_m() - else: - raise ValueError(f"Unsupported dim_order: {o_spec.dim_order}") + from executorch.backends.nxp.tests.cortex_m_benchmarking import ( + CortexMNXPBenchmarkTester, + ) - all_outputs.append(current_outputs) + numpy_samples = read_prepared_samples(calibration_dataset_dir, input_spec) + calibration_samples = [ + process_input_sample(input_spec, sample) for sample in numpy_samples + ] - # Store all the results. - store_results(all_outputs, cpu_results_dir, npu_results_dir) + example_inputs = get_example_input(input_spec) + + tester = CortexMNXPBenchmarkTester( + model, + example_inputs, + ) + cortex_m_delegated_program = tester.run_benchmark( + calibration_samples, + input_spec, + output_spec, + testing_dataset_dir, + cpu_results_dir, + npu_results_dir, + ) + save_pte_program( + cortex_m_delegated_program, test_name + "_cortex_m_delegated", test_dir + ) + + +def assert_cortex_m(): + # Follow backends/cortex_m/README.md to install the required dependencies. + # Build Arm executor runner with target="cortex-m33": + # ./backends/cortex_m/test/build_test_runner.sh --target="cortex-m33" + # FVP Corstone-300 simulator needs to be added to PATH. + + import sysconfig + + suffix = sysconfig.get_config_var( + "EXT_SUFFIX" + ) # e.g. ".cpython-312-x86_64-linux-gnu.so" + cmsis_nn_lib_path = os.path.join( + PROJECT_DIR, "backends/cortex_m/library/_cmsis_nn", f"cmsis_nn{suffix}" + ) + + assert os.path.exists( + cmsis_nn_lib_path + ), "CMSIS-NN lib is not available, check if ET is built correctly." + fvp_simulator_path = os.path.join( + PROJECT_DIR, + "examples/arm/arm-scratch/FVP-corstone300/models/Linux64_GCC-9.3/FVP_Corstone_SSE-300_Ethos-U55", + ) + assert os.path.exists(fvp_simulator_path), "Arm FVP Corstone-300 is not installed." + arm_executor_runner_path = os.path.join( + PROJECT_DIR, + "arm_test/arm_semihosting_executor_runner_corstone-300_cortex-m33/arm_executor_runner", + ) + assert os.path.exists(arm_executor_runner_path), "Arm executor is not installed." def assert_NSYS(): @@ -590,11 +545,44 @@ def lower_run_compare( npu_results_dir, ) + case ReferenceModel.QUANTIZED_CORTEX_M: + if use_qat: + raise ValueError( + "Flag use_qat is not applicable to QUANTIZED_CORTEX_M reference model " + "as it doesn't support QAT. Run with use_qat=False." + ) + if remove_quant_io_ops: + raise ValueError( + "Flag remove_quant_io_ops is not applicable to QUANTIZED_CORTEX_M reference model " + "as it works with float data only. Run with remove_quant_io_ops=False." + ) + if any( + spec.dim_order != torch.channels_last + for spec in input_spec + if len(spec.shape) == 4 + ): + raise ValueError( + "Cortex-M backend supports only channel last dim order for 4D inputs." + ) + + model_to_delegate_cortex_m = deepcopy(model) + + # Lower to quantized Cortex-M program and run on Arm simulator. + _run_cortex_m_program( + model_to_delegate_cortex_m, + test_dir, + test_name, + calibration_dataset_dir, + testing_dataset_dir, + input_spec, + output_spec, + cpu_results_dir, + npu_results_dir, + ) + case _: raise ValueError(f"Unsupported reference model: `{reference_model}`.") - output_tensor_spec = _get_program_output_spec(delegated_program) - if logging.root.isEnabledFor(logging.DEBUG): _generate_txt_test_data( calibration_dataset_dir, testing_dataset_dir, list(input_spec) @@ -602,9 +590,7 @@ def lower_run_compare( dump_debug_test_summary(test_name, test_dir) npu_results_dir = os.path.join(test_dir, "results_npu") cpu_results_dir = os.path.join(test_dir, "results_cpu") - output_comparator.compare_results( - cpu_results_dir, npu_results_dir, output_tensor_spec - ) + output_comparator.compare_results(cpu_results_dir, npu_results_dir, output_spec) def lower_run_compare_ptq_qat( diff --git a/backends/nxp/tests/profiling_utils.py b/backends/nxp/tests/profiling_utils.py new file mode 100644 index 00000000000..f43cb5455cd --- /dev/null +++ b/backends/nxp/tests/profiling_utils.py @@ -0,0 +1,139 @@ +# Copyright 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 logging +import re +import subprocess + +from executorch.devtools.etdump.serialize import deserialize_from_etdump_flatcc + + +def get_neutron_driver_version(etdump_path: str) -> str: + """ + Extract the Neutron Driver version from an ETDump file. + + The Neutron Driver version is stored in the metadata of the last Neutron + delegate event. This event is emitted when the profiling dump is generated. + The version is encoded as a 16-bit value in little-endian format: + - 4 bits - major version + - 4 bits - minor version + - 4 bits - patch version + - 4 bits - reserved + + :param etdump_path: Path to the ETDump binary file. + :return: Neutron Driver version string (e.g. "1.2.3") if successfully decoded, + otherwise empty string. Errors are logged instead of raised. + """ + + try: + with open(etdump_path, "rb") as f: + data = f.read() + etdump = deserialize_from_etdump_flatcc(data) + except Exception as e: + logging.exception("Failed to load ETDump: %s", e) + return "" + + events = [] + try: + for run in etdump.run_data: + for event in run.events: + profile_event = getattr(event, "profile_event", None) + if ( + profile_event is not None + and getattr(profile_event, "delegate_debug_id_int", 0) > 0 + ): + events.append(event) + except Exception as e: + logging.exception("Failed while processing events: %s", e) + return "" + + try: + metadata = events[-1].profile_event.delegate_debug_metadata + if not metadata or len(metadata) < 2: + logging.error("Invalid delegate_debug_metadata") + return "" + + major, minor, patch = [ + (int.from_bytes(metadata, "little") >> shift) & 0xF for shift in (8, 4, 0) + ] + return f"{major}.{minor}.{patch}" + + except Exception as e: + logging.exception("Failed to extract version from metadata: %s", e) + return "" + + +def get_neutron_compiler_version() -> str: + """ + Get the Neutron Compiler version reported by the neutron_compiler tool. + + Executes `neutron_compiler --version` and returns the version as + {major}.{minor}.{patch} string. + + :return: The version string returned by neutron_compiler, or empty string if + the command fails, times out, or the executable is not available. + Errors are logged instead of being raised. + """ + + try: + # Use neutron_compiler because neutron_converter executable is unavailable since NS 3.2.1. + proc = subprocess.Popen( + ["neutron_compiler", "--version"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + stdout, stderr = proc.communicate(timeout=10) + if proc.returncode != 0: + logging.error( + "Failed to get compiler version: %s", + stderr.strip(), + ) + return "" + version_match = re.search(r"version\s(\d+\.\d+\.\d+)", stdout) + if version_match: + return version_match.group(1) + else: + logging.exception("Unexpected error while getting neutron compiler version") + return "" + except Exception: + logging.exception("Error while getting neutron compiler version") + return "" + + +def get_neutron_kernel_kinds(target: str = "imxrt700") -> dict[int, str]: + """ + Retrieve kernel kinds supported by neutron_compiler for the specified target. + + Executes the neutron_compiler command with the --show-kernel-kinds option, + parses its output, and returns a dictionary mapping kernel IDs to kernel + names. + + :param target: Target platform for which kernel kinds should be queried. + Defaults to "imxrt700". + :return: Returns empty dict if neutron_compiler exits with an error. + Otherwise, a dictionary where: + - key: kernel ID (int) + - value: kernel name (str) + """ + + # Use neutron_compiler because neutron_converter executable is unavailable since NS 3.2.1. + proc = subprocess.Popen( + ["neutron_compiler", "--target", target, "--show-kernel-kinds"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + stdout, stderr = proc.communicate(timeout=10) + if proc.returncode != 0: + logging.error( + "Failed to get kernrl kinds from neutron_compiler: %s", + stderr.strip(), + ) + return {} + return { + int(op_id): name + for op_id, name in re.findall(r"\[\s*(\d+)\s*\]\s+(.+)", stdout) + } diff --git a/backends/nxp/tests/models.py b/backends/nxp/tests/simple_models.py similarity index 99% rename from backends/nxp/tests/models.py rename to backends/nxp/tests/simple_models.py index eaa5e57cee6..85158b77625 100644 --- a/backends/nxp/tests/models.py +++ b/backends/nxp/tests/simple_models.py @@ -456,6 +456,16 @@ def forward(self, x): return self.avg_pool(x) +class AdaptiveAvgPool1dModule(torch.nn.Module): + def __init__(self, output_size): + super().__init__() + + self.adaptive_avg_pool = torch.nn.AdaptiveAvgPool1d(output_size=output_size) + + def forward(self, x): + return self.adaptive_avg_pool(x) + + class AdaptiveAvgPool2dModule(torch.nn.Module): def __init__(self, output_size): super().__init__() diff --git a/backends/nxp/tests/use_qat.py b/backends/nxp/tests/use_qat.py index 7a63270ae21..c7996f8aad1 100644 --- a/backends/nxp/tests/use_qat.py +++ b/backends/nxp/tests/use_qat.py @@ -3,6 +3,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +# noinspection PyUnusedImports import pytest diff --git a/backends/nxp/tests/utils.py b/backends/nxp/tests/utils.py index 00b7c364a31..ef0fe956b4c 100644 --- a/backends/nxp/tests/utils.py +++ b/backends/nxp/tests/utils.py @@ -11,7 +11,11 @@ import numpy as np +import torch +from executorch.backends.nxp.backend.edge_helper import is_channels_last_dim_order + from executorch.backends.nxp.backend.ir.converter.conversion.translator import ( + dims_to_channels_last, torch_type_to_numpy_type, ) from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec @@ -64,3 +68,165 @@ def store_txt_input_tensor( def archive_test_dir(test_dir: str): shutil.make_archive(test_dir, "zip", test_dir) + + +def read_prepared_samples( + dataset_dir: str, input_spec: list[ModelInputSpec] +) -> list[tuple[np.ndarray, ...]]: + """Read numpy arrays generated by a `DatasetCreator`. + + :param dataset_dir: Directory containing the generated samples + :param input_spec: List of ModelInputSpec defining the shape and type of each input + + :return: List of tuples, where each tuple contains numpy arrays for one sample + """ + all_samples = [] + + # Multi-input: samples are in numbered subdirectories + if len(input_spec) > 1: + sample_dirs = sorted( + [ + d + for d in os.listdir(dataset_dir) + if os.path.isdir(os.path.join(dataset_dir, d)) + ] + ) + + for sample_name in sample_dirs: + sample_dir = os.path.join(dataset_dir, sample_name) + current_samples = [] + + for spec_idx, spec in enumerate(input_spec): + bin_file_path = os.path.join( + sample_dir, f"{str(spec_idx).zfill(2)}.bin" + ) + sample_vector = np.fromfile( + bin_file_path, dtype=torch_type_to_numpy_type(spec.dtype) + ).reshape(spec.shape) + current_samples.append(sample_vector) + + all_samples.append(tuple(current_samples)) + + # Single-input: binary files are directly in dataset_dir + else: + bin_files = sorted([f for f in os.listdir(dataset_dir) if f.endswith(".bin")]) + + for bin_file in bin_files: + bin_file_path = os.path.join(dataset_dir, bin_file) + sample_vector = np.fromfile( + bin_file_path, dtype=torch_type_to_numpy_type(input_spec[0].dtype) + ).reshape(input_spec[0].shape) + all_samples.append((sample_vector,)) + + return all_samples + + +def store_results( + results: list[tuple[np.ndarray, ...]], output_dir: str, reference_dir: str +): + """Store a list of output arrays in the directory structure matching the reference directory. + + :param results: List of tuples, where each tuple contains numpy arrays (outputs for one sample) + :param output_dir: Directory where results will be stored + + Directory structure created matches reference_dir: + output_dir/ + ├── sample_0/ + │ ├── 0000.bin + │ └── 0001.bin + ├── some_other_sample/ + │ ├── 0000.bin + │ └── 0001.bin + """ + os.makedirs(output_dir, exist_ok=True) + + # Get subdirectories from reference directory + sample_dirs = sorted( + [ + d + for d in os.listdir(reference_dir) + if os.path.isdir(os.path.join(reference_dir, d)) + ] + ) + + assert len(sample_dirs) == len( + results + ), f"Number of samples ({len(results)}) must match number of subdirectories in reference_dir ({len(sample_dirs)})" + + for _sample_idx, (sample_name, sample_outputs) in enumerate( + zip(sample_dirs, results) + ): + sample_dir = os.path.join(output_dir, sample_name) + os.makedirs(sample_dir, exist_ok=True) + + # Store each output tensor + for output_idx, output_array in enumerate(sample_outputs): + bin_file_name = f"{str(output_idx).zfill(4)}.bin" + bin_file_path = os.path.join(sample_dir, bin_file_name) + output_array.tofile(bin_file_path) + + +def process_input_sample( + input_spec: list[ModelInputSpec], input_samples: tuple[np.ndarray, ...] +) -> list[torch.Tensor]: + """Process input samples by converting them to PyTorch tensors with correct dimension order. + + :param input_spec: List of ModelInputSpec defining the shape, type, and dimension order of each input + :param input_samples: Tuple of numpy arrays representing one sample + + :return: List of PyTorch tensors with correct dimension order + """ + current_input_samples = [] + for spec, sample in zip(input_spec, input_samples, strict=True): + match spec.dim_order: + case torch.contiguous_format: + # Use the data as is, just turn it into a PyTorch tensor. + sample = torch.tensor(sample) + + case torch.channels_last: + # The tensor data was stored by the DatasetCreator as channels last (NHWC), but it was now + # incorrectly parsed as contiguous/channels first (NCHW). Transpose it to channels last to preserve + # the semantics. + channels_last_shape = dims_to_channels_last(list(spec.shape)) + sample = np.moveaxis(sample.reshape(channels_last_shape), -1, 1) + sample = torch.tensor(sample).to(memory_format=torch.channels_last) + + case _: + raise ValueError(f"Unsupported dim_order: {spec.dim_order}") + + current_input_samples.append(sample) + + return current_input_samples + + +def process_output_sample( + output: tuple[torch.Tensor, ...] | torch.Tensor, output_spec: list[torch.Tensor] +) -> list[np.ndarray]: + """Process output tensors by converting them to numpy arrays with correct dimension order. + + :param output: Model output - either a single tensor or tuple of tensors + :param output_spec: List of output tensor specifications + + :return: List of numpy arrays with correct dimension order matching NPU output format + """ + + if isinstance(output, torch.Tensor): + output = (output,) + + current_outputs = [] + + for o, o_spec in zip(output, output_spec, strict=True): + dim_order = list(o_spec.dim_order()) # ExecuTorch dim order. + rank = len(o_spec.shape) + if dim_order == list(range(rank)): # Contiguous dim order. + current_outputs.append(o.detach().numpy()) + + elif is_channels_last_dim_order(dim_order): # Channels last dim order. + # The NPU variant outputs channels last (NHWC). We need to convert the CPU output to match. + o = o.detach().numpy().reshape(o_spec.shape) + current_outputs.append(np.moveaxis(o, 1, -1)) + + else: + raise ValueError(f"Unsupported dim_order: {o_spec.dim_order()}") + + return current_outputs diff --git a/backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py b/backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py index a94766ab335..7a42b0e4a63 100644 --- a/backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py +++ b/backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py @@ -51,17 +51,18 @@ def get_passes_dependency_for_capture_program(cls): { DecomposeHardsigmoid: [RemoveRedundancy], DecomposeReciprocal: [RemoveRedundancy], - LpaiPartitionFallbackSupport: [TagQuantIO], - ResolveDebugHandle: [LpaiPartitionFallbackSupport], + LpaiPartitionFallbackSupport: [TagQuantIO, ResolveDebugHandle], } ) return deps def _validate_edge_passes(self) -> None: - super()._validate_edge_passes() assert isinstance( - self.passes[-2], LpaiPartitionFallbackSupport - ), "Please ensure LpaiPartitionFallbackSupport is the last edge pass before ResolveDebugHandle." + self.passes[-2], ResolveDebugHandle + ), "Please ensure ResolveDebugHandle is the last edge pass before LpaiPartitionFallbackSupport." + assert isinstance( + self.passes[-1], LpaiPartitionFallbackSupport + ), "Please ensure LpaiPartitionFallbackSupport is the last pass." @classmethod def get_annotation_passes(cls): diff --git a/backends/qualcomm/_passes/lpai_partition_fallback_support.py b/backends/qualcomm/_passes/lpai_partition_fallback_support.py index 5983c145749..ad4a611f24b 100644 --- a/backends/qualcomm/_passes/lpai_partition_fallback_support.py +++ b/backends/qualcomm/_passes/lpai_partition_fallback_support.py @@ -254,7 +254,9 @@ def insert_partition_qdq( output_dq_node.meta[QCOM_BYPASS_NODE] = True graph_module.graph.eliminate_dead_code() - def handle_back_to_back_nodes(self, graph_module: torch.fx.GraphModule): + def handle_back_to_back_nodes( + self, graph_module: torch.fx.GraphModule, unsupported_nodes: set[torch.fx.Node] + ): """ This function takes care of following cases: 1. When 2 contiguous fall back nodes ``a`` and ``b`` (both @@ -279,6 +281,7 @@ def handle_back_to_back_nodes(self, graph_module: torch.fx.GraphModule): input_node for input_node in node.all_input_nodes if input_node.op == "call_function" + and input_node not in unsupported_nodes ] assert all( input_node.target in dq_ops for input_node in input_call_func_nodes @@ -327,7 +330,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: unsupported_nodes = self.get_unsupported_nodes(graph_module) for node in unsupported_nodes: self.insert_partition_qdq(graph_module, node) - self.handle_back_to_back_nodes(graph_module) + self.handle_back_to_back_nodes(graph_module, unsupported_nodes) graph_module.graph.eliminate_dead_code() graph_module.recompile() - return PassResult(graph_module, bool(unsupported_nodes)) + return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index 055567c2be5..2bf5e3f28b3 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/backends/qualcomm/_passes/qnn_pass_manager.py @@ -323,9 +323,7 @@ def get_passes_dependency_for_capture_program(cls): RecomposePixelUnshuffle: [RemoveRedundancy], RecomposeRmsNorm: [RemoveRedundancy], TagQuantIO: [LayoutTransform], - ResolveDebugHandle: [ - TagQuantIO - ], # IMPORTANT: Please always ensure ResolveDebugHandle is the last executed pass. + ResolveDebugHandle: [TagQuantIO], } @classmethod diff --git a/backends/qualcomm/builders/op_batch_norm.py b/backends/qualcomm/builders/op_batch_norm.py index e9675bf2397..a2623cb37b2 100644 --- a/backends/qualcomm/builders/op_batch_norm.py +++ b/backends/qualcomm/builders/op_batch_norm.py @@ -29,6 +29,7 @@ class BatchNorm(NodeVisitor): target = [ "aten._native_batch_norm_legit_no_training.default", "aten._native_batch_norm_legit.no_stats", + "aten._native_batch_norm_legit_functional.default", ] def __init__(self, *args) -> None: diff --git a/backends/qualcomm/debugger/BUCK b/backends/qualcomm/debugger/BUCK index 28b5e68a879..ac2e13cebd0 100644 --- a/backends/qualcomm/debugger/BUCK +++ b/backends/qualcomm/debugger/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "utils", diff --git a/backends/qualcomm/qnn_preprocess.py b/backends/qualcomm/qnn_preprocess.py index 0d8fcd6fe3b..34d1ce05b94 100644 --- a/backends/qualcomm/qnn_preprocess.py +++ b/backends/qualcomm/qnn_preprocess.py @@ -243,6 +243,7 @@ def preprocess_multimethod( # noqa: C901 (handle_id := node.meta.get(DEBUG_HANDLE_KEY)) and QCOM_TENSOR_NAME in node.meta and len(node.meta[QCOM_TENSOR_NAME]) == 1 + and node.op == "call_function" ): debug_handle_builder.insert_delegate_mapping_entry( handles=handle_id, diff --git a/backends/qualcomm/quantizer/annotators/lpai_rules.py b/backends/qualcomm/quantizer/annotators/lpai_rules.py index 178654f18fc..8b30c9427a9 100644 --- a/backends/qualcomm/quantizer/annotators/lpai_rules.py +++ b/backends/qualcomm/quantizer/annotators/lpai_rules.py @@ -129,7 +129,7 @@ class AvgPool2d(GeneralOpDef): # TODO: Batch_norm op cannot directly map to QNN OpBatchnorm due to the number of input doesn't match. @register_annotator( - [torch.ops.aten.batch_norm.default, torch.ops.aten.instance_norm.default], + [torch.ops.aten.batch_norm.default], qnn_op=None, ) class BatchNorm(GeneralOpDef): @@ -420,7 +420,8 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None: torch.ops.aten.topk.default, torch.ops.aten.sort.default, ): - out_act_quantization_spec = SharedQuantizationSpec(node.args[0]) + # assign to None since they are not supported so far + out_act_quantization_spec = None node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( output_qspec=out_act_quantization_spec, _annotated=True, @@ -807,21 +808,6 @@ class ReluMinMax(GeneralOpDef): pass -# TODO: Expand_as op cannot directly map to QNN OpTile due to the number of input doesn't match. -@register_annotator( - [ - torch.ops.aten.expand_as.default, - ], - qnn_op=None, -) -class ExpandAs(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) - - @register_annotator( [ torch.ops.aten.flatten.using_ints, @@ -854,7 +840,6 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None: return act_node = node.args[0] - weight_node = node.args[2] # TODO current only support 16a16w annotate_input_qspec_map( @@ -863,94 +848,23 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None: quantization_config.input_activation, ) - annotate_input_qspec_map( - node, - weight_node, - quantization_config.input_activation, - ) + if len(node.args) > 2 and node.args[2] is not None: + weight_node = node.args[2] + annotate_input_qspec_map( + node, + weight_node, + quantization_config.input_activation, + ) nodes_to_mark_annotated = [node] annotate_output_qspec(node, quantization_config.output_activation) _mark_nodes_as_annotated(nodes_to_mark_annotated) -# TODO: There is a bug in the BackendOpInfo library, so it is bypassed now. -@register_annotator([torch.ops.aten.rsqrt.default], qnn_op=None) -class Rsqrt(GeneralOpDef): - pass - - @register_annotator([torch.ops.aten.scaled_dot_product_attention.default], qnn_op=None) class ScaledDotProductAttention(GeneralOpDef): pass -@register_annotator( - [ - torch.ops.aten.scatter.src, - torch.ops.aten.scatter.value, - torch.ops.aten.scatter_add.default, - torch.ops.aten.scatter_reduce.two, - ], - 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.sort.default], QnnConstants.OpTopK.op_name) -class Sort(GeneralOpDef): - @staticmethod - def annotate(node: Node, quantization_config: QuantizationConfig) -> None: - if _is_annotated([node]): - return - - input_qspec_map = {} - input_act_qspec = quantization_config.input_activation - out_act_quantization_spec = None - if input_act_qspec is not None: - if _is_float_tensor(node.args[0]): - input_act = node.args[0] - assert isinstance(input_act, Node) - input_qspec_map[input_act] = input_act_qspec - out_act_quantization_spec = SharedQuantizationSpec((input_act, node)) - - node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( - input_qspec_map=input_qspec_map, - output_qspec=out_act_quantization_spec, - _annotated=True, - ) - - @register_annotator( [torch.ops.aten.sigmoid, torch.ops.aten.sigmoid.default], QnnConstants.OpSigmoid.op_name, diff --git a/backends/qualcomm/quantizer/quant_recipe.py b/backends/qualcomm/quantizer/quant_recipe.py index b2eb41841f0..347c82c1f82 100644 --- a/backends/qualcomm/quantizer/quant_recipe.py +++ b/backends/qualcomm/quantizer/quant_recipe.py @@ -18,6 +18,10 @@ QuantizationConfig, ) from executorch.backends.qualcomm.quantizer.rules import OpQuantRule +from executorch.backends.qualcomm.utils.check_qnn_version import ( + get_sdk_build_id, + is_qnn_sdk_version_less_than, +) from tabulate import tabulate from torch._ops import OpOverload from torchao.quantization.pt2e import UniformQuantizationObserverBase @@ -91,6 +95,7 @@ def __init__( is_qat=self.is_qat, is_conv_per_channel=True, is_linear_per_channel=True, + is_embedding_per_channel=True, act_observer=self.act_observer, act_symmetric=self.act_symmetric, ) @@ -108,6 +113,15 @@ def get_quant_config(self, node: torch.fx.Node) -> Optional[QuantizationConfig]: if self.granularity == QuantGranularity.PER_TENSOR: return self.quant_config.quant_config elif self.granularity == QuantGranularity.PER_CHANNEL: + if op == torch.ops.aten.embedding.default and is_qnn_sdk_version_less_than( + "2.48" + ): + raise RuntimeError( + "Per-channel embedding quantization requires QNN SDK version " + f">= 2.48. Current QNN SDK version: {get_sdk_build_id()}. Please " + "upgrade your QNN SDK, or use QuantGranularity.PER_TENSOR for the " + "embedding layer instead." + ) ch_axis = self.quant_config.use_per_channel_weight_quant_ops.get(op) assert ( ch_axis is not None diff --git a/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp b/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp index 8bbe047a967..76b67f911cd 100644 --- a/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp +++ b/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp @@ -139,7 +139,41 @@ Error QnnExecuTorchBackend::execute( std::vector input_tensor_structs; std::vector output_tensor_structs; - int args_index = 0; + // The loops below walk the tensor lists recovered from the context binary and + // index args[] with a running counter, so the number of bindable tensors the + // binary declares has to agree with what the program passes. When it does not + // -- a stale binary, or an AOT bug that publishes extra graph I/O -- the walk + // runs off the end of the Span. Count first and fail with both numbers rather + // than reading out of bounds. + size_t bindable_inputs = 0; + for (const auto& input_tensor : input_tensors) { + const auto& name = input_tensor->GetName(); + if (name.find("mutbuf_") == std::string::npos) { + ++bindable_inputs; + } + } + size_t bindable_outputs = 0; + for (const auto& output_tensor : output_tensors) { + const auto& name = output_tensor->GetName(); + if (name.rfind("output_", 0) == 0 && + name.find("mutbuf_") == std::string::npos) { + ++bindable_outputs; + } + } + ET_CHECK_OR_RETURN_ERROR( + bindable_inputs + bindable_outputs == args.size(), + Internal, + "Method %s: the QNN context binary binds %zu tensors (%zu bindable inputs, " + "%zu bindable outputs) but ExecuTorch passed %zu arguments. The binary and " + "the program disagree on the delegate signature; the model has to be " + "re-exported.", + method_name.c_str(), + bindable_inputs + bindable_outputs, + bindable_inputs, + bindable_outputs, + args.size()); + + size_t args_index = 0; input_tensor_structs.reserve(input_tensors.size()); for (const auto& input_tensor : input_tensors) { if (input_tensor->GetName().find("mutbuf_") == std::string::npos) { diff --git a/backends/qualcomm/serialization/qc_compiler_spec.fbs b/backends/qualcomm/serialization/qc_compiler_spec.fbs index 57708c959e9..100404329af 100644 --- a/backends/qualcomm/serialization/qc_compiler_spec.fbs +++ b/backends/qualcomm/serialization/qc_compiler_spec.fbs @@ -47,10 +47,13 @@ enum QcomChipset: int { UNKNOWN_SM = 0, SA8295 = 39, SA8797 = 72, + SC8380XP = 60, + SM7675 = 70, SM8350 = 30, SM8450 = 36, SM8475 = 42, SM8550 = 43, + SM8635 = 68, SM8650 = 57, SM8750 = 69, SM8850 = 87, diff --git a/backends/qualcomm/serialization/qc_schema.py b/backends/qualcomm/serialization/qc_schema.py index aeffbc069b6..692c3f3f217 100644 --- a/backends/qualcomm/serialization/qc_schema.py +++ b/backends/qualcomm/serialization/qc_schema.py @@ -54,10 +54,13 @@ class QcomChipset(IntEnum): UNKNOWN_SM = 0 SA8295 = 39 # v68 SA8797 = 72 # v81 + SC8380XP = 60 # v73 + SM7675 = 70 # v73 SM8350 = 30 # v68 SM8450 = 36 # v69 SM8475 = 42 # v69 SM8550 = 43 # v73 + SM8635 = 68 # v73 SM8650 = 57 # v75 SM8750 = 69 # v79 SM8850 = 87 # v81 @@ -84,11 +87,14 @@ class SocInfo: _soc_info_table = { QcomChipset.SA8295: SocInfo(QcomChipset.SA8295, HtpInfo(HtpArch.V68, 8)), QcomChipset.SA8797: SocInfo(QcomChipset.SA8797, HtpInfo(HtpArch.V81, 16)), + QcomChipset.SC8380XP: SocInfo(QcomChipset.SC8380XP, HtpInfo(HtpArch.V73, 8)), + QcomChipset.SM7675: SocInfo(QcomChipset.SM7675, HtpInfo(HtpArch.V73, 4)), QcomChipset.SM8350: SocInfo(QcomChipset.SM8350, HtpInfo(HtpArch.V68, 4)), QcomChipset.SM8450: SocInfo(QcomChipset.SM8450, HtpInfo(HtpArch.V69, 8)), QcomChipset.SM8475: SocInfo(QcomChipset.SM8475, HtpInfo(HtpArch.V69, 8)), QcomChipset.SM8550: SocInfo(QcomChipset.SM8550, HtpInfo(HtpArch.V73, 8)), QcomChipset.SA8255: SocInfo(QcomChipset.SA8255, HtpInfo(HtpArch.V73, 8)), + QcomChipset.SM8635: SocInfo(QcomChipset.SM8635, HtpInfo(HtpArch.V73, 4)), QcomChipset.SM8650: SocInfo(QcomChipset.SM8650, HtpInfo(HtpArch.V75, 8)), QcomChipset.SM8750: SocInfo(QcomChipset.SM8750, HtpInfo(HtpArch.V79, 8)), QcomChipset.SM8850: SocInfo( diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index ad85bafec06..56b53c6f6c9 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -2943,6 +2943,16 @@ def forward(self, x): ) +class ConvRelu(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, kernel_size=3, padding=1) + self.relu = torch.nn.ReLU() + + def forward(self, x): + return self.relu(self.conv(x)) + + class TopKandIndex(torch.nn.Module): def __init__(self): super().__init__() diff --git a/backends/qualcomm/tests/rework/conftest.py b/backends/qualcomm/tests/rework/conftest.py index 257166bd7c2..9730b120c16 100644 --- a/backends/qualcomm/tests/rework/conftest.py +++ b/backends/qualcomm/tests/rework/conftest.py @@ -32,6 +32,7 @@ get_qnn_context_binary_alignment, prepare_pt2e, QnnConfig, + QnnExecuTorchBackendType, QnnQuantizer, setup_common_args_and_variables, SimpleADB, @@ -269,7 +270,7 @@ def qnn_config(global_setup, request): f'invalid configuration detected, fall back to emulator workload:\n"{e}"' ) config = QnnConfig( - soc_model="unknown", build_folder="build-x86", compile_only=True + soc_model="unknown", build_folder="build-x86", enable_x86_64=True ) return config @@ -349,6 +350,7 @@ def invoke_remote( qnn_config: QnnConfig, executorch_prog: ExecutorchProgramManager, callback: callable, + inputs: Tuple[torch.Tensor] = None, ): with tempfile.TemporaryDirectory() as tmp_dir: pte_fname = f"{tmp_dir}/qnn_executorch_test.pte" @@ -363,7 +365,7 @@ def invoke_remote( pte_path=[pte_fname], workspace=f"/data/local/tmp/{device_workspace}", ) - adb.push() + adb.push(inputs=[inputs] if inputs is not None else None) callback(adb) @@ -478,7 +480,23 @@ def export_and_verify( metrics: Metrics, ): with calibrate(module, [inputs], quantizer) as exported_module: - if quantizer is not None: + fake_tensors = ( + [ + node.meta["val"] + for node in exported_module.graph.nodes + if node.op == "call_function" and "val" in node.meta + ] + if quantizer + else [] + ) + dtypes = set() + for tensor in fake_tensors: + if isinstance(tensor, (tuple, list)): + dtypes.update([n.dtype for n in tensor]) + else: + dtypes.add(tensor.dtype) + + if quantizer and {torch.float, torch.float32} & dtypes: nodes = {node.target for node in exported_module.graph.nodes} q_and_dq = { torch.ops.quantized_decomposed.quantize_per_tensor.default, @@ -505,15 +523,34 @@ def export_and_verify( ) ) execution_plan = executorch_prog.executorch_program.execution_plan[0] + + def validate(): + match qnn_config.backend: + case QnnExecuTorchBackendType.kHtpBackend: + return len(execution_plan.operators) == 0 + case QnnExecuTorchBackendType.kGpuBackend: + return len(execution_plan.operators) == 0 + case QnnExecuTorchBackendType.kLpaiBackend: + aten_op_names = { + op.name + for op in execution_plan.operators + if "quantize" not in op.name + } + return len(aten_op_names) == 0 + case _: + return True + assert all( [ - len(execution_plan.delegates) == 1, - execution_plan.delegates[0].id == "QnnBackend", - len(execution_plan.operators) == 0, + ( + len(execution_plan.delegates) == 1 + and execution_plan.delegates[0].id == "QnnBackend" + ), + validate(), ] ), EXPECT_NOT_FULLY_DELEGATED - mode = "emulator" if qnn_config.build_folder == "build-x86" else "remote" + mode = "emulator" if qnn_config.enable_x86_64 else "remote" globals()[f"verify_output_{mode}"]( module=module, inputs=inputs, diff --git a/backends/qualcomm/tests/rework/gpu/conftest.py b/backends/qualcomm/tests/rework/gpu/conftest.py index b5f86874fd4..4aaf73933cb 100644 --- a/backends/qualcomm/tests/rework/gpu/conftest.py +++ b/backends/qualcomm/tests/rework/gpu/conftest.py @@ -3,3 +3,40 @@ # # 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 pytest + +from executorch.backends.qualcomm.export_utils import ( + generate_gpu_compiler_spec, + generate_qnn_executorch_compiler_spec, + QcomChipset, +) + + +def with_gpu_context(func): + def wrapper(request, kwargs): + preserved = {k: kwargs.pop(k) for k in ["expected"]} + qnn_config = request.getfixturevalue("qnn_config") + fixtures = { + "quantizer": None, + "compile_spec": generate_qnn_executorch_compiler_spec( + soc_model=getattr(QcomChipset, qnn_config.soc_model), + backend_options=generate_gpu_compiler_spec(), + online_prepare=True, + ), + } + return func(request, fixtures | preserved) + + return wrapper + + +def enumerate_fp_dtype(metric: Any): + def wrapper(test_body): + return pytest.mark.parametrize( + "kwargs", + [pytest.param({"act": None, "expected": metric}, id="fp")], + )(test_body) + + return wrapper diff --git a/backends/qualcomm/tests/rework/gpu/feature/conftest.py b/backends/qualcomm/tests/rework/gpu/feature/conftest.py new file mode 100644 index 00000000000..5c3483a7537 --- /dev/null +++ b/backends/qualcomm/tests/rework/gpu/feature/conftest.py @@ -0,0 +1,37 @@ +# 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 inspect +from functools import lru_cache + +import pytest + +from executorch.backends.qualcomm.export_utils import ( + generate_gpu_compiler_spec, + generate_qnn_executorch_compiler_spec, +) + + +@pytest.fixture(scope="session") +def compile_specs(): + @lru_cache() + def _build(kwargs_config): + kwargs = dict(kwargs_config) + et_compile_spec_sig = set( + inspect.signature(generate_qnn_executorch_compiler_spec).parameters.keys() + ) + et_compile_spec_kwargs = { + k: kwargs[k] for k in kwargs.keys() if k in et_compile_spec_sig + } + for k in et_compile_spec_kwargs.keys(): + kwargs.pop(k) + + return generate_qnn_executorch_compiler_spec( + backend_options=generate_gpu_compiler_spec(**kwargs), + **et_compile_spec_kwargs, + ) + + return lambda kwargs_config: _build(kwargs_config) diff --git a/backends/qualcomm/tests/rework/gpu/feature/test.py b/backends/qualcomm/tests/rework/gpu/feature/test.py index b5f86874fd4..f500a0de71d 100644 --- a/backends/qualcomm/tests/rework/gpu/feature/test.py +++ b/backends/qualcomm/tests/rework/gpu/feature/test.py @@ -3,3 +3,85 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.skip(reason="multiple graphs is not supported with online-prepare") +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +# GPU requires online_prepare=True +@pytest.mark.parametrize("kwargs", [pytest.param({"expected": Tolerance()}, id="e2e")]) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize("kwargs", [pytest.param({"expected": Tolerance()}, id="e2e")]) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +# SpillFill is HTP-specific (uses use_multi_contexts / SRAM spill-fill) +@pytest.mark.skip(reason="SpillFill is HTP-specific; not applicable to GPU backend") +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.skip(reason="TBD on native GPU support") +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 + + +# MultiGraph weight sharing is not supported on GPU +@pytest.mark.skip( + reason="Weight sharing across multiple graphs is not supported on GPU backend" +) +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/gpu/op/test.py b/backends/qualcomm/tests/rework/gpu/op/test.py index b5f86874fd4..4f64c6f547b 100644 --- a/backends/qualcomm/tests/rework/gpu/op/test.py +++ b/backends/qualcomm/tests/rework/gpu/op/test.py @@ -3,3 +3,1282 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_FULLY_DELEGATED, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.gpu.conftest import ( + enumerate_fp_dtype, + with_gpu_context, +) + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "expected": pytest.raises( + Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM) + ), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "expected": pytest.raises( + Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM) + ), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "expected": pytest.raises( + AssertionError, match=EXPECT_NOT_FULLY_DELEGATED + ), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +# 3D pooling is not supported on GPU +@enumerate_fp_dtype(pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED)) +@with_gpu_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED)) +@with_gpu_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +# 3D pooling is not supported on GPU +@enumerate_fp_dtype(pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED)) +@with_gpu_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +# 3D convolution is not supported on GPU +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": pytest.raises( + AssertionError, match=EXPECT_NOT_FULLY_DELEGATED + ), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +# 3D convolution is not supported on GPU +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": pytest.raises( + AssertionError, match=EXPECT_NOT_FULLY_DELEGATED + ), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype( + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)) +) +@with_gpu_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance()) +@with_gpu_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +# test_linear_block_quant has no fp variant (lpbq is quantization-specific) +@pytest.mark.skip(reason="LPBQ quantization is not applicable to GPU fp mode") +@pytest.mark.parametrize("kwargs", [pytest.param({}, id="16a4w_lpbq")]) +@with_gpu_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +# 3D pooling is not supported on GPU +@enumerate_fp_dtype(pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED)) +@with_gpu_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED)) +@with_gpu_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(pytest.raises(AssertionError)) +@with_gpu_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(pytest.raises(AssertionError)) +@with_gpu_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype( + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)) +) +@with_gpu_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/feature/v69/test.py b/backends/qualcomm/tests/rework/htp/feature/v69/test.py index b5f86874fd4..85269adaa6a 100644 --- a/backends/qualcomm/tests/rework/htp/feature/v69/test.py +++ b/backends/qualcomm/tests/rework/htp/feature/v69/test.py @@ -3,3 +3,110 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/feature/v73/test.py b/backends/qualcomm/tests/rework/htp/feature/v73/test.py index b5f86874fd4..85269adaa6a 100644 --- a/backends/qualcomm/tests/rework/htp/feature/v73/test.py +++ b/backends/qualcomm/tests/rework/htp/feature/v73/test.py @@ -3,3 +3,110 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/feature/v75/test.py b/backends/qualcomm/tests/rework/htp/feature/v75/test.py index b5f86874fd4..85269adaa6a 100644 --- a/backends/qualcomm/tests/rework/htp/feature/v75/test.py +++ b/backends/qualcomm/tests/rework/htp/feature/v75/test.py @@ -3,3 +3,110 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/feature/v79/test.py b/backends/qualcomm/tests/rework/htp/feature/v79/test.py index b5f86874fd4..85269adaa6a 100644 --- a/backends/qualcomm/tests/rework/htp/feature/v79/test.py +++ b/backends/qualcomm/tests/rework/htp/feature/v79/test.py @@ -3,3 +3,110 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/feature/v81/test.py b/backends/qualcomm/tests/rework/htp/feature/v81/test.py index b5f86874fd4..85269adaa6a 100644 --- a/backends/qualcomm/tests/rework/htp/feature/v81/test.py +++ b/backends/qualcomm/tests/rework/htp/feature/v81/test.py @@ -3,3 +3,110 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/op/v68/test.py b/backends/qualcomm/tests/rework/htp/op/v68/test.py index 03a418a41c8..420f58781d4 100644 --- a/backends/qualcomm/tests/rework/htp/op/v68/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v68/test.py @@ -15,7 +15,6 @@ CosineSimilarity, EXCEPTION_EXIR_PROGRAM, EXCEPTION_FROM_PASSES, - EXPECT_NOT_ANNOTATED, EXPECT_NOT_FULLY_DELEGATED, SkipOutputCheck, Tolerance, @@ -151,25 +150,13 @@ def test_amin(request, kwargs): AMin.test(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(rtol=1e-1), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_any(request, kwargs): Any.test(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(rtol=1e-1), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_arange_dtype_int(request, kwargs): Arange.test_dtype_int(request, kwargs) # noqa: F405 @@ -255,8 +242,8 @@ def test_batchnorm_2d(request, kwargs): @enumerate_activation_dtype( [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), ] ) @@ -265,13 +252,7 @@ def test_bitwise_and_numeric(request, kwargs): BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(rtol=1e-1), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_bitwise_and_bool(request, kwargs): BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 @@ -279,8 +260,8 @@ def test_bitwise_and_bool(request, kwargs): @enumerate_activation_dtype( [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), ] ) @@ -289,13 +270,7 @@ def test_bitwise_or_numeric(request, kwargs): BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(rtol=1e-1), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_bitwise_or_bool(request, kwargs): BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 @@ -303,8 +278,8 @@ def test_bitwise_or_bool(request, kwargs): @enumerate_activation_dtype( [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), ] ) @@ -313,13 +288,7 @@ def test_bitwise_xor_numeric(request, kwargs): BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(rtol=1e-1), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_bitwise_xor_bool(request, kwargs): BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 @@ -818,25 +787,13 @@ def test_interpolate_nearest(request, kwargs): Interpolate.test_nearest(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(rtol=1e-1), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_is_inf(request, kwargs): IsInf.test(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance()]) @with_htp_context def test_is_nan(request, kwargs): IsNan.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/op/v69/test.py b/backends/qualcomm/tests/rework/htp/op/v69/test.py index b5f86874fd4..3babe431a5c 100644 --- a/backends/qualcomm/tests/rework/htp/op/v69/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v69/test.py @@ -3,3 +3,1411 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import re +from functools import partial +from pathlib import Path + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + CosineSimilarity, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_FULLY_DELEGATED, + SkipOutputCheck, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.htp.conftest import ( + enumerate_activation_dtype, + with_htp_context, +) + +# e.g. get 69 from ".../rework/htp/unit_test/op/v69/test.py" +HTP_ARCH = int(re.search(r".*v([0-9]+)$", Path(__file__).parent.name).group(1)) +with_htp_context = partial(with_htp_context, hw_arch=HTP_ARCH) + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +# addmm 16a requires V73+; still fails on V69 +@enumerate_activation_dtype( + [ + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + Tolerance(rtol=1e-1), + ] +) +@with_htp_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +# bmm 16a requires V73+; still fails on V69 +@enumerate_activation_dtype( + [ + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + Tolerance(rtol=1e-1), + ] +) +@with_htp_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +# LPBQ (QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION) requires V69+; enabled here +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "lpbq": True, + "block_sz_map": {"conv2d": (1, 32, 1, 1)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + ], +) +@with_htp_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +# einsum 16a requires V73+; still fails on V69 +@enumerate_activation_dtype( + [ + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + Tolerance(rtol=1e-1), + ] +) +@with_htp_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, False, Tolerance(), "8a8w_ptq"), + (16, 16, False, Tolerance(), "16a16w_ptq"), + (16, 8, True, Tolerance(), "16a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + CosineSimilarity(0.95), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance()]) +@with_htp_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "pcq": False, + "lpbq": True, + "block_sz_map": {"linear": (1, 32)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + ], +) +@with_htp_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 16, "param": 16, "pcq": False, "expected": Tolerance()}, + id="16a16w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 4, "pcq": True, "expected": CosineSimilarity(0.95)}, + id="16a4w_pcq", + ), + pytest.param( + {"act": 16, "param": 8, "pcq": True, "expected": Tolerance()}, + id="16a8w_pcq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 2, "pcq": True, "expected": CosineSimilarity(0.9)}, + id="16a2w_pcq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 16, + "pcq": False, + "expected": pytest.raises(AssertionError, match=Tolerance()), + }, + id="16a16w_ptq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + SkipOutputCheck(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +# sdpa 16a requires V73+; still fails on V69 +@enumerate_activation_dtype( + [ + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + Tolerance(rtol=1e-1), + ] +) +@with_htp_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/op/v73/test.py b/backends/qualcomm/tests/rework/htp/op/v73/test.py index b5f86874fd4..6362690f5fd 100644 --- a/backends/qualcomm/tests/rework/htp/op/v73/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v73/test.py @@ -3,3 +3,1387 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import re +from functools import partial +from pathlib import Path + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + CosineSimilarity, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_FULLY_DELEGATED, + SkipOutputCheck, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.htp.conftest import ( + enumerate_activation_dtype, + with_htp_context, +) + +# e.g. get 73 from ".../rework/htp/unit_test/op/v73/test.py" +HTP_ARCH = int(re.search(r".*v([0-9]+)$", Path(__file__).parent.name).group(1)) +with_htp_context = partial(with_htp_context, hw_arch=HTP_ARCH) + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +# addmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +# bmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +# LPBQ (QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION) requires V69+; enabled here +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "lpbq": True, + "block_sz_map": {"conv2d": (1, 32, 1, 1)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + ], +) +@with_htp_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +# einsum 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, False, Tolerance(), "8a8w_ptq"), + (16, 16, False, Tolerance(), "16a16w_ptq"), + (16, 8, True, Tolerance(), "16a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + CosineSimilarity(0.95), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance()]) +@with_htp_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "pcq": False, + "lpbq": True, + "block_sz_map": {"linear": (1, 32)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + ], +) +@with_htp_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 16, "param": 16, "pcq": False, "expected": Tolerance()}, + id="16a16w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 4, "pcq": True, "expected": CosineSimilarity(0.95)}, + id="16a4w_pcq", + ), + pytest.param( + {"act": 16, "param": 8, "pcq": True, "expected": Tolerance()}, + id="16a8w_pcq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 2, "pcq": True, "expected": CosineSimilarity(0.9)}, + id="16a2w_pcq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 16, + "pcq": False, + "expected": pytest.raises(AssertionError, match=Tolerance()), + }, + id="16a16w_ptq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + SkipOutputCheck(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +# sdpa 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/op/v75/test.py b/backends/qualcomm/tests/rework/htp/op/v75/test.py index b5f86874fd4..6362690f5fd 100644 --- a/backends/qualcomm/tests/rework/htp/op/v75/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v75/test.py @@ -3,3 +3,1387 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import re +from functools import partial +from pathlib import Path + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + CosineSimilarity, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_FULLY_DELEGATED, + SkipOutputCheck, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.htp.conftest import ( + enumerate_activation_dtype, + with_htp_context, +) + +# e.g. get 73 from ".../rework/htp/unit_test/op/v73/test.py" +HTP_ARCH = int(re.search(r".*v([0-9]+)$", Path(__file__).parent.name).group(1)) +with_htp_context = partial(with_htp_context, hw_arch=HTP_ARCH) + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +# addmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +# bmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +# LPBQ (QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION) requires V69+; enabled here +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "lpbq": True, + "block_sz_map": {"conv2d": (1, 32, 1, 1)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + ], +) +@with_htp_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +# einsum 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, False, Tolerance(), "8a8w_ptq"), + (16, 16, False, Tolerance(), "16a16w_ptq"), + (16, 8, True, Tolerance(), "16a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + CosineSimilarity(0.95), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance()]) +@with_htp_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "pcq": False, + "lpbq": True, + "block_sz_map": {"linear": (1, 32)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + ], +) +@with_htp_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 16, "param": 16, "pcq": False, "expected": Tolerance()}, + id="16a16w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 4, "pcq": True, "expected": CosineSimilarity(0.95)}, + id="16a4w_pcq", + ), + pytest.param( + {"act": 16, "param": 8, "pcq": True, "expected": Tolerance()}, + id="16a8w_pcq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 2, "pcq": True, "expected": CosineSimilarity(0.9)}, + id="16a2w_pcq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 16, + "pcq": False, + "expected": pytest.raises(AssertionError, match=Tolerance()), + }, + id="16a16w_ptq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + SkipOutputCheck(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +# sdpa 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/op/v79/test.py b/backends/qualcomm/tests/rework/htp/op/v79/test.py index b5f86874fd4..6362690f5fd 100644 --- a/backends/qualcomm/tests/rework/htp/op/v79/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v79/test.py @@ -3,3 +3,1387 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import re +from functools import partial +from pathlib import Path + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + CosineSimilarity, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_FULLY_DELEGATED, + SkipOutputCheck, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.htp.conftest import ( + enumerate_activation_dtype, + with_htp_context, +) + +# e.g. get 73 from ".../rework/htp/unit_test/op/v73/test.py" +HTP_ARCH = int(re.search(r".*v([0-9]+)$", Path(__file__).parent.name).group(1)) +with_htp_context = partial(with_htp_context, hw_arch=HTP_ARCH) + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +# addmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +# bmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +# LPBQ (QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION) requires V69+; enabled here +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "lpbq": True, + "block_sz_map": {"conv2d": (1, 32, 1, 1)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + ], +) +@with_htp_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +# einsum 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, False, Tolerance(), "8a8w_ptq"), + (16, 16, False, Tolerance(), "16a16w_ptq"), + (16, 8, True, Tolerance(), "16a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + CosineSimilarity(0.95), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance()]) +@with_htp_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "pcq": False, + "lpbq": True, + "block_sz_map": {"linear": (1, 32)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + ], +) +@with_htp_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 16, "param": 16, "pcq": False, "expected": Tolerance()}, + id="16a16w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 4, "pcq": True, "expected": CosineSimilarity(0.95)}, + id="16a4w_pcq", + ), + pytest.param( + {"act": 16, "param": 8, "pcq": True, "expected": Tolerance()}, + id="16a8w_pcq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 2, "pcq": True, "expected": CosineSimilarity(0.9)}, + id="16a2w_pcq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 16, + "pcq": False, + "expected": pytest.raises(AssertionError, match=Tolerance()), + }, + id="16a16w_ptq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + SkipOutputCheck(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +# sdpa 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/op/v81/test.py b/backends/qualcomm/tests/rework/htp/op/v81/test.py index b5f86874fd4..6362690f5fd 100644 --- a/backends/qualcomm/tests/rework/htp/op/v81/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v81/test.py @@ -3,3 +3,1387 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import re +from functools import partial +from pathlib import Path + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + CosineSimilarity, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_FULLY_DELEGATED, + SkipOutputCheck, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.htp.conftest import ( + enumerate_activation_dtype, + with_htp_context, +) + +# e.g. get 73 from ".../rework/htp/unit_test/op/v73/test.py" +HTP_ARCH = int(re.search(r".*v([0-9]+)$", Path(__file__).parent.name).group(1)) +with_htp_context = partial(with_htp_context, hw_arch=HTP_ARCH) + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +# addmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +# bmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +# LPBQ (QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION) requires V69+; enabled here +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "lpbq": True, + "block_sz_map": {"conv2d": (1, 32, 1, 1)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + ], +) +@with_htp_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +# einsum 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, False, Tolerance(), "8a8w_ptq"), + (16, 16, False, Tolerance(), "16a16w_ptq"), + (16, 8, True, Tolerance(), "16a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + CosineSimilarity(0.95), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance()]) +@with_htp_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "pcq": False, + "lpbq": True, + "block_sz_map": {"linear": (1, 32)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + ], +) +@with_htp_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 16, "param": 16, "pcq": False, "expected": Tolerance()}, + id="16a16w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 4, "pcq": True, "expected": CosineSimilarity(0.95)}, + id="16a4w_pcq", + ), + pytest.param( + {"act": 16, "param": 8, "pcq": True, "expected": Tolerance()}, + id="16a8w_pcq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 2, "pcq": True, "expected": CosineSimilarity(0.9)}, + id="16a2w_pcq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 16, + "pcq": False, + "expected": pytest.raises(AssertionError, match=Tolerance()), + }, + id="16a16w_ptq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + SkipOutputCheck(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +# sdpa 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/lpai/conftest.py b/backends/qualcomm/tests/rework/lpai/conftest.py index b5f86874fd4..9553051b08f 100644 --- a/backends/qualcomm/tests/rework/lpai/conftest.py +++ b/backends/qualcomm/tests/rework/lpai/conftest.py @@ -3,3 +3,100 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from functools import lru_cache +from typing import Any, List + +import pytest + +from executorch.backends.qualcomm.export_utils import ( + generate_lpai_compiler_spec, + generate_qnn_executorch_compiler_spec, + make_quantizer, + QcomChipset, + QnnExecuTorchBackendType, + QuantDtype, +) +from executorch.backends.qualcomm.serialization.qc_schema import ( + LpaiHardwareVersion, + QnnExecuTorchLpaiTargetEnv, +) + + +def with_lpai_context(func, hw_arch): + def wrapper(request, kwargs): + # extend this if necessary + preserved = {k: kwargs.pop(k) for k in ["expected"]} + callbacks_and_args = { + # extract objects from callback + "quantizers": {"arch": hw_arch} | kwargs, + "compile_specs": {"arch": hw_arch}, + } + fixtures = { + k[:-1]: request.getfixturevalue(k)(**v) + for k, v in callbacks_and_args.items() + } + return func(request, fixtures | preserved) + + return wrapper + + +def enumerate_activation_dtype(metrics: List[Any]): + def wrapper(test_body): + return pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"act": act, "expected": metrics[i]}, id=id) + for i, (act, id) in enumerate( + [ + (8, "8a"), + ] + ) + ], + )(test_body) + + return wrapper + + +def _get_lpai_arch(): + # hardcoded lpai architecture with corresponding premium soc + return [ + (LpaiHardwareVersion.V6, "SM8850"), + ] + + +@pytest.fixture(scope="session") +def quantizers(): + arch_to_soc = dict(_get_lpai_arch()) + + @lru_cache() + def _build(arch, act, param, per_ch): + attr = f"use_{act}a{param}w" + if quant_dtype := getattr(QuantDtype, attr, None): + return make_quantizer( + quant_dtype=quant_dtype, + per_channel_conv=per_ch, + per_channel_linear=per_ch, + backend=QnnExecuTorchBackendType.kLpaiBackend, + soc_model=arch_to_soc[arch], + ) + + def get_quantizer(arch, act, param=None, pcq=False, **_): + param = 8 if (param is None and act is not None) else param + return _build(arch, act, param, pcq) + + return get_quantizer + + +@pytest.fixture(scope="session") +def compile_specs(): + compile_spec = { + arch: generate_qnn_executorch_compiler_spec( + soc_model=getattr(QcomChipset, soc_model), + backend_options=generate_lpai_compiler_spec( + target_env=QnnExecuTorchLpaiTargetEnv.kX86, + ), + ) + for (arch, soc_model) in _get_lpai_arch() + } + return lambda arch: compile_spec[arch] diff --git a/backends/qualcomm/tests/rework/lpai/feature/conftest.py b/backends/qualcomm/tests/rework/lpai/feature/conftest.py new file mode 100644 index 00000000000..b168d759654 --- /dev/null +++ b/backends/qualcomm/tests/rework/lpai/feature/conftest.py @@ -0,0 +1,37 @@ +# 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 inspect +from functools import lru_cache + +import pytest + +from executorch.backends.qualcomm.export_utils import ( + generate_lpai_compiler_spec, + generate_qnn_executorch_compiler_spec, +) + + +@pytest.fixture(scope="session") +def compile_specs(): + @lru_cache() + def _build(kwargs_config): + kwargs = dict(kwargs_config) + et_compile_spec_sig = set( + inspect.signature(generate_qnn_executorch_compiler_spec).parameters.keys() + ) + et_compile_spec_kwargs = { + k: kwargs[k] for k in kwargs.keys() if k in et_compile_spec_sig + } + for k in et_compile_spec_kwargs.keys(): + kwargs.pop(k) + + return generate_qnn_executorch_compiler_spec( + backend_options=generate_lpai_compiler_spec(**kwargs), + **et_compile_spec_kwargs, + ) + + return lambda kwargs_config: _build(kwargs_config) diff --git a/backends/qualcomm/tests/rework/lpai/feature/v6/test.py b/backends/qualcomm/tests/rework/lpai/feature/v6/test.py index b5f86874fd4..9e49d0b9ba3 100644 --- a/backends/qualcomm/tests/rework/lpai/feature/v6/test.py +++ b/backends/qualcomm/tests/rework/lpai/feature/v6/test.py @@ -3,3 +3,85 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize("kwargs", [pytest.param({"expected": Tolerance()}, id="e2e")]) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +# LPAI forbids online_prepare=True at runtime +@pytest.mark.skip( + reason="LPAI backend only supports offline_prepare; online_prepare is forbidden" +) +@pytest.mark.parametrize("kwargs", [pytest.param({"expected": Tolerance()}, id="e2e")]) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.skip(reason="TBD on native LPAI support") +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize("kwargs", [pytest.param({"expected": Tolerance()}, id="e2e")]) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +# SpillFill is HTP-specific (uses use_multi_contexts / SRAM spill-fill) +@pytest.mark.skip(reason="SpillFill is HTP-specific; not applicable to LPAI backend") +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 + + +# MultiGraph weight sharing requires use_weight_sharing in generate_lpai_compiler_spec (not supported) +@pytest.mark.skip( + reason="Weight sharing across multiple graphs is not supported on LPAI backend" +) +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/lpai/op/v6/test.py b/backends/qualcomm/tests/rework/lpai/op/v6/test.py index b5f86874fd4..81eb670eebe 100644 --- a/backends/qualcomm/tests/rework/lpai/op/v6/test.py +++ b/backends/qualcomm/tests/rework/lpai/op/v6/test.py @@ -3,3 +3,1611 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import re +from functools import partial +from pathlib import Path + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_ANNOTATED, + EXPECT_NOT_FULLY_DELEGATED, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.lpai.conftest import ( + enumerate_activation_dtype, + with_lpai_context, +) + + +# e.g. get 68 from ".../rework/htp/unit_test/op/v68/test.py" +LPAI_ARCH = int(re.search(r".*v([0-9]+)$", Path(__file__).parent.name).group(1)) +with_lpai_context = partial(with_lpai_context, hw_arch=LPAI_ARCH) + + +# abs not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +# acos not in lpai_rules but will be decomposed into equivalent ops +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +# adaptive_avg_pool3d not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +# addmm decomposes to mm+add before annotation; both in lpai_rules +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +# int64 cast for indices is not supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +# int64 cast for indices is not supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +# asin not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +# some decomposed ops are not supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +# some decomposed ops are not supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +# avg_pool3d not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +# ceil not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +# channel_shuffle not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +# clone not in lpai_rules but will be omitted +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + ], +) +@with_lpai_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + ], +) +@with_lpai_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + ], +) +@with_lpai_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + ], +) +@with_lpai_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + ], +) +@with_lpai_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +# conv3d not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +# conv3d_transpose not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +# cos not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +# cumsum not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +# einsum decomposes to bmm/matmul before annotation; both in lpai_rules +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +# elu not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + ], +) +@with_lpai_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +# equal not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +# expand not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +# expand_as not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +# expm1 not in lpai_rules but will be decomposed +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +# fill translates to static tensor in QNN; backend-agnostic +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +# floor not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +# floor_divide not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +# fold uses col2im which is in lpai_rules (ColIm, qnn_op=None) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_lpai_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +# full/full_like translate to static tensors in QNN; backend-agnostic +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +# gather is in lpai_rules (Embedding class handles index/gather/index_select) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +# glu decomposes to chunk+sigmoid+mul before annotation; all in lpai_rules +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +# greater not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +# greater_equal not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +# grid_sample not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +# grid_sample not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +# group_norm not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +# hardsigmoid: DecomposeHardsigmoid runs before annotation; decomposed ops in lpai_rules +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +# decomposed ops might not be supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +# decomposed ops might not be supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +# decomposed ops might not be supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +# instance_norm not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +# registered by the partitioner to not be decomposed but failed to be delegated +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_lpai_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +# maps to prelu +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +# less_equal not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +# less_than not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +# decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +# LPBQ not applicable to LPAI v6 (requires HTP V69+ feature) +@pytest.mark.skip(reason="LPBQ quantization is not supported on LPAI v6") +@with_lpai_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + ], +) +@with_lpai_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + ], +) +@with_lpai_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +# log10 not in lpai_rules but will be decomposed +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +# log1p not in lpai_rules but will be decomposed +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +# log2 not in lpai_rules but will be decomposed +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +# logical_and not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +# logical_not not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +# decomposed ops are not supported with invalid weight fallback triggered +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_lpai_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +# cast op for indices is not supported +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +# max_pool3d not in lpai_rules and the decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +# cast op for indices is not supported +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +# neg not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +# not_equal not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +# registered by the partitioner to not be decomposed but failed to be delegated +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_lpai_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +# rand not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +# reciprocal: DecomposeReciprocal decomposes to div(1, x); div is in lpai_rules +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +# registered by the partitioner to not be decomposed but failed to be delegated +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_lpai_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +# registered by the partitioner to not be decomposed but failed to be delegated +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_lpai_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +# reflection_pad3d not in lpai_rules but will be decomposed +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +# relu6 decomposes to hardtanh(0, 6); hardtanh is in lpai_rules (ReluMinMax) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +# decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +# repeat not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +# roll not in lpai_rules but decomposed ops are fully delegated +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +# round not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +# rsqrt not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +# scatter.src not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +# select_copy maps to aten.select.int which is in lpai_rules (StrideSlice) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +# select_scatter not in lpai_rules and decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +# sign not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +# sin not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +# slice_copy maps to aten.slice.Tensor which is in lpai_rules (StrideSlice) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +# slice_scatter not in lpai_rules and decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +# scatter not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +# sort not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +# square is in lpai_rules (Pow class handles square.default) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +# stack maps to OpPack which is HTP-specific +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +# tan not in lpai_rules and the decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +# threshold not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +# triu not in lpai_rulesdecomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +# triu not in lpai_rulesdecomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +# trunc not in lpai_rules and decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +# topk not in lpai_rules, use EXPECT_NOT_FULLY_DELEGATED for there are +# other ops in the test body +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +# unbind maps to OpUnpack which is HTP-specific +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +# unfold uses im2col which is in lpai_rules (ColIm, qnn_op=None) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_lpai_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +# where not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +# var not in lpai_rules but will be decomposed +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/src/feature.py b/backends/qualcomm/tests/rework/src/feature.py index 8ef6d1abec4..49499fc04e6 100644 --- a/backends/qualcomm/tests/rework/src/feature.py +++ b/backends/qualcomm/tests/rework/src/feature.py @@ -14,15 +14,24 @@ import torch +from executorch.backends.qualcomm.debugger.qcom_numerical_comparator_sample import ( + QcomCosineSimilarityComparator, +) +from executorch.backends.qualcomm.debugger.qnn_intermediate_debugger import ( + QNNIntermediateDebugger, +) from executorch.backends.qualcomm.export_utils import ( make_quantizer, QcomChipset, + QnnConfig, QnnExecuTorchBackendType, QnnExecuTorchHtpPerformanceMode, SimpleADB, to_edge_transform_and_lower_to_qnn, ) from executorch.backends.qualcomm.serialization.qc_schema import ( + QnnExecuTorchGpuPerformanceMode, + QnnExecuTorchLpaiClientPerf, QnnExecuTorchProfileLevel, ) from executorch.backends.qualcomm.tests.rework.conftest import ( @@ -51,6 +60,17 @@ def wrapper(request, kwargs): return wrapper +def get_quantizer(qnn_config: QnnConfig): + return ( + make_quantizer( + backend=qnn_config.backend, + soc_model=qnn_config.soc_model, + ) + if qnn_config.backend != QnnExecuTorchBackendType.kGpuBackend + else None + ) + + class Logging: class Model(torch.nn.Module): def __init__(self): @@ -62,6 +82,15 @@ def example_inputs(self): def forward(self, x): return torch.nn.ReLU()(x) + @staticmethod + def _get_log_pattern(backend): + return { + QnnExecuTorchBackendType.kHtpBackend: "QnnDsp ", + QnnExecuTorchBackendType.kGpuBackend: "OpenCL", + # looks like no special keyword appears + QnnExecuTorchBackendType.kLpaiBackend: "", + }[backend] + @staticmethod def _test(qnn_config, compile_specs, expected, aot): def callback(adb: SimpleADB, pattern): @@ -78,9 +107,7 @@ def verify(log): model = __class__.Model() inputs = model.example_inputs() # perform ptq - with calibrate( - model, [inputs], make_quantizer(soc_model=qnn_config.soc_model) - ) as model: + with calibrate(model, [inputs], get_quantizer(qnn_config)) as model: # start lowering executorch_prog_mgr = to_edge_transform_and_lower_to_qnn( module=model, @@ -91,28 +118,28 @@ def verify(log): invoke_remote( qnn_config=qnn_config, executorch_prog=executorch_prog_mgr, - callback=partial(callback, pattern="QnnDsp "), + callback=partial( + callback, + pattern=Logging._get_log_pattern(qnn_config.backend), + ), ) @staticmethod @unpack_fixtures def test(subtests, qnn_config, compile_specs, expected): - # extend this for other backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { QnnExecuTorchBackendType.kHtpBackend: [ - compile_specs(tuple(d.items())) - for d in [ - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "debug": True, - "use_fp16": False, - }, - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "debug": False, - "use_fp16": False, - }, - ] + {"soc_model": soc_model, "debug": True, "use_fp16": False}, + {"soc_model": soc_model, "debug": False, "use_fp16": False}, + ], + QnnExecuTorchBackendType.kGpuBackend: [ + {"soc_model": soc_model, "debug": True, "online_prepare": True}, + {"soc_model": soc_model, "debug": False, "online_prepare": True}, + ], + QnnExecuTorchBackendType.kLpaiBackend: [ + {"soc_model": soc_model, "debug": True}, + {"soc_model": soc_model, "debug": False}, ], } @@ -120,7 +147,9 @@ def test(subtests, qnn_config, compile_specs, expected): with subtests.test(msg=config): __class__._test( qnn_config=qnn_config, - compile_specs=backend_compile_specs[qnn_config.backend][i], + compile_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend][i].items()) + ), expected=expected, aot=config == "compile_time_option", ) @@ -152,7 +181,7 @@ def compile(models, compile_specs): with calibrate( models[i], [inputs], - make_quantizer(soc_model=qnn_config.soc_model), + get_quantizer(qnn_config), ) as model: modules_dict[graph_name] = model sample_inputs_dict[graph_name] = inputs @@ -221,21 +250,24 @@ def test_weight_sharing(qnn_config, compile_specs, expected): @staticmethod @unpack_fixtures def test_inference(qnn_config, compile_specs, expected): - # extend this for other backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { - QnnExecuTorchBackendType.kHtpBackend: compile_specs( - tuple( - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "use_fp16": False, - }.items() - ) - ), + QnnExecuTorchBackendType.kHtpBackend: { + "soc_model": soc_model, + "use_fp16": False, + }, + QnnExecuTorchBackendType.kGpuBackend: { + "soc_model": soc_model, + "online_prepare": True, + }, + QnnExecuTorchBackendType.kLpaiBackend: {"soc_model": soc_model}, } __class__._test( qnn_config=qnn_config, - compile_specs=backend_compile_specs[qnn_config.backend], + compile_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend].items()) + ), expected=expected, ) @@ -254,27 +286,28 @@ def forward(self, x): @staticmethod @unpack_fixtures def test(qnn_config, compile_specs, expected): - # extend this for other backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { - QnnExecuTorchBackendType.kHtpBackend: compile_specs( - tuple( - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "online_prepare": True, - "use_fp16": False, - }.items() - ) - ), + QnnExecuTorchBackendType.kHtpBackend: { + "soc_model": soc_model, + "online_prepare": True, + "use_fp16": False, + }, + QnnExecuTorchBackendType.kGpuBackend: { + "soc_model": soc_model, + "online_prepare": True, + }, } module = __class__.Model() - qnn_config.online_prepare = True export_and_verify( module=module, inputs=module.example_inputs(), qnn_config=qnn_config, - quantizer=make_quantizer(soc_model=qnn_config.soc_model), - compile_specs=backend_compile_specs[qnn_config.backend], + quantizer=get_quantizer(qnn_config), + compile_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend].items()) + ), metrics=expected, ) @@ -304,53 +337,82 @@ def verify(log): adb.extra_cmds += "" if aot else " --htp_performance_mode 6" adb.execute(output_callback=verify) + # TODO: extend performance check for following backends + def callback_gpu(adb: SimpleADB): + adb.execute() + + def callback_lpai(adb: SimpleADB): + adb.execute() + with expected: # model declaration model = __class__.Model() inputs = model.example_inputs() # perform ptq - with calibrate( - model, [inputs], make_quantizer(soc_model=qnn_config.soc_model) - ) as model: + with calibrate(model, [inputs], get_quantizer(qnn_config)) as model: # start lowering executorch_prog_mgr = to_edge_transform_and_lower_to_qnn( module=model, inputs=inputs, compiler_specs=compile_specs, ).to_executorch() - # verifier per backend dispatcher = { QnnExecuTorchBackendType.kHtpBackend: callback_htp, + QnnExecuTorchBackendType.kGpuBackend: callback_gpu, + QnnExecuTorchBackendType.kLpaiBackend: callback_lpai, } # remote testing invoke_remote( qnn_config=qnn_config, executorch_prog=executorch_prog_mgr, - callback=partial(dispatcher[qnn_config.backend], voltage=80), + callback=( + partial(dispatcher[qnn_config.backend], voltage=80) + if qnn_config.backend == QnnExecuTorchBackendType.kHtpBackend + else dispatcher[qnn_config.backend] + ), ) @staticmethod @unpack_fixtures def test(subtests, qnn_config, compile_specs, expected): - # extend this for other backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { QnnExecuTorchBackendType.kHtpBackend: [ - compile_specs(tuple(d.items())) - for d in [ - # compile_time option - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "debug": True, - "use_fp16": False, - "htp_performance_mode": QnnExecuTorchHtpPerformanceMode.kHtpHighPowerSaver, - }, - # runtime_option (performance mode defaults to kHtpBurst) - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "debug": True, - "use_fp16": False, - }, - ] + # compile_time option + { + "soc_model": soc_model, + "debug": True, + "use_fp16": False, + "htp_performance_mode": QnnExecuTorchHtpPerformanceMode.kHtpHighPowerSaver, + }, + # runtime_option (performance mode defaults to kHtpBurst) + {"soc_model": soc_model, "debug": True, "use_fp16": False}, + ], + QnnExecuTorchBackendType.kGpuBackend: [ + # compile_time option: set low perf hint to GPU + { + "soc_model": soc_model, + "online_prepare": True, + "performance_mode": QnnExecuTorchGpuPerformanceMode.kGpuPerfHintLow, + }, + # runtime_option: scaffold — GPUruntime perf hint not yet wired in C++ + # TODO: extend GPU runtime to accept dynamic performance settings + { + "soc_model": soc_model, + "debug": True, + "online_prepare": True, + }, + ], + QnnExecuTorchBackendType.kLpaiBackend: [ + { + "soc_model": soc_model, + "fps": 30, + "ftrt_ratio": 10, + "client_perf_type": QnnExecuTorchLpaiClientPerf.kRealTime, + }, + # runtime_option: scaffold — LPAI runtime perf hint not yet wired in C++ + # TODO: extend LPAI runtime to accept dynamic performance settings + {"soc_model": soc_model, "debug": True}, ], } @@ -358,7 +420,9 @@ def test(subtests, qnn_config, compile_specs, expected): with subtests.test(msg=config): __class__._test( qnn_config=qnn_config, - compile_specs=backend_compile_specs[qnn_config.backend][i], + compile_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend][i].items()) + ), expected=expected, aot=config == "compile_time_option", ) @@ -409,9 +473,7 @@ def callback(adb: SimpleADB, executorch_prog_mgr, expected_profile_events): model = __class__.Model() inputs = model.example_inputs() # perform ptq - with calibrate( - model, [inputs], make_quantizer(soc_model=qnn_config.soc_model) - ) as model: + with calibrate(model, [inputs], get_quantizer(qnn_config)) as model: # start lowering executorch_prog_mgr = to_edge_transform_and_lower_to_qnn( module=model, @@ -426,30 +488,43 @@ def callback(adb: SimpleADB, executorch_prog_mgr, expected_profile_events): callback=partial( callback, executorch_prog_mgr=executorch_prog_mgr, - expected_profile_events=20, + expected_profile_events=2, ), ) @staticmethod @unpack_fixtures def test(subtests, qnn_config, compile_specs, expected): - # extend this for other backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { QnnExecuTorchBackendType.kHtpBackend: [ - compile_specs(tuple(d.items())) - for d in [ - # compile_time option - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "profile_level": QnnExecuTorchProfileLevel.kProfileDetailed, - "use_fp16": False, - }, - # runtime_option - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "use_fp16": False, - }, - ] + # compile_time option + { + "soc_model": soc_model, + "profile_level": QnnExecuTorchProfileLevel.kProfileDetailed, + "use_fp16": False, + }, + # runtime_option + {"soc_model": soc_model, "use_fp16": False}, + ], + QnnExecuTorchBackendType.kGpuBackend: [ + # compile_time option + { + "soc_model": soc_model, + "profile_level": QnnExecuTorchProfileLevel.kProfileDetailed, + "online_prepare": True, + }, + # runtime_option + {"soc_model": soc_model, "online_prepare": True}, + ], + QnnExecuTorchBackendType.kLpaiBackend: [ + # compile_time option + { + "soc_model": soc_model, + "profile_level": QnnExecuTorchProfileLevel.kProfileDetailed, + }, + # runtime_option + {"soc_model": soc_model}, ], } @@ -457,7 +532,9 @@ def test(subtests, qnn_config, compile_specs, expected): with subtests.test(msg=config): __class__._test( qnn_config=qnn_config, - compile_specs=backend_compile_specs[qnn_config.backend][i], + compile_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend][i].items()) + ), expected=expected, aot=config == "compile_time_option", ) @@ -482,17 +559,23 @@ def test(qnn_config, compile_specs, expected): option_to_flatbuffer, ) - # extend this for other backends + # saver=True is a top-level QnnExecuTorchOptions field; works across backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { - QnnExecuTorchBackendType.kHtpBackend: compile_specs( - tuple( - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "saver": True, - "use_fp16": False, - }.items() - ) - ), + QnnExecuTorchBackendType.kHtpBackend: { + "soc_model": soc_model, + "saver": True, + "use_fp16": False, + }, + QnnExecuTorchBackendType.kGpuBackend: { + "soc_model": soc_model, + "saver": True, + "online_prepare": True, + }, + QnnExecuTorchBackendType.kLpaiBackend: { + "soc_model": soc_model, + "saver": True, + }, } with expected: @@ -500,13 +583,13 @@ def test(qnn_config, compile_specs, expected): model = __class__.Model() inputs = model.example_inputs() # perform ptq - with calibrate( - model, [inputs], make_quantizer(soc_model=qnn_config.soc_model) - ) as model: + with calibrate(model, [inputs], get_quantizer(qnn_config)) as model: # start lowering with tempfile.TemporaryDirectory() as tmp_dir: # hack saver output folder - cs = backend_compile_specs[qnn_config.backend] + cs = compile_specs( + tuple(backend_compile_specs[qnn_config.backend].items()) + ) option = flatbuffer_to_option(cs[0].value) option.saver_output_dir = f"{tmp_dir}/saver_output" cs[0].value = option_to_flatbuffer(option) @@ -539,17 +622,23 @@ def forward(self, x): @staticmethod @unpack_fixtures def test(qnn_config, compile_specs, expected): - # extend this for other backends + # shared_buffer=True is a top-level QnnExecuTorchOptions field; works across backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { - QnnExecuTorchBackendType.kHtpBackend: compile_specs( - tuple( - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "shared_buffer": True, - "use_fp16": False, - }.items() - ) - ), + QnnExecuTorchBackendType.kHtpBackend: { + "soc_model": soc_model, + "shared_buffer": True, + "use_fp16": False, + }, + QnnExecuTorchBackendType.kGpuBackend: { + "soc_model": soc_model, + "shared_buffer": True, + "online_prepare": True, + }, + QnnExecuTorchBackendType.kLpaiBackend: { + "soc_model": soc_model, + "shared_buffer": True, + }, } module = __class__.Model() @@ -558,8 +647,10 @@ def test(qnn_config, compile_specs, expected): module=module, inputs=module.example_inputs(), qnn_config=qnn_config, - quantizer=make_quantizer(soc_model=qnn_config.soc_model), - compile_specs=backend_compile_specs[qnn_config.backend], + quantizer=get_quantizer(qnn_config), + compile_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend].items()) + ), metrics=expected, ) @@ -605,9 +696,7 @@ def test(qnn_config, compile_specs, expected): # perform ptq model = __class__.Model() inputs = model.example_inputs() - with calibrate( - model, [inputs], make_quantizer(soc_model=qnn_config.soc_model) - ) as model: + with calibrate(model, [inputs], get_quantizer(qnn_config)) as model: # start lowering edge_prog_mgr = to_edge_transform_and_lower_to_qnn( module=model, @@ -621,22 +710,23 @@ def test(qnn_config, compile_specs, expected): class TensorDump: + # Simple Conv2d+ReLU model that is supported by all backends class Model(torch.nn.Module): def __init__(self): super().__init__() - self.idx_source = torch.rand(10, 3) + self.conv = torch.nn.Conv2d(3, 8, kernel_size=3, padding=1) + self.relu = torch.nn.ReLU() def example_inputs(self): - return (torch.randn(3, 10),) + return (torch.randn(1, 3, 8, 8),) def forward(self, x): - a, b = torch.topk(x, 3) - return a + self.idx_source[b] + return self.relu(self.conv(x)) @staticmethod @unpack_fixtures def test(qnn_config, compile_specs, expected): - def callback(adb: SimpleADB, expected_intermediate_events): + def callback(adb: SimpleADB, debugger, expected_compared_events): with tempfile.TemporaryDirectory() as tmp_dir: etdump_path = f"{tmp_dir}/etdump.etdp" debug_output_path = f"{tmp_dir}/debug_output.bin" @@ -644,49 +734,82 @@ def callback(adb: SimpleADB, expected_intermediate_events): adb.pull_debug_output( etdump_path=etdump_path, debug_buffer_path=debug_output_path ) - inspector = Inspector( - etdump_path=etdump_path, debug_buffer_path=debug_output_path + debugger.setup_inspector( + etdump_path=etdump_path, + debug_buffer_path=debug_output_path, ) - for event_block in inspector.event_blocks: - if event_block.name == "Execute": - assert ( - len(event_block.events) == expected_intermediate_events - ), ( - f"unexpected number of intermediate events, expecting " - f"{expected_intermediate_events}, but has {len(event_block.events)} events.", - ) + comparator = debugger.create_comparator(QcomCosineSimilarityComparator) + numeric_results = debugger.inspector.calculate_numeric_gap( + distance=comparator, + reference_graph=debugger.reference_graph_name, + ) + numeric_results = numeric_results.set_index("runtime_debug_handle") + assert len(numeric_results) == expected_compared_events, ( + f"unexpected number of compared events, expecting " + f"{expected_compared_events}, but has {len(numeric_results)} events." + ) + for _, row in numeric_results.iterrows(): + assert comparator.is_valid_score(row.gap[0]), ( + f"Node {row.aot_ops} is failing " + f"{comparator.metric_name()} test, {row.gap[0]} is lower " + f"than {comparator.threshold}." + ) - # extend this for other backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) + # dump_intermediate_outputs=True is a top-level QnnExecuTorchOptions field; works across backends backend_compile_specs = { - QnnExecuTorchBackendType.kHtpBackend: compile_specs( - tuple( - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "dump_intermediate_outputs": True, - "use_fp16": False, - }.items() - ) - ), + QnnExecuTorchBackendType.kHtpBackend: { + "soc_model": soc_model, + "dump_intermediate_outputs": True, + "use_fp16": False, + }, + QnnExecuTorchBackendType.kGpuBackend: { + "soc_model": soc_model, + "dump_intermediate_outputs": True, + "online_prepare": True, + }, + QnnExecuTorchBackendType.kLpaiBackend: { + "soc_model": soc_model, + "dump_intermediate_outputs": True, + }, } with expected: # perform ptq model = __class__.Model() inputs = model.example_inputs() - with calibrate( - model, [inputs], make_quantizer(soc_model=qnn_config.soc_model) - ) as model: + with calibrate(model, [inputs], get_quantizer(qnn_config)) as model: # start lowering executorch_prog_mgr = to_edge_transform_and_lower_to_qnn( module=model, inputs=inputs, - compiler_specs=backend_compile_specs[qnn_config.backend], + compiler_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend].items()) + ), generate_etrecord=True, ).to_executorch() - # remote testing - qnn_config.dump_intermediate_outputs = True - invoke_remote( - qnn_config=qnn_config, - executorch_prog=executorch_prog_mgr, - callback=partial(callback, expected_intermediate_events=9), - ) + + with tempfile.TemporaryDirectory() as etrecord_dir: + etrecord_path = f"{etrecord_dir}/etrecord.bin" + etrecord = executorch_prog_mgr.get_etrecord() + debugger = QNNIntermediateDebugger(inputs) + debugger.set_etrecord_file_path(etrecord_path) + debugger.set_edge_ep( + edge_ep=etrecord.graph_map[debugger.reference_graph_name] + ) + etrecord.update_representative_inputs(debugger.sample_input) + etrecord.save(etrecord_path) + + # remote testing + qnn_config.dump_intermediate_outputs = True + invoke_remote( + qnn_config=qnn_config, + executorch_prog=executorch_prog_mgr, + inputs=inputs, + # conv + relu = 2 intermediate outputs + callback=partial( + callback, + debugger=debugger, + expected_compared_events=2, + ), + ) diff --git a/backends/qualcomm/tests/rework/src/op.py b/backends/qualcomm/tests/rework/src/op.py index 963285725cb..a782235732f 100644 --- a/backends/qualcomm/tests/rework/src/op.py +++ b/backends/qualcomm/tests/rework/src/op.py @@ -3198,7 +3198,7 @@ def forward(self, x): @unpack_fixtures def test(subtests, qnn_config, quantizer, compile_spec, expected): inputs = (torch.randn(1, 4, 8, 8),) - dims = [-1, 1, 2] + dims = [-1, 3] for dim in dims: with subtests.test(msg=f"dim:{dim}"): with expected as metrics: diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index 1139085ec31..3c05a816c6d 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -7787,6 +7787,32 @@ def test_qnn_backend_dump_intermediate_outputs_simple_model(self): expected_compared_events=expected_compared_events, ) + def test_qnn_backend_dump_intermediate_outputs_conv_relu(self): + match get_backend_type(self.backend): + case QnnExecuTorchBackendType.kHtpBackend: + backend_options = generate_htp_compiler_spec(use_fp16=False) + case QnnExecuTorchBackendType.kLpaiBackend: + backend_options = generate_lpai_compiler_spec( + target_env=self.get_lpai_target_env() + ) + case _: + raise ValueError("Backend is not implemented yet") + TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=self.chipset_table[TestQNN.soc_model], + backend_options=backend_options, + dump_intermediate_outputs=True, + ) + sample_input = (torch.randn(1, 3, 8, 8),) + module = ConvRelu() # noqa: F405 + module = self.get_qdq_module(module, sample_input) + + self.lower_module_and_test_output( + module, + sample_input, + expected_partitions=1, + expected_compared_events=2, + ) + def test_qnn_backend_dump_intermediate_outputs_topk(self): torch.manual_seed(8) backend_options = generate_htp_compiler_spec(use_fp16=False) diff --git a/backends/qualcomm/utils/qnn_sdk_setup.py b/backends/qualcomm/utils/qnn_sdk_setup.py index 9ef1b66392c..bce3b761f9e 100644 --- a/backends/qualcomm/utils/qnn_sdk_setup.py +++ b/backends/qualcomm/utils/qnn_sdk_setup.py @@ -166,7 +166,8 @@ def disable_mkldnn_on_amd() -> None: import torch - torch.backends.mkldnn.enabled = False + if not torch.backends.flags_frozen(): + torch.backends.mkldnn.enabled = False def _host_is_amd() -> bool: diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index 4a1715e8c03..f78d47525ad 100644 --- a/backends/qualcomm/utils/utils.py +++ b/backends/qualcomm/utils/utils.py @@ -34,6 +34,7 @@ QnnExecuTorchBackendOptions, QnnExecuTorchBackendType, QnnExecuTorchGpuBackendOptions, + QnnExecuTorchGpuPerformanceMode, QnnExecuTorchGpuPrecision, QnnExecuTorchHtpBackendOptions, QnnExecuTorchHtpPerformanceMode, @@ -1018,6 +1019,7 @@ def draw_graph(title, path, graph_module: torch.fx.GraphModule, format=DrawForma def generate_gpu_compiler_spec( + performance_mode: QnnExecuTorchGpuPerformanceMode = QnnExecuTorchGpuPerformanceMode.kGpuPerfHintHigh, precision: QnnExecuTorchGpuPrecision = QnnExecuTorchGpuPrecision.kGpuPrecisionUserProvided, use_memory_optimizations: bool = True, use_node_optimizations: bool = True, @@ -1028,6 +1030,8 @@ def generate_gpu_compiler_spec( Helper function generating backend options for QNN HTP Args: + performance_mode: + kGpuPerfHintHigh / kGpuPerfHintNormal / kGpuPerfHintLow precision: kGpuPrecisionFp32 - Sets the precision mode to floating point 32-bit (FP32). kGpuPrecisionFp16 - Sets the precision mode to floating point 16-bit (FP16). @@ -1046,6 +1050,7 @@ def generate_gpu_compiler_spec( """ # TODO: enable performance hint mechanism in runtime and make this as an option gpu_options = QnnExecuTorchGpuBackendOptions() + gpu_options.performance_mode = performance_mode gpu_options.precision = precision gpu_options.use_memory_optimizations = use_memory_optimizations gpu_options.use_node_optimizations = use_node_optimizations @@ -1194,9 +1199,11 @@ def generate_qnn_executorch_compiler_spec( # noqa: C901 Args: soc_model: The SoC you plan to run the compiled model. Please check QcomChipset for supported SoC. + SM7675(Snapdragon 7+ Gen 3) SM8450 (Snapdragon 8 Gen 1) SM8475(Snapdragon 8 Gen 1+) SM8550(Snapdragon 8 Gen 2) + SM8635(Snapdragon 8s Gen 3) SM8650(Snapdragon 8 Gen 3) SM8750(Snapdragon 8 Elite) SM8850(Snapdragon 8 Elite Gen 5) @@ -1322,11 +1329,14 @@ def get_soc_to_htp_arch_map(): return { "SA8295": HtpArch.V68, "SA8797": HtpArch.V81, + "SC8380XP": HtpArch.V73, "SM8350": HtpArch.V68, "SM8450": HtpArch.V69, "SM8475": HtpArch.V69, "SM8550": HtpArch.V73, + "SM7675": HtpArch.V73, "SA8255": HtpArch.V73, + "SM8635": HtpArch.V73, "SM8650": HtpArch.V75, "SM8750": HtpArch.V79, "SM8850": HtpArch.V81, @@ -1355,11 +1365,14 @@ def get_soc_to_chipset_map(): return { "SA8295": QcomChipset.SA8295, "SA8797": QcomChipset.SA8797, + "SC8380XP": QcomChipset.SC8380XP, + "SM7675": QcomChipset.SM7675, "SM8350": QcomChipset.SM8350, "SM8450": QcomChipset.SM8450, "SM8475": QcomChipset.SM8475, "SM8550": QcomChipset.SM8550, "SA8255": QcomChipset.SA8255, + "SM8635": QcomChipset.SM8635, "SM8650": QcomChipset.SM8650, "SM8750": QcomChipset.SM8750, "SM8850": QcomChipset.SM8850, diff --git a/backends/samsung/CMakeLists.txt b/backends/samsung/CMakeLists.txt index 1f647c5bbe2..c8b4bedbc4a 100644 --- a/backends/samsung/CMakeLists.txt +++ b/backends/samsung/CMakeLists.txt @@ -146,24 +146,6 @@ if(${ANDROID}) executorch_target_link_options_shared_lib(enn_backend) target_compile_options(enn_backend PRIVATE -Wno-deprecated-declarations) - set(__enn_executor_runner_srcs - ${EXECUTORCH_SOURCE_DIR}/examples/samsung/executor_runner/enn_executor_runner.cpp - ) - add_executable(enn_executor_runner ${__enn_executor_runner_srcs}) - add_dependencies(enn_executor_runner enn_backend) - target_link_libraries( - enn_executor_runner - PRIVATE enn_logging - enn_backend - gflags - executorch - extension_data_loader - portable_ops_lib - android - ) - set_target_properties( - enn_executor_runner PROPERTIES CXX_VISIBILITY_PRESET hidden - ) install( TARGETS enn_backend enn_logging EXPORT ExecuTorchTargets diff --git a/backends/samsung/README.md b/backends/samsung/README.md index bc48bad830a..30c855fe365 100644 --- a/backends/samsung/README.md +++ b/backends/samsung/README.md @@ -49,13 +49,16 @@ Generates python artifacts that allow user call `Compile` interface to lower a m ./backends/samsung/build.sh -b x86_64 ``` -### Build ENN Executor Runner +### Build Backend Delegate ```bash ./backends/samsung/build.sh -b android --ndk ${ANDROID_NDK} ``` -ANDROID_ABI=arm64-v8a is default, necessary runtime executable generated in `build_exynos_android` directory. +ANDROID_ABI=arm64-v8a is default, necessary runtime backend library generated in `build_samsung_android` directory. -### Build Anroid Extension +### Build Executable +Please see the [README.md](../../examples/samsung/README.md). + +### Build Android Extension This is later exposed Java app. Please turn on CMake option `EXECUTORCH_BUILD_ENN`, and ENN runtime will be added. ```bash cmake extension/android \ @@ -64,7 +67,7 @@ cmake extension/android \ -DCMAKE_INSTALL_PREFIX=cmake-android-out \ -Bcmake-android-out/extension/android -cmake --build cmake-android-out/extension/android -j8 +cmake --build cmake-android-out/extension/android -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) ``` ## Examples diff --git a/backends/samsung/_passes/remove_useless_ops.py b/backends/samsung/_passes/remove_useless_ops.py index c88a2d4a5d8..9c965844749 100644 --- a/backends/samsung/_passes/remove_useless_ops.py +++ b/backends/samsung/_passes/remove_useless_ops.py @@ -15,7 +15,6 @@ class RemoveUselessOpPass(ExportPass): USELESS_OP_SET = { exir_ops.edge.aten._to_copy.default, exir_ops.edge.aten.clone.default, - exir_ops.edge.aten.clone.default, exir_ops.edge.aten.alias.default, exir_ops.edge.aten.lift_fresh_copy.default, exir_ops.edge.dim_order_ops._to_dim_order_copy.default, diff --git a/backends/samsung/_passes/replace_scalar_ops.py b/backends/samsung/_passes/replace_scalar_ops.py index 8ae54b0dc98..22a74c15f61 100644 --- a/backends/samsung/_passes/replace_scalar_ops.py +++ b/backends/samsung/_passes/replace_scalar_ops.py @@ -38,9 +38,16 @@ def call_operator( if op not in self._ops_with_scalar: return super().call_operator(op, args, kwargs, meta) + # For pow operation, convert int scalar to float32 tensor + # because the PowVisitor requires both inputs to be float32 + if op == exir_ops.edge.aten.pow.Tensor_Scalar and isinstance(args[1], int): + args1 = torch.tensor(float(args[1]), dtype=torch.float32) + else: + args1 = torch.tensor(args[1]) + return super().call_operator( op=self._ops_with_scalar.get(op, op), - args=(args[0], torch.tensor(args[1])), + args=(args[0], args1), kwargs=kwargs, meta=meta, ) diff --git a/backends/samsung/build.sh b/backends/samsung/build.sh index a4871feb50c..e6258152785 100755 --- a/backends/samsung/build.sh +++ b/backends/samsung/build.sh @@ -66,7 +66,17 @@ function build_android() { ANDROID_ABI=arm64-v8a ANDROID_PLATFORM=android-28 # Trace requires over android-23 + local host_flatcc=${X86_64_BUILD_DIR}/third-party/flatcc_ep/bin/flatcc + local flatcc_args=() + if [[ -x ${host_flatcc} ]]; then + flatcc_args=(-DFLATCC_EXECUTABLE=${host_flatcc}) + else + echo "Warning: ${host_flatcc} not found. Build the x86_64 target first" \ + "('-b x86_64' or '-b all') if executor_runner fails to link libflatccrt.a." + fi + cmake \ + "${flatcc_args[@]}" \ -DCMAKE_INSTALL_PREFIX=${ANDROID_BUILD_DIR} \ -DCMAKE_TOOLCHAIN_FILE="${ANDROID_NDK_ROOT}/build/cmake/android.toolchain.cmake" \ -DANDROID_NDK=${ANDROID_NDK} \ diff --git a/backends/samsung/builders/__init__.py b/backends/samsung/builders/__init__.py index 14b9a17a6c9..60f97cf60f7 100644 --- a/backends/samsung/builders/__init__.py +++ b/backends/samsung/builders/__init__.py @@ -18,6 +18,7 @@ op_dequantize, op_div, op_embedding, + op_exp, op_expand_copy, op_gelu, op_getitem, @@ -48,6 +49,7 @@ op_select, op_sigmoid, op_sin, + op_skip, op_slice_copy, op_softmax, op_split_with_sizes_copy, @@ -77,6 +79,7 @@ op_dequantize, op_div, op_embedding, + op_exp, op_expand_copy, op_gelu, op_getitem, @@ -107,6 +110,7 @@ op_select, op_sigmoid, op_sin, + op_skip, op_slice_copy, op_softmax, op_split_with_sizes_copy, diff --git a/backends/samsung/builders/node_visitor.py b/backends/samsung/builders/node_visitor.py index 0d2707da8f5..cb7d4db690a 100644 --- a/backends/samsung/builders/node_visitor.py +++ b/backends/samsung/builders/node_visitor.py @@ -31,7 +31,7 @@ def __init__(self, exported_program: ExportedProgram) -> None: def exported_program(self) -> ExportedProgram: return self._exported_program - def define_node(self, node: torch.fx.Node, enn_graph: EnnGraph): + def define_node(self, node: torch.fx.Node, enn_graph: EnnGraph) -> bool: raise NotImplementedError("NodeVisitor must be extended!") def define_tensor( @@ -58,7 +58,9 @@ def define_tensor( if is_param_node(self.exported_program, node): if swap_nc_for_weights: tensor = torch.swapdims(tensor, 0, 1) - const_data = tensor.contiguous().detach().numpy() + if not isinstance(tensor, torch._subclasses.fake_tensor.FakeTensor): + # .numpy() is not supported for tensor subclasses if the tensor is a fake tensor. + const_data = tensor.contiguous().detach().numpy() dims = [1] if len(tensor.size()) == 0 else list(tensor.size()) diff --git a/backends/samsung/builders/op_add.py b/backends/samsung/builders/op_add.py index a6eb79897dd..177f700f7e1 100644 --- a/backends/samsung/builders/op_add.py +++ b/backends/samsung/builders/op_add.py @@ -4,6 +4,7 @@ # 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 from typing import Dict import torch @@ -26,16 +27,22 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input1 = node.args[0] input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) params = {} self._update_params_qdtype(node, params) input2 = node.args[1] input_id_2 = self.define_tensor(input2, enn_graph, vals_to_ids) + alpha = node.kwargs.get("alpha", 1.0) + if alpha != 1.0: + logging.warning("Currently, only alpha 1 for add is supported.") + return False output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op( node.name, "ELTSUM", [input_id_1, input_id_2], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_avg_pool2d.py b/backends/samsung/builders/op_avg_pool2d.py index bfca8b89b22..529a3156030 100644 --- a/backends/samsung/builders/op_avg_pool2d.py +++ b/backends/samsung/builders/op_avg_pool2d.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -52,10 +52,6 @@ def define_node( params["explicit_padding"] = explicit_padding self._update_params_qdtype(node, params) - if len(node.args) > 4: - ceil_mode = cast(bool, node.args[4]) - assert not ceil_mode, "Not support ceil_mode = True." - if len(node.args) > 5: params["count_include_pad"] = cast(bool, node.args[5]) else: @@ -68,3 +64,5 @@ def define_node( ), "Not supported divisor_override which is not equal to pooling region." output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "AVGPOOL2D", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_batch_norm.py b/backends/samsung/builders/op_batch_norm.py index e5373a8223a..990b8e0ca20 100644 --- a/backends/samsung/builders/op_batch_norm.py +++ b/backends/samsung/builders/op_batch_norm.py @@ -25,7 +25,16 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: + index_zero_getitems = [] + for user in node.users.keys(): + if user.target.__name__ != "getitem": + continue + if user.args[1] == 0: + index_zero_getitems.append(user) + elif len(user.users) > 0: + return False + all_input_tensors = [] input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -51,6 +60,11 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids, output_idx=0) + for getitem in index_zero_getitems: + vals_to_ids[getitem] = output_id + enn_graph.define_op( node.name, "BatchNormalization", all_input_tensors, [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_bmm.py b/backends/samsung/builders/op_bmm.py index 13e0d19cb14..2ac96a4a1b7 100644 --- a/backends/samsung/builders/op_bmm.py +++ b/backends/samsung/builders/op_bmm.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input1 = node.args[0] input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) @@ -41,3 +41,5 @@ def define_node( enn_graph.define_op( node.name, "BATCH_MATMUL", [input_id_1, input_id_2], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_cat.py b/backends/samsung/builders/op_cat.py index 09387f2e361..82762cfcdbf 100644 --- a/backends/samsung/builders/op_cat.py +++ b/backends/samsung/builders/op_cat.py @@ -28,7 +28,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: tensors = cast(List[torch.fx.Node], node.args[0]) input_tensor_ids = [] constant_idx = None @@ -48,3 +48,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "CONCAT", input_tensor_ids, [output_id], params) + + return True diff --git a/backends/samsung/builders/op_clamp.py b/backends/samsung/builders/op_clamp.py index 74af83212a5..b69066c3337 100644 --- a/backends/samsung/builders/op_clamp.py +++ b/backends/samsung/builders/op_clamp.py @@ -4,6 +4,7 @@ # 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 from typing import cast, Dict import torch @@ -11,6 +12,7 @@ NodeVisitor, register_node_visitor, ) +from executorch.backends.samsung.builders.utils import get_tensor from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph @@ -26,9 +28,13 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) + input_tensor = get_tensor(self.exported_program, input) + if input_tensor.dtype == torch.int64: + logging.warning("Currently, int64 clip is unsupported!") + return False # The default value of lower bound and upper bound output_min = torch.finfo(torch.float32).min @@ -45,3 +51,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "CLIP", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_constant_pad_nd.py b/backends/samsung/builders/op_constant_pad_nd.py index 006f52619ff..bc72305bc13 100644 --- a/backends/samsung/builders/op_constant_pad_nd.py +++ b/backends/samsung/builders/op_constant_pad_nd.py @@ -29,7 +29,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -54,3 +54,5 @@ def define_node( } self._update_params_qdtype(node, params) enn_graph.define_op(node.name, "PAD", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_conv2d.py b/backends/samsung/builders/op_conv2d.py index ab77d8df626..87f76f610bf 100644 --- a/backends/samsung/builders/op_conv2d.py +++ b/backends/samsung/builders/op_conv2d.py @@ -4,6 +4,7 @@ # 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 from typing import cast, Dict, List import torch @@ -27,7 +28,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: all_input_tensors = [] input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -52,9 +53,24 @@ def define_node( padding = cast(List[int], node.args[4]) dilation = cast(List[int], node.args[5]) groups = cast(int, node.args[8]) + if is_transpose_conv and groups != 1: + logging.warning("Don't support groups for transpose conv.") + return False + output_padding = cast(List[int], node.args[7]) + if is_transpose_conv and output_padding != [0, 0]: + logging.warning("Don't support output padding for transpose conv.") + return False + if len(padding) < 2: + logging.warning( + "For conv1d decomposed to conv2d(with conv1d params), Conv1dToConv2d pass will update the params." + ) + return True explicit_padding = [padding[0], padding[1], padding[0], padding[1]] input_shape = get_shape(input) + if len(input_shape) > 4: + logging.warning("Currently, only conv2d is supported.") + return False kernel_shape = get_shape(weight_node) params = {} self._update_params_qdtype(node, params) @@ -72,7 +88,7 @@ def define_node( params["explicit_padding"] = explicit_padding params["in_channels"] = input_shape[1] params["out_channels"] = kernel_shape[0] * kernel_shape[1] * groups - params["out_channels"] //= input_shape[1] * input_shape[0] + params["out_channels"] //= input_shape[1] output_id = self.define_tensor(node, enn_graph, vals_to_ids) @@ -87,3 +103,5 @@ def define_node( enn_graph.define_op( node.name, conv_type, all_input_tensors, [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_cos.py b/backends/samsung/builders/op_cos.py index bd746db91cd..f1d4d917b7d 100644 --- a/backends/samsung/builders/op_cos.py +++ b/backends/samsung/builders/op_cos.py @@ -23,9 +23,11 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "Cos", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_div.py b/backends/samsung/builders/op_div.py index 8b0e7cdd5af..7afc23220fe 100644 --- a/backends/samsung/builders/op_div.py +++ b/backends/samsung/builders/op_div.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input1 = node.args[0] input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) @@ -40,3 +40,5 @@ def define_node( enn_graph.define_op( node.name, "ELTDIV", [input_id_1, input_id_2], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_embedding.py b/backends/samsung/builders/op_embedding.py index a500ea051fd..bc28ae4e55e 100644 --- a/backends/samsung/builders/op_embedding.py +++ b/backends/samsung/builders/op_embedding.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: weight_node = node.args[0] weight_id = self.define_tensor(weight_node, enn_graph, vals_to_ids) @@ -40,3 +40,5 @@ def define_node( enn_graph.define_op( node.name, "GATHER", [weight_id, input_id], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_exp.py b/backends/samsung/builders/op_exp.py new file mode 100644 index 00000000000..63dd9ca3219 --- /dev/null +++ b/backends/samsung/builders/op_exp.py @@ -0,0 +1,33 @@ +# Copyright (c) 2026 Samsung Electronics Co. LTD +# 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 torch +from executorch.backends.samsung.builders.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph + + +@register_node_visitor +class ExpVisitor(NodeVisitor): + target = "aten.exp.default" + + def define_node( + self, + node: torch.fx.Node, + enn_graph: EnnGraph, + vals_to_ids: Dict[torch.Tensor, int], + ) -> bool: + input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) + + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + enn_graph.define_op(node.name, "Exp", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_expand_copy.py b/backends/samsung/builders/op_expand_copy.py index f4c707b8e62..5fb6b6c5166 100644 --- a/backends/samsung/builders/op_expand_copy.py +++ b/backends/samsung/builders/op_expand_copy.py @@ -4,6 +4,7 @@ # 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 from typing import cast, Dict, List import torch @@ -27,7 +28,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ): + ) -> bool: # inputs input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -35,6 +36,8 @@ def define_node( in_shape = get_shape(input) sizes = cast(List[int], node.args[1]) expand_dims = self.check_expand_dims(sizes, in_shape) + if expand_dims is None: + return False # output output_id = self.define_tensor(node, enn_graph, vals_to_ids) @@ -53,7 +56,10 @@ def define_node( params, ) else: - raise NotImplementedError("Don't support expanding at more than one axes.") + logging.warning("Don't support expanding at more than one axes.") + return False + + return True def check_expand_dims(self, sizes, in_shape): expand_dims = [] @@ -72,6 +78,8 @@ def check_expand_dims(self, sizes, in_shape): while new_size_index > 0: new_size_index -= 1 - assert sizes[new_size_index] == 1, "Current expand is unsupported!" + if sizes[new_size_index] != 1: + logging.warning("Current expand is unsupported!") + return None return expand_dims diff --git a/backends/samsung/builders/op_gelu.py b/backends/samsung/builders/op_gelu.py index 88417f688f9..6a064194561 100644 --- a/backends/samsung/builders/op_gelu.py +++ b/backends/samsung/builders/op_gelu.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: # input1 input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -38,3 +38,5 @@ def define_node( self._update_params_qdtype(node, params) enn_graph.define_op(node.name, "GELU", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_getitem.py b/backends/samsung/builders/op_getitem.py index 901ec73cf7d..bc7c0441e3c 100644 --- a/backends/samsung/builders/op_getitem.py +++ b/backends/samsung/builders/op_getitem.py @@ -28,5 +28,5 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: - return + ) -> bool: + return True diff --git a/backends/samsung/builders/op_group_norm.py b/backends/samsung/builders/op_group_norm.py index 55c7bb6732a..b0509005053 100644 --- a/backends/samsung/builders/op_group_norm.py +++ b/backends/samsung/builders/op_group_norm.py @@ -23,7 +23,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: all_input_tensors = [] input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) all_input_tensors.append(input_id) @@ -44,3 +44,5 @@ def define_node( enn_graph.define_op( node.name, "GROUPNORM", all_input_tensors, [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_hardsigmoid.py b/backends/samsung/builders/op_hardsigmoid.py index 3a50d65da41..58cc0a12d5b 100644 --- a/backends/samsung/builders/op_hardsigmoid.py +++ b/backends/samsung/builders/op_hardsigmoid.py @@ -26,10 +26,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) params = {} self._update_params_qdtype(node, params) enn_graph.define_op(node.name, "HardSigmoid", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_hardswish.py b/backends/samsung/builders/op_hardswish.py index 8c30125e8a4..fc9ec134418 100644 --- a/backends/samsung/builders/op_hardswish.py +++ b/backends/samsung/builders/op_hardswish.py @@ -26,10 +26,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) params = {} self._update_params_qdtype(node, params) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "HARDSWISH", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_hardtanh.py b/backends/samsung/builders/op_hardtanh.py index 7d65e97a566..4c60d7227dc 100644 --- a/backends/samsung/builders/op_hardtanh.py +++ b/backends/samsung/builders/op_hardtanh.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -40,3 +40,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "CLIP", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_index.py b/backends/samsung/builders/op_index.py index b7765e35b3f..0145616de56 100644 --- a/backends/samsung/builders/op_index.py +++ b/backends/samsung/builders/op_index.py @@ -23,7 +23,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -47,3 +47,5 @@ def define_node( enn_graph.define_op( node.name, "GATHER", [input_id, indices_id], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_layer_norm.py b/backends/samsung/builders/op_layer_norm.py index 098bc92dc84..937168c36e9 100644 --- a/backends/samsung/builders/op_layer_norm.py +++ b/backends/samsung/builders/op_layer_norm.py @@ -25,7 +25,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: all_input_tensors = [] input_node = node.args[0] input_id = self.define_tensor(input_node, enn_graph, vals_to_ids) @@ -51,3 +51,5 @@ def define_node( enn_graph.define_op( node.name, "LAYERNORM", all_input_tensors, [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_leaky_relu.py b/backends/samsung/builders/op_leaky_relu.py index c7ed37d12e5..28f4a7851d4 100644 --- a/backends/samsung/builders/op_leaky_relu.py +++ b/backends/samsung/builders/op_leaky_relu.py @@ -27,7 +27,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: all_input_tensors = [] input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) all_input_tensors.append(input_id) @@ -56,3 +56,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "PRELU", all_input_tensors, [output_id]) + + return True diff --git a/backends/samsung/builders/op_linear.py b/backends/samsung/builders/op_linear.py index 720439de976..dffb6b0108c 100644 --- a/backends/samsung/builders/op_linear.py +++ b/backends/samsung/builders/op_linear.py @@ -27,7 +27,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: all_input_tensors = [] input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -49,3 +49,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "FC", all_input_tensors, [output_id], params) + + return True diff --git a/backends/samsung/builders/op_log.py b/backends/samsung/builders/op_log.py index 97127dd94ba..a20de9bd95d 100644 --- a/backends/samsung/builders/op_log.py +++ b/backends/samsung/builders/op_log.py @@ -23,10 +23,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "LOG", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_log_softmax.py b/backends/samsung/builders/op_log_softmax.py index f2d87601cbb..f26a36e2af7 100644 --- a/backends/samsung/builders/op_log_softmax.py +++ b/backends/samsung/builders/op_log_softmax.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ): + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -37,3 +37,5 @@ def define_node( meta_data = {"axis": axis} enn_graph.define_op(node.name, "LOGSOFTMAX", [input_id], [output_id], meta_data) + + return True diff --git a/backends/samsung/builders/op_max_pool2d.py b/backends/samsung/builders/op_max_pool2d.py index 57b716fcb34..bec8da5f683 100644 --- a/backends/samsung/builders/op_max_pool2d.py +++ b/backends/samsung/builders/op_max_pool2d.py @@ -25,7 +25,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -75,10 +75,6 @@ def define_node( params["dilation_w"] = dilation[1] self._update_params_qdtype(node, params) - if len(node.args) > 5: - ceil_mode = cast(bool, node.args[5]) - assert not ceil_mode, "Not support ceil_mode = True." - if not is_indices: output_id = self.define_tensor( node, @@ -94,3 +90,5 @@ def define_node( ) enn_graph.define_op(node.name, "MAXPOOL2D", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_maximum.py b/backends/samsung/builders/op_maximum.py index d3358d736f3..d81dfdd4a35 100644 --- a/backends/samsung/builders/op_maximum.py +++ b/backends/samsung/builders/op_maximum.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: # inputs input_id_1 = self.define_tensor(node.args[0], enn_graph, vals_to_ids) input_id_2 = self.define_tensor(node.args[1], enn_graph, vals_to_ids) @@ -35,3 +35,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "MAXIMUM", [input_id_1, input_id_2], [output_id]) + + return True diff --git a/backends/samsung/builders/op_mean_dim.py b/backends/samsung/builders/op_mean_dim.py index 3d0377703a7..3396102c045 100644 --- a/backends/samsung/builders/op_mean_dim.py +++ b/backends/samsung/builders/op_mean_dim.py @@ -4,6 +4,7 @@ # 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 from typing import cast, Dict, List import torch @@ -11,6 +12,7 @@ NodeVisitor, register_node_visitor, ) +from executorch.backends.samsung.builders.utils import get_tensor from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph from executorch.backends.transforms import get_shape @@ -27,7 +29,11 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: + output_tensor = get_tensor(self.exported_program, node) + if output_tensor.dtype == torch.float64: + logging.warning("float64 for mean has not supported yet.") + return False # input input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -37,6 +43,9 @@ def define_node( dims = cast(List[int], node.args[1]) reduce_axes = [] in_shape = get_shape(input) + if dims is None: + logging.warning("dims is None for this case.") + return False for dim in dims: reduce_axes.append(dim % len(in_shape)) @@ -47,3 +56,5 @@ def define_node( params = {"keep_dims": keep_dim, "axis": reduce_axes} self._update_params_qdtype(node, params) enn_graph.define_op(node.name, "REDUCEMEAN", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_minimum.py b/backends/samsung/builders/op_minimum.py index a32b462d45f..4c612df34c7 100644 --- a/backends/samsung/builders/op_minimum.py +++ b/backends/samsung/builders/op_minimum.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: # inputs input1 = node.args[0] input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) @@ -37,3 +37,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "MIN", [input_id_1, input_id_2], [output_id]) + + return True diff --git a/backends/samsung/builders/op_mul.py b/backends/samsung/builders/op_mul.py index 6dd7c0dd9f0..e703ddff253 100644 --- a/backends/samsung/builders/op_mul.py +++ b/backends/samsung/builders/op_mul.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input1 = node.args[0] input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) @@ -41,3 +41,5 @@ def define_node( enn_graph.define_op( node.name, "ELTMUL", [input_id_1, input_id_2], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_permute.py b/backends/samsung/builders/op_permute.py index 646eac4c06a..42286dddfef 100644 --- a/backends/samsung/builders/op_permute.py +++ b/backends/samsung/builders/op_permute.py @@ -25,13 +25,17 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) # permutation permute_order = cast(List[int], node.args[1]) + # to prevent negative values + permute_order = [x % len(permute_order) for x in permute_order] params = {"perm": permute_order} output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "TRANSPOSE", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_pixel_shuffle.py b/backends/samsung/builders/op_pixel_shuffle.py index 28259299c81..db0aaaaef08 100644 --- a/backends/samsung/builders/op_pixel_shuffle.py +++ b/backends/samsung/builders/op_pixel_shuffle.py @@ -25,7 +25,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) scale_factor = cast(int, node.args[1]) @@ -36,3 +36,5 @@ def define_node( enn_graph.define_op( node.name, "DEPTH_TO_SPACE", [input_id], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_placeholder.py b/backends/samsung/builders/op_placeholder.py index b4b606f56ea..8c6a89a5eb5 100644 --- a/backends/samsung/builders/op_placeholder.py +++ b/backends/samsung/builders/op_placeholder.py @@ -31,7 +31,9 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: if is_param_node(self.exported_program, node): return self.define_tensor(node, enn_graph, vals_to_ids) + + return True diff --git a/backends/samsung/builders/op_pow.py b/backends/samsung/builders/op_pow.py index cd6ec7f81ef..d417685126e 100644 --- a/backends/samsung/builders/op_pow.py +++ b/backends/samsung/builders/op_pow.py @@ -4,6 +4,7 @@ # 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 from typing import Dict import torch @@ -24,15 +25,17 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input1 = node.args[0] input2 = node.args[1] input_tensor_1 = get_tensor(self.exported_program, input1) input_tensor_2 = get_tensor(self.exported_program, input2) - assert ( - input_tensor_1.dtype == torch.float32 - and input_tensor_2.dtype == torch.float32 - ), "Requires the two inputs are all float type" + if ( + input_tensor_1.dtype != torch.float32 + or input_tensor_2.dtype != torch.float32 + ): + logging.warning("Requires the two inputs are all float type.") + return False input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) input_id_2 = self.define_tensor(input2, enn_graph, vals_to_ids) @@ -40,3 +43,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "POW", [input_id_1, input_id_2], [output_id]) + + return True diff --git a/backends/samsung/builders/op_quantize.py b/backends/samsung/builders/op_quantize.py index dcf30e291f9..771ab419888 100644 --- a/backends/samsung/builders/op_quantize.py +++ b/backends/samsung/builders/op_quantize.py @@ -24,32 +24,39 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: # input input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) scales = node.args[1] - if isinstance(scales, torch.Tensor): - scales = scales.tolist() - elif not isinstance(scales, list): - scales = torch.tensor(scales).reshape([1]).tolist() zero_points = node.args[2] - if isinstance(zero_points, torch.Tensor): - zero_points = zero_points.tolist() - elif not isinstance(zero_points, list): - zero_points = torch.tensor(zero_points).reshape([1]).tolist() + if not isinstance(scales, torch.fx.Node) and not isinstance( + zero_points, torch.fx.Node + ): + if isinstance(scales, torch.Tensor): + scales = scales.tolist() + elif not isinstance(scales, list): + scales = torch.tensor(scales).reshape([1]).tolist() + if isinstance(zero_points, torch.Tensor): + zero_points = zero_points.tolist() + elif not isinstance(zero_points, list): + zero_points = torch.tensor(zero_points).reshape([1]).tolist() - output_id = self.define_tensor(node, enn_graph, vals_to_ids) + output_id = self.define_tensor(node, enn_graph, vals_to_ids) - params = {"scales": scales, "zero_points": zero_points} + params = {"scales": scales, "zero_points": zero_points} - if node.target in QuantConstants.QUANT_OPS_KEY_MAP: - enn_graph.define_op(node.name, "QUANTIZE", [input_id], [output_id], params) - else: - enn_graph.define_op( - node.name, "DEQUANTIZE", [input_id], [output_id], params - ) + if node.target in QuantConstants.QUANT_OPS_KEY_MAP: + enn_graph.define_op( + node.name, "QUANTIZE", [input_id], [output_id], params + ) + else: + enn_graph.define_op( + node.name, "DEQUANTIZE", [input_id], [output_id], params + ) + + return True @register_node_visitor diff --git a/backends/samsung/builders/op_relu.py b/backends/samsung/builders/op_relu.py index a4a2b6bc4f0..fb2668e46fc 100644 --- a/backends/samsung/builders/op_relu.py +++ b/backends/samsung/builders/op_relu.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -35,3 +35,5 @@ def define_node( self._update_params_qdtype(node, params) enn_graph.define_op(node.name, "RELU", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_reshape.py b/backends/samsung/builders/op_reshape.py index 1f4e85ac059..bb413ed793d 100644 --- a/backends/samsung/builders/op_reshape.py +++ b/backends/samsung/builders/op_reshape.py @@ -7,6 +7,7 @@ NodeVisitor, register_node_visitor, ) +from executorch.backends.samsung.builders.utils import get_tensor from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph @@ -22,13 +23,18 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) - new_shape = node.args[1] + # node.args[1] may contain "sym_size" + tensor = get_tensor(self.exported_program, node) + shape = [1] if len(tensor.size()) == 0 else list(tensor.size()) + enn_graph.define_op( - node.name, "RESHAPE", [input_id], [output_id], {"new_shape": new_shape} + node.name, "RESHAPE", [input_id], [output_id], {"new_shape": shape} ) + + return True diff --git a/backends/samsung/builders/op_rms_norm.py b/backends/samsung/builders/op_rms_norm.py index 6a58d62a5ce..0ff01701d1b 100644 --- a/backends/samsung/builders/op_rms_norm.py +++ b/backends/samsung/builders/op_rms_norm.py @@ -24,7 +24,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: # args of node : ['input', 'normalized_shape', 'weight', 'eps'] input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -50,3 +50,5 @@ def define_node( enn_graph.define_op( node.name, "RMSNORM", [input_id, gamma_id], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_rsqrt.py b/backends/samsung/builders/op_rsqrt.py index b3600d41ee2..55e9ddf54f9 100644 --- a/backends/samsung/builders/op_rsqrt.py +++ b/backends/samsung/builders/op_rsqrt.py @@ -26,10 +26,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "RSQRT", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_select.py b/backends/samsung/builders/op_select.py index 26f455b2548..3f3550d3cf8 100644 --- a/backends/samsung/builders/op_select.py +++ b/backends/samsung/builders/op_select.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ): + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -50,3 +50,5 @@ def define_node( } enn_graph.define_op(node.name, "STRIDEDSLICE", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_sigmoid.py b/backends/samsung/builders/op_sigmoid.py index e87973f9a85..aef9d90ec52 100644 --- a/backends/samsung/builders/op_sigmoid.py +++ b/backends/samsung/builders/op_sigmoid.py @@ -23,10 +23,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "SIGMOID", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_sin.py b/backends/samsung/builders/op_sin.py index 5fd22e8275e..11dc6bccb46 100644 --- a/backends/samsung/builders/op_sin.py +++ b/backends/samsung/builders/op_sin.py @@ -23,9 +23,11 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "Sin", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_skip.py b/backends/samsung/builders/op_skip.py new file mode 100644 index 00000000000..3413e603983 --- /dev/null +++ b/backends/samsung/builders/op_skip.py @@ -0,0 +1,32 @@ +# Copyright (c) 2026 Samsung Electronics Co. LTD +# 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 torch +from executorch.backends.samsung.builders.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph + + +@register_node_visitor +class OpSkipVisitor(NodeVisitor): + target = ["sym_size.int", "add", "floordiv"] + """ + do nothing + """ + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + enn_graph: EnnGraph, + vals_to_ids: Dict[torch.Tensor, int], + ) -> bool: + return True diff --git a/backends/samsung/builders/op_slice_copy.py b/backends/samsung/builders/op_slice_copy.py index e85b6bf60c3..4b837db1b12 100644 --- a/backends/samsung/builders/op_slice_copy.py +++ b/backends/samsung/builders/op_slice_copy.py @@ -27,7 +27,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ): + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -61,3 +61,5 @@ def define_node( params = {"begin": begin, "end": end, "strides": strides} enn_graph.define_op(node.name, "STRIDEDSLICE", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_softmax.py b/backends/samsung/builders/op_softmax.py index 7f569cea6fc..e96b5638db3 100644 --- a/backends/samsung/builders/op_softmax.py +++ b/backends/samsung/builders/op_softmax.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ): + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -37,3 +37,5 @@ def define_node( params = {"axis": axis} self._update_params_qdtype(node, params) enn_graph.define_op(node.name, "SOFTMAX", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_split_with_sizes_copy.py b/backends/samsung/builders/op_split_with_sizes_copy.py index b67b5331627..48612ba9a6d 100644 --- a/backends/samsung/builders/op_split_with_sizes_copy.py +++ b/backends/samsung/builders/op_split_with_sizes_copy.py @@ -23,7 +23,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -39,6 +39,10 @@ def define_node( ) all_output_tensors.append(output_id) + for user in node.users.keys(): + if user.target.__name__ == "getitem" and len(user.args) > 1: + vals_to_ids[user] = all_output_tensors[user.args[1]] + axis = node.args[2] if len(node.args) > 2 else 0 params = {} @@ -46,3 +50,5 @@ def define_node( params["point"] = node.args[1] enn_graph.define_op(node.name, "SPLIT", [input_id], all_output_tensors, params) + + return True diff --git a/backends/samsung/builders/op_sqrt.py b/backends/samsung/builders/op_sqrt.py index 3560542a0bc..77793e4589d 100644 --- a/backends/samsung/builders/op_sqrt.py +++ b/backends/samsung/builders/op_sqrt.py @@ -26,10 +26,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "SQRT", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_squeeze.py b/backends/samsung/builders/op_squeeze.py index 82fa17fbc95..ff5286fee28 100644 --- a/backends/samsung/builders/op_squeeze.py +++ b/backends/samsung/builders/op_squeeze.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -35,3 +35,5 @@ def define_node( params = {"new_shape": [*node.meta["val"].shape]} enn_graph.define_op(node.name, "RESHAPE", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_sub.py b/backends/samsung/builders/op_sub.py index 7dc97bfa7ca..c12f4c85f4d 100644 --- a/backends/samsung/builders/op_sub.py +++ b/backends/samsung/builders/op_sub.py @@ -4,6 +4,7 @@ # 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 from typing import Dict import torch @@ -26,12 +27,16 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: # inputs input1 = node.args[0] input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) input2 = node.args[1] input_id_2 = self.define_tensor(input2, enn_graph, vals_to_ids) + alpha = node.kwargs.get("alpha", 1.0) + if alpha != 1.0: + logging.warning("Currently, only alpha 1 for sub is supported.") + return False # output output_id = self.define_tensor(node, enn_graph, vals_to_ids) @@ -41,3 +46,5 @@ def define_node( enn_graph.define_op( node.name, "SUB", [input_id_1, input_id_2], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_sum_int_list.py b/backends/samsung/builders/op_sum_int_list.py index 7743e6632dd..a8c5367371c 100644 --- a/backends/samsung/builders/op_sum_int_list.py +++ b/backends/samsung/builders/op_sum_int_list.py @@ -24,7 +24,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -37,3 +37,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "REDUCESUM", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_tanh.py b/backends/samsung/builders/op_tanh.py index 5b002890075..106256c791b 100644 --- a/backends/samsung/builders/op_tanh.py +++ b/backends/samsung/builders/op_tanh.py @@ -23,10 +23,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "TANH", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_to_copy.py b/backends/samsung/builders/op_to_copy.py index c770602bb5f..143ab007cba 100644 --- a/backends/samsung/builders/op_to_copy.py +++ b/backends/samsung/builders/op_to_copy.py @@ -28,7 +28,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: memory_format_target = node.kwargs.get("memory_format", torch.contiguous_format) to_contiguous = bool(memory_format_target == torch.contiguous_format) assert to_contiguous, "Don't support other param in _to_copy" @@ -42,3 +42,5 @@ def define_node( params["out_dtype"] = get_map_dtype(out_tensor.dtype) enn_graph.define_op(node.name, "CAST", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_topk.py b/backends/samsung/builders/op_topk.py index e4cda0ef148..6a4de9ccc91 100644 --- a/backends/samsung/builders/op_topk.py +++ b/backends/samsung/builders/op_topk.py @@ -24,7 +24,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -72,3 +72,5 @@ def define_node( raise AssertionError("Not supported sorted = False.") enn_graph.define_op(node.name, "TopK", [input_id], all_output_tensors, params) + + return True diff --git a/backends/samsung/builders/op_unsqueeze.py b/backends/samsung/builders/op_unsqueeze.py index 61fa06e6310..18f7c2d33b2 100644 --- a/backends/samsung/builders/op_unsqueeze.py +++ b/backends/samsung/builders/op_unsqueeze.py @@ -25,7 +25,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -33,3 +33,5 @@ def define_node( params = {"new_shape": [*node.meta["val"].shape]} enn_graph.define_op(node.name, "RESHAPE", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_upsample_bilinear2d.py b/backends/samsung/builders/op_upsample_bilinear2d.py index d4b040460e3..7374e687101 100644 --- a/backends/samsung/builders/op_upsample_bilinear2d.py +++ b/backends/samsung/builders/op_upsample_bilinear2d.py @@ -4,6 +4,7 @@ # 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 from typing import cast, Dict, List import torch @@ -27,11 +28,14 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) in_shape = get_shape(input) output_size = cast(List[int], node.args[1]) + if output_size is None: + logging.warning("output is None for this case.") + return False scale_factor = [ output_size[0] * 1.0 / in_shape[-2], output_size[1] * 1.0 / in_shape[-1], @@ -44,10 +48,12 @@ def define_node( params = { "align_corners": align_corners, "upsampling_factor": scale_factor, - "half_pixel_centers": True, + "half_pixel_centers": not align_corners, } self._update_params_qdtype(node, params) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op( node.name, "RESIZE_BILINEAR", [input_id], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_upsample_nearest2d.py b/backends/samsung/builders/op_upsample_nearest2d.py index 9859cd8f07e..6af5402d56c 100644 --- a/backends/samsung/builders/op_upsample_nearest2d.py +++ b/backends/samsung/builders/op_upsample_nearest2d.py @@ -4,6 +4,7 @@ # 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 from typing import cast, Dict, List import torch @@ -27,11 +28,14 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) in_shape = get_shape(input) output_size = cast(List[int], node.args[1]) + if output_size is None: + logging.warning("output is None for this case.") + return False scale_factor = [ output_size[0] * 1.0 / in_shape[-2], output_size[1] * 1.0 / in_shape[-1], @@ -50,3 +54,5 @@ def define_node( enn_graph.define_op( node.name, "RESIZE_NEAREST_NEIGHBOR", [input_id], [output_id], params ) + + return True diff --git a/backends/samsung/partition/enn_partitioner.py b/backends/samsung/partition/enn_partitioner.py index 91f496e7a5c..3e450cdf1bd 100644 --- a/backends/samsung/partition/enn_partitioner.py +++ b/backends/samsung/partition/enn_partitioner.py @@ -15,6 +15,7 @@ from executorch.backends.samsung.serialization.compile_options import ( ENN_COMPILE_OPTION_TITLE, ) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph from executorch.backends.samsung.utils.utils import get_compile_spec from executorch.exir.backend.backend_details import CompileSpec from executorch.exir.backend.canonical_partitioners.pattern_op_partitioner import ( @@ -70,9 +71,16 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool: ]: return False - if node.target in SUPPORTED_OPS or node.target.__name__ in self.node_visitors: + if node.target in SUPPORTED_OPS: return True + if node.target.__name__ in self.node_visitors: + enn_graph = EnnGraph() + vals_to_ids: Dict[torch.fx.Node, int] = {} + return self.node_visitors[node.target.__name__].define_node( + node, enn_graph, vals_to_ids + ) + supported = self.enn_wrapper.IsNodeSupportedByBackend() return supported @@ -91,10 +99,19 @@ def generate_partitions( self, edge_program: torch.export.ExportedProgram ) -> List[Any]: self.op_support_checker = EnnOperatorSupport(edge_program, self.compile_specs) - return generate_partitions_from_list_of_nodes( + partition_list = generate_partitions_from_list_of_nodes( edge_program.graph_module, op_support=self.op_support_checker, ) + if len(partition_list) == 1 and partition_list[0].size() == 1: + first_node = list(partition_list[0].nodes.keys())[0] + # If there is only one partition graph containing a single "aten.clone.default" that is a useless operation, + # the RemoveUselessOpPass will remove this operation and cause a graph error. + # Therefore, we delete this node to prevent this graph error. + # For example, in the test_index_put_in_place_dtype case partition_list is [{aten_clone_default: 2}] + if first_node.target == exir_ops.edge.aten.clone.default: + del partition_list[0] + return partition_list def tag_nodes(self, partitions: List[Partition]) -> None: partition_tags: Dict[str, DelegationSpec] = {} @@ -127,8 +144,6 @@ def ops_to_not_decompose( torch.ops.aten.max_pool2d.default, torch.ops.aten.linear.default, torch.ops.aten._safe_softmax.default, - torch.ops.aten.upsample_bilinear2d.vec, - torch.ops.aten.upsample_nearest2d.vec, torch.ops.aten.prelu.default, torch.ops.aten.layer_norm.default, torch.ops.aten.pixel_shuffle.default, diff --git a/backends/samsung/runtime/CMakeLists.txt b/backends/samsung/runtime/CMakeLists.txt index deb93f31bc8..aec4a71de5d 100644 --- a/backends/samsung/runtime/CMakeLists.txt +++ b/backends/samsung/runtime/CMakeLists.txt @@ -7,7 +7,7 @@ # logging target_sources( enn_logging - PUBLIC ${CMAKE_CURRENT_LIST_DIR}/logging.h + PUBLIC $ PRIVATE ${CMAKE_CURRENT_LIST_DIR}/logging.cpp ) @@ -17,6 +17,12 @@ if(${ANDROID}) enn_backend PRIVATE ${CMAKE_CURRENT_LIST_DIR}/enn_backend.cpp ${CMAKE_CURRENT_LIST_DIR}/enn_executor.cpp + ${CMAKE_CURRENT_LIST_DIR}/enn_shared_memory_manager.cpp ${CMAKE_CURRENT_LIST_DIR}/enn_api_implementation.cpp + ${CMAKE_CURRENT_LIST_DIR}/extension/exynos_file_data_loader.cpp ) + + if(BUILD_TESTING) + add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/extension/test) + endif() endif() diff --git a/backends/samsung/runtime/enn_api_implementation.cpp b/backends/samsung/runtime/enn_api_implementation.cpp index bbef883fd10..0cdf810d5c4 100644 --- a/backends/samsung/runtime/enn_api_implementation.cpp +++ b/backends/samsung/runtime/enn_api_implementation.cpp @@ -33,31 +33,35 @@ void* loadApiFunction(void* handle, const char* name, bool optional) { return fn; } -std::mutex EnnApi::instance_mutex_; - EnnApi* EnnApi::getEnnApiInstance() { - std::lock_guard lgd(instance_mutex_); static EnnApi enn_api; - if (!enn_api.getInitialize()) { - auto status = enn_api.loadApiLib(); - if (status == Error::Ok) { - ENN_LOG_INFO("Loading ENN API library Completed.") - enn_api.initialize_ = true; - } else { - ENN_LOG_ERROR("Failed to load enn api library. %s", dlerror()); - } - } return &enn_api; } +EnnApi::EnnApi() { + auto status = loadApiLib(); + if (status != Error::Ok) { + ET_LOG(Error, "Failed to load enn api library. %s", dlerror()); + return; + } + auto ret = EnnInitialize(); + if (ret != ENN_RET_SUCCESS) { + ET_LOG(Error, "EnnInitialize failed: %d", static_cast(ret)); + unloadApiLib(); + return; + } + ET_LOG(Info, "Loading ENN API library Completed."); + initialize_ = true; +} + EnnApi::~EnnApi() { - std::lock_guard lgd(instance_mutex_); - if (getInitialize()) { + if (initialize_) { + EnnDeinitialize(); unloadApiLib(); } } -bool EnnApi::getInitialize() const { +bool EnnApi::isInitialized() const { return initialize_; } @@ -87,13 +91,18 @@ Error EnnApi::loadApiLib() { ENN_LOAD_API_FUNC(libenn_public_api_, EnnBufferCommit, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnGetBuffersInfo, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnReleaseBuffers, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnCreateBuffer, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnReleaseBuffer, this); + ENN_LOAD_API_FUNC( + libenn_public_api_, EnnGetFileDescriptorFromEnnBuffer, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnOpenModelFromFd, this); return Error::Ok; } Error EnnApi::unloadApiLib() { if (dlclose(libenn_public_api_) != 0) { - ENN_LOG_ERROR("Failed to close ENN API library. %s", dlerror()); + ET_LOG(Error, "Failed to close ENN API library. %s", dlerror()); return Error::Internal; } libenn_public_api_ = nullptr; diff --git a/backends/samsung/runtime/enn_api_implementation.h b/backends/samsung/runtime/enn_api_implementation.h index e6f8df01f7a..3e0bf1c5f2a 100644 --- a/backends/samsung/runtime/enn_api_implementation.h +++ b/backends/samsung/runtime/enn_api_implementation.h @@ -28,6 +28,7 @@ class EnnApi { ~EnnApi(); static EnnApi* getEnnApiInstance(); + bool isInitialized() const; EnnReturn (*EnnInitialize)(void); EnnReturn (*EnnSetPreferencePerfMode)(const uint32_t val); @@ -37,6 +38,7 @@ class EnnApi { const char* va, const uint32_t size, EnnModelId* model_id); + EnnReturn (*EnnOpenModelFromFd)(int _fd, EnnModelId* model_id); EnnReturn (*EnnSetFastIpc)(void); EnnReturn (*EnnUnsetFastIpc)(void); EnnReturn (*EnnExecuteModelFastIpc)( @@ -67,6 +69,13 @@ class EnnApi { NumberOfBuffersInfo* buffers_info); EnnReturn ( *EnnReleaseBuffers)(EnnBufferPtr* buffers, const int32_t numOfBuffers); + EnnReturn (*EnnCreateBuffer)( + const uint32_t req_size, + const uint32_t ion_flag, + EnnBufferPtr* out); + EnnReturn (*EnnReleaseBuffer)(EnnBufferPtr buf); + EnnReturn ( + *EnnGetFileDescriptorFromEnnBuffer)(EnnBufferPtr buffer, int32_t* fd); private: static std::mutex instance_mutex_; @@ -75,8 +84,7 @@ class EnnApi { void* libenn_public_api_ = nullptr; static std::atomic ref_count_; - EnnApi() = default; - bool getInitialize() const; + EnnApi(); Error loadApiLib(); Error unloadApiLib(); }; @@ -120,6 +128,14 @@ typedef EnnReturn (*EnnGetBuffersInfo_fn)( NumberOfBuffersInfo* buffers_info); typedef EnnReturn ( *EnnReleaseBuffers_fn)(EnnBufferPtr* buffers, const int32_t numOfBuffers); +typedef EnnReturn (*EnnCreateBuffer_fn)( + const uint32_t req_size, + const uint32_t ion_flag, + EnnBufferPtr* out); +typedef EnnReturn (*EnnReleaseBuffer_fn)(EnnBufferPtr buf); +typedef EnnReturn ( + *EnnGetFileDescriptorFromEnnBuffer_fn)(EnnBufferPtr buffer, int32_t* fd); +typedef EnnReturn (*EnnOpenModelFromFd_fn)(int _fd, EnnModelId* model_id); } // namespace enn } // namespace executor diff --git a/backends/samsung/runtime/enn_backend.cpp b/backends/samsung/runtime/enn_backend.cpp index 44838342013..81484afc44c 100644 --- a/backends/samsung/runtime/enn_backend.cpp +++ b/backends/samsung/runtime/enn_backend.cpp @@ -9,6 +9,7 @@ #include #include #include + #include #include #include diff --git a/backends/samsung/runtime/enn_executor.cpp b/backends/samsung/runtime/enn_executor.cpp index f6c1b08a8a6..99a2ce11e5e 100644 --- a/backends/samsung/runtime/enn_executor.cpp +++ b/backends/samsung/runtime/enn_executor.cpp @@ -7,12 +7,12 @@ * */ #include +#include #include #include -#include -#include -#include +#include +#include #include namespace torch { @@ -21,21 +21,39 @@ namespace enn { Error EnnExecutor::initialize(const char* binary_buf_addr, size_t buf_size) { EXYNOS_ATRACE_FUNCTION_LINE(); + auto sm_instance = executorch::backends::enn::shared_memory_manager:: + SharedMemoryManager::getInstance(); const EnnApi* enn_api_inst = EnnApi::getEnnApiInstance(); - auto ret = enn_api_inst->EnnInitialize(); ET_CHECK_OR_RETURN_ERROR( - ret == ENN_RET_SUCCESS, Internal, "Enn initialize failed."); + enn_api_inst->isInitialized(), Internal, "Enn initialize failed."); + EnnReturn ret = ENN_RET_SUCCESS; - ET_LOG(Info, "Start to open model %p, %ld", binary_buf_addr, buf_size); - ret = enn_api_inst->EnnOpenModelFromMemory( - binary_buf_addr, buf_size, &model_id_); + ET_LOG(Info, "Start to open model %p, %zu", binary_buf_addr, buf_size); + EnnBufferPtr shared_buffer = nullptr; + if (sm_instance->query(&shared_buffer, binary_buf_addr, buf_size)) { + int32_t fd; + if (shared_buffer->va == binary_buf_addr && + !enn_api_inst->EnnGetFileDescriptorFromEnnBuffer(shared_buffer, &fd)) { + ret = enn_api_inst->EnnOpenModelFromFd(fd, &model_id_); + if (ret == ENN_RET_SUCCESS) { + ET_LOG(Info, "Opened model from file descriptor, so fd is closed"); + sm_instance->free(shared_buffer->va); + } + } + } + if (!model_id_) { + ET_LOG(Info, "Open model from memory"); + ret = enn_api_inst->EnnOpenModelFromMemory( + binary_buf_addr, buf_size, &model_id_); + } ET_CHECK_OR_RETURN_ERROR( ret == ENN_RET_SUCCESS, Internal, "Failed to load Enn model from buffer %d", (int)ret); ET_LOG(Info, "Open successfully."); + NumberOfBuffersInfo buffers_info; ret = enn_api_inst->EnnAllocateAllBuffersWithSessionId( model_id_, &alloc_buffer_, &buffers_info, 0, true); diff --git a/backends/samsung/runtime/enn_executor.h b/backends/samsung/runtime/enn_executor.h index 902b420a036..68ee1f56847 100644 --- a/backends/samsung/runtime/enn_executor.h +++ b/backends/samsung/runtime/enn_executor.h @@ -40,7 +40,7 @@ class EnnExecutor { ~EnnExecutor(); private: - EnnModelId model_id_; + EnnModelId model_id_ = 0ULL; EnnBufferPtr* alloc_buffer_ = nullptr; int32_t num_of_inputs_ = 0; int32_t num_of_outputs_ = 0; diff --git a/backends/samsung/runtime/enn_shared_memory_manager.cpp b/backends/samsung/runtime/enn_shared_memory_manager.cpp new file mode 100644 index 00000000000..9babf34de9f --- /dev/null +++ b/backends/samsung/runtime/enn_shared_memory_manager.cpp @@ -0,0 +1,119 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * 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. + * + */ +#include + +#include +#include + +#include +#include +#include + +namespace executorch { +namespace backends { +namespace enn { +namespace shared_memory_manager { + +using torch::executor::enn::EnnApi; + +static std::mutex instance_mutex_; + +SharedMemoryManager* SharedMemoryManager::getInstance() { + // Touch the EnnApi singleton first so that it outlives this instance: the + // destructor below releases buffers through the ENN API. + EnnApi::getEnnApiInstance(); + static SharedMemoryManager instance; + return &instance; +} + +SharedMemoryManager::~SharedMemoryManager() { + std::lock_guard lgd(instance_mutex_); + auto enn_api_inst = EnnApi::getEnnApiInstance(); + for (auto& buffer : buffers_) { + if (enn_api_inst->EnnReleaseBuffer(buffer)) { + ET_LOG(Error, "Failed to destroy buffer: %p", buffer->va); + } + } + buffers_.clear(); +} + +void* SharedMemoryManager::alloc(const size_t size) { + if (size > std::numeric_limits::max()) { + ET_LOG( + Error, "Requested size %zu exceeds ENN buffer limit (uint32_t)", size); + return nullptr; + } + std::lock_guard lgd(instance_mutex_); + auto enn_api_inst = EnnApi::getEnnApiInstance(); + EnnBufferPtr bufferPtr; + auto ret = enn_api_inst->EnnCreateBuffer(size, 0, &bufferPtr); + if (ret) { + ET_LOG(Error, "Buffer Creation Error"); + return nullptr; + } + buffers_.emplace_back(bufferPtr); + return bufferPtr->va; +} + +bool SharedMemoryManager::query( + EnnBufferPtr* out, + const void* ptr, + const size_t size) { + std::lock_guard lgd(instance_mutex_); + auto enn_api_inst = EnnApi::getEnnApiInstance(); + for (const auto& buffer : buffers_) { + if (buffer->va <= ptr && + ptr < static_cast(buffer->va) + buffer->size) { + int32_t fd; + auto ret = enn_api_inst->EnnGetFileDescriptorFromEnnBuffer(buffer, &fd); + if (ret) { + ET_LOG( + Info, + "va: %p, size: %zu is in LUT, but failed to get FileDescriptor", + ptr, + size); + return false; + } + *out = buffer; + return true; + } + } + ET_LOG(Info, "va: %p, size: %zu is not in LUT", ptr, size); + *out = nullptr; + return false; +} + +void SharedMemoryManager::free(void* ptr) { + std::lock_guard lgd(instance_mutex_); + auto enn_api_inst = EnnApi::getEnnApiInstance(); + for (auto it = buffers_.begin(); it != buffers_.end(); ++it) { + if ((*it)->va == ptr) { + if (enn_api_inst->EnnReleaseBuffer(*it)) { + ET_LOG( + Error, + "Failed to destroy buffer: %p, keeping tracked for retry", + ptr); + return; + } + ET_LOG( + Info, + "va(%p), size(%" PRIu32 "), offset(%" PRIu32 ") is erased from LUT", + ptr, + (*it)->size, + (*it)->offset); + buffers_.erase(it); + return; + } + } +} + +} // namespace shared_memory_manager +} // namespace enn +} // namespace backends +} // namespace executorch diff --git a/backends/samsung/runtime/enn_shared_memory_manager.h b/backends/samsung/runtime/enn_shared_memory_manager.h new file mode 100644 index 00000000000..301b18f8cd5 --- /dev/null +++ b/backends/samsung/runtime/enn_shared_memory_manager.h @@ -0,0 +1,43 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * 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 + +#include + +#include +#include + +namespace executorch { +namespace backends { +namespace enn { +namespace shared_memory_manager { + +class SharedMemoryManager { + public: + static SharedMemoryManager* getInstance(); + + SharedMemoryManager() = default; + ~SharedMemoryManager(); + SharedMemoryManager(const SharedMemoryManager&) = delete; + SharedMemoryManager& operator=(const SharedMemoryManager&) = delete; + SharedMemoryManager(SharedMemoryManager&&) = delete; + SharedMemoryManager& operator=(SharedMemoryManager&&) = delete; + + void* alloc(const size_t size); + void free(void* ptr); + bool query(EnnBufferPtr* out, const void* ptr, const size_t size); + + private: + std::vector buffers_; +}; + +} // namespace shared_memory_manager +} // namespace enn +} // namespace backends +} // namespace executorch diff --git a/backends/samsung/runtime/extension/exynos_file_data_loader.cpp b/backends/samsung/runtime/extension/exynos_file_data_loader.cpp new file mode 100644 index 00000000000..e7bd252f866 --- /dev/null +++ b/backends/samsung/runtime/extension/exynos_file_data_loader.cpp @@ -0,0 +1,265 @@ +/* + * 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. + */ + +#include + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include + +// Some platforms (e.g. Xtensa) do not support pread() that we use to read the +// file at different offsets simultaneously from multiple threads not affecting +// each other. We list them below and use a workaround for them. +#if defined(__xtensa__) || defined(__hexagon__) +#define ET_HAVE_PREAD 0 +#endif // defined(__xtensa__) || defined(__hexagon__) + +#ifndef ET_HAVE_PREAD +#define ET_HAVE_PREAD 1 +#endif // !ET_HAVE_PREAD + +using executorch::backends::enn::shared_memory_manager::SharedMemoryManager; +using executorch::runtime::Error; +using executorch::runtime::FreeableBuffer; +using executorch::runtime::Result; + +namespace executorch { +namespace backends { +namespace enn { + +namespace { + +/** + * Returns true if the value is an integer power of 2. + */ +bool is_power_of_2(size_t value) { + return value > 0 && (value & ~(value - 1)) == value; +} + +/** + * FreeableBuffer::FreeFn-compatible callback. + * + * `data` is the original buffer pointer. `context` and `size` are unused. + */ +void FreeSegment(ET_UNUSED void* context, void* data, ET_UNUSED size_t size) { + SharedMemoryManager::getInstance()->free(data); +} + +} // namespace + +ExynosFileDataLoader::~ExynosFileDataLoader() { + // file_name_ can be nullptr if this instance was moved from, but freeing a + // null pointer is safe. + std::free(const_cast(file_name_)); + // fd_ can be -1 if this instance was moved from, but closing a negative fd is + // safe (though it will return an error). + if (fd_ == -1) { + return; + } + ::close(fd_); +} + +Result ExynosFileDataLoader::from( + const char* file_name, + size_t alignment) { + ET_CHECK_OR_RETURN_ERROR( + is_power_of_2(alignment), + InvalidArgument, + "Alignment %zu is not a power of 2", + alignment); + + const size_t page_size = static_cast(sysconf(_SC_PAGESIZE)); + ET_CHECK_OR_RETURN_ERROR( + alignment <= page_size, + InvalidArgument, + "Alignment %zu exceeds page size %zu; ENN shared memory buffers " + "are only guaranteed to be page-aligned", + alignment, + page_size); + + ET_CHECK_OR_RETURN_ERROR( + file_name != nullptr, InvalidArgument, "File name cannot be empty."); + + // Use open() instead of fopen() to avoid the layer of buffering that + // fopen() does. We will be reading large portions of the file in one shot, + // so buffering does not help. + int fd = ::open(file_name, O_RDONLY); + if (fd < 0) { + ET_LOG( + Error, "Failed to open %s: %s (%d)", file_name, strerror(errno), errno); + return Error::AccessFailed; + } + + // Cache the file size. + struct stat st; + int err = ::fstat(fd, &st); + if (err < 0) { + ET_LOG( + Error, + "Could not get length of %s: %s (%d)", + file_name, + ::strerror(errno), + errno); + ::close(fd); + return Error::AccessFailed; + } + size_t file_size = st.st_size; + + // Copy the filename so we can print better debug messages if reads fail. + const char* file_name_copy = ::strdup(file_name); + + if (file_name_copy == nullptr) { + ET_LOG(Error, "strdup(%s) failed", file_name); + ::close(fd); + return Error::MemoryAllocationFailed; + } + + return ExynosFileDataLoader(fd, file_size, alignment, file_name_copy); +} + +Result ExynosFileDataLoader::load( + size_t offset, + size_t size, + ET_UNUSED const DataLoader::SegmentInfo& segment_info) const { + ET_CHECK_OR_RETURN_ERROR( + // Probably had its value moved to another instance. + fd_ >= 0, + InvalidState, + "Uninitialized"); + ET_CHECK_OR_RETURN_ERROR( + size <= file_size_ && offset <= file_size_ - size, + InvalidArgument, + "File %s: offset %zu + size %zu > file_size_ %zu", + file_name_, + offset, + size, + file_size_); + + // Don't bother allocating/freeing for empty segments. + if (size == 0) { + return FreeableBuffer(nullptr, 0, /*free_fn=*/nullptr); + } + + auto* shared_memory = SharedMemoryManager::getInstance(); + void* buffer = shared_memory->alloc(size); + if (buffer == nullptr) { + ET_LOG( + Error, + "Reading from %s at offset %zu: alloc(%zu) failed", + file_name_, + offset, + size); + return Error::MemoryAllocationFailed; + } + + auto err = load_into(offset, size, segment_info, buffer); + if (err != Error::Ok) { + shared_memory->free(buffer); + return err; + } + + return FreeableBuffer(buffer, size, FreeSegment); +} + +Result ExynosFileDataLoader::size() const { + ET_CHECK_OR_RETURN_ERROR( + // Probably had its value moved to another instance. + fd_ >= 0, + InvalidState, + "Uninitialized"); + return file_size_; +} + +ET_NODISCARD Error ExynosFileDataLoader::load_into( + size_t offset, + size_t size, + ET_UNUSED const SegmentInfo& segment_info, + void* buffer) const { + ET_CHECK_OR_RETURN_ERROR( + // Probably had its value moved to another instance. + fd_ >= 0, + InvalidState, + "Uninitialized"); + ET_CHECK_OR_RETURN_ERROR( + size <= file_size_ && offset <= file_size_ - size, + InvalidArgument, + "File %s: offset %zu + size %zu > file_size_ %zu", + file_name_, + offset, + size, + file_size_); + ET_CHECK_OR_RETURN_ERROR( + buffer != nullptr, InvalidArgument, "Provided buffer cannot be null"); + + // Read the data into the aligned address. + size_t needed = size; + uint8_t* buf = reinterpret_cast(buffer); + + // Make a duplicate fd if pread() is not available and we have to seek(). + // Cannot use the standard dup() or fcntl() calls because the returned + // duplicate will share the underlying file record and affect the original fd + // when seeking on multiple threads simultaneously. + const auto dup_fd = ET_HAVE_PREAD ? fd_ : ::open(file_name_, O_RDONLY); + + while (needed > 0) { + // Reads on macOS will fail with EINVAL if size > INT32_MAX. + const auto chunk_size = std::min( + needed, static_cast(std::numeric_limits::max())); + const auto nread = +#if ET_HAVE_PREAD + ::pread(dup_fd, buf, chunk_size, offset); +#else + (::lseek(dup_fd, offset, SEEK_SET) == (off_t)-1) + ? -1 + : ::read(dup_fd, buf, chunk_size); +#endif + if (nread < 0 && errno == EINTR) { + // Interrupted by a signal; zero bytes read. + continue; + } + if (nread <= 0) { + // nread == 0 means EOF, which we shouldn't see if we were able to read + // the full amount. nread < 0 means an error occurred. + ET_LOG( + Error, + "Reading from %s: failed to read %zu bytes at offset %zu: %s", + file_name_, + size, + offset, + nread == 0 ? "EOF" : strerror(errno)); + if (!ET_HAVE_PREAD) { + ::close(dup_fd); + } + return Error::AccessFailed; + } + needed -= nread; + buf += nread; + offset += nread; + } + if (!ET_HAVE_PREAD) { + ::close(dup_fd); + } + return Error::Ok; +} + +} // namespace enn +} // namespace backends +} // namespace executorch diff --git a/backends/samsung/runtime/extension/exynos_file_data_loader.h b/backends/samsung/runtime/extension/exynos_file_data_loader.h new file mode 100644 index 00000000000..d33453e92bf --- /dev/null +++ b/backends/samsung/runtime/extension/exynos_file_data_loader.h @@ -0,0 +1,88 @@ +/* + * 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 + +#include + +#include +#include +#include + +namespace executorch { +namespace backends { +namespace enn { + +/** + * Loads a file from disk into buffers backed by ENN shared memory (dmabuf), so + * that segments handed to the ENN backend can be opened without an extra copy. + * + * Mirrors executorch::extension::FileDataLoader, which cannot be reused + * directly because it is final and always allocates from the heap. + */ +class ExynosFileDataLoader final : public executorch::runtime::DataLoader { + public: + // `alignment` must not exceed the system page size: ENN shared memory + // buffers are only guaranteed to be page-aligned, and larger requests are + // rejected rather than silently under-aligned. + static executorch::runtime::Result from( + const char* file_name, + size_t alignment = alignof(std::max_align_t)); + + ExynosFileDataLoader(ExynosFileDataLoader&& rhs) noexcept + : file_name_(rhs.file_name_), + file_size_(rhs.file_size_), + alignment_(rhs.alignment_), + fd_(rhs.fd_) { + const_cast(rhs.file_name_) = nullptr; + const_cast(rhs.file_size_) = 0; + const_cast(rhs.alignment_) = {}; + const_cast(rhs.fd_) = -1; + } + + ~ExynosFileDataLoader() override; + + ET_NODISCARD + executorch::runtime::Result load( + size_t offset, + size_t size, + const DataLoader::SegmentInfo& segment_info) const override; + + ET_NODISCARD executorch::runtime::Result size() const override; + + ET_NODISCARD executorch::runtime::Error load_into( + size_t offset, + size_t size, + ET_UNUSED const SegmentInfo& segment_info, + void* buffer) const override; + + private: + ExynosFileDataLoader( + int fd, + size_t file_size, + size_t alignment, + const char* file_name) + : file_name_(file_name), + file_size_(file_size), + alignment_{alignment}, + fd_(fd) {} + + // Not safely copyable. + ExynosFileDataLoader(const ExynosFileDataLoader&) = delete; + ExynosFileDataLoader& operator=(const ExynosFileDataLoader&) = delete; + ExynosFileDataLoader& operator=(ExynosFileDataLoader&&) = delete; + + const char* const file_name_; // Owned by the instance. + const size_t file_size_; + const std::align_val_t alignment_; + const int fd_; // Owned by the instance. +}; + +} // namespace enn +} // namespace backends +} // namespace executorch diff --git a/backends/samsung/runtime/extension/test/CMakeLists.txt b/backends/samsung/runtime/extension/test/CMakeLists.txt new file mode 100644 index 00000000000..139a2f4959e --- /dev/null +++ b/backends/samsung/runtime/extension/test/CMakeLists.txt @@ -0,0 +1,20 @@ +# Copyright (c) 2025 Samsung Electronics Co. LTD +# 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. + +# logging + +cmake_minimum_required(VERSION 3.19) + +set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../../..) + +include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) + +set(_test_srcs exynos_file_data_loader_test.cpp) + +et_cxx_test( + exynos_file_data_loader_test SOURCES ${_test_srcs} EXTRA_LIBS enn_backend + enn_logging +) diff --git a/backends/samsung/runtime/extension/test/exynos_file_data_loader_test.cpp b/backends/samsung/runtime/extension/test/exynos_file_data_loader_test.cpp new file mode 100644 index 00000000000..777c5665acc --- /dev/null +++ b/backends/samsung/runtime/extension/test/exynos_file_data_loader_test.cpp @@ -0,0 +1,266 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * 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. + * + */ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace ::testing; +using executorch::backends::enn::ExynosFileDataLoader; +using executorch::backends::enn::shared_memory_manager::SharedMemoryManager; +using executorch::extension::testing::TempFile; +using executorch::runtime::DataLoader; +using executorch::runtime::Error; +using executorch::runtime::FreeableBuffer; +using executorch::runtime::Result; +using torch::executor::enn::EnnApi; + +class ExynosFileDataLoaderTest : public ::testing::TestWithParam { + protected: + void SetUp() override { + executorch::runtime::runtime_init(); + // Constructing the singleton initializes the ENN API. + EnnApi::getEnnApiInstance(); + } + + size_t alignment() const { + return GetParam(); + } +}; + +TEST_P(ExynosFileDataLoaderTest, InBoundsLoadsSucceed) { + uint8_t data[256]; + for (int i = 0; i < sizeof(data); ++i) { + data[i] = i; + } + TempFile tf(data, sizeof(data)); + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str(), alignment()); + ASSERT_EQ(fdl.error(), Error::Ok); + + // size() should succeed and reflect the total size. + Result size = fdl->size(); + ASSERT_EQ(size.error(), Error::Ok); + EXPECT_EQ(*size, sizeof(data)); + + // Load the first bytes of the data. + { + Result fb = fdl->load( + /*offset=*/0, + /*size=*/8, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + EXPECT_ALIGNED(fb->data(), alignment()); + EXPECT_EQ(fb->size(), 8); + EXPECT_EQ( + 0, + std::memcmp( + fb->data(), + "\x00\x01\x02\x03" + "\x04\x05\x06\x07", + fb->size())); + + // Freeing should release the buffer and clear out the segment. + fb->Free(); + EXPECT_EQ(fb->size(), 0); + EXPECT_EQ(fb->data(), nullptr); + + // Safe to call multiple times. + fb->Free(); + } + + // Load the last few bytes of the data, a different size than the first time. + { + Result fb = fdl->load( + /*offset=*/sizeof(data) - 3, + /*size=*/3, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + EXPECT_ALIGNED(fb->data(), alignment()); + EXPECT_EQ(fb->size(), 3); + EXPECT_EQ(0, std::memcmp(fb->data(), "\xfd\xfe\xff", fb->size())); + } + + // Loading all of the data succeeds. + { + Result fb = fdl->load( + /*offset=*/0, + /*size=*/sizeof(data), + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + EXPECT_ALIGNED(fb->data(), alignment()); + EXPECT_EQ(fb->size(), sizeof(data)); + EXPECT_EQ(0, std::memcmp(fb->data(), data, fb->size())); + } + + // Loading zero-sized data succeeds, even at the end of the data. + { + Result fb = fdl->load( + /*offset=*/sizeof(data), + /*size=*/0, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + EXPECT_EQ(fb->size(), 0); + } +} + +TEST_P(ExynosFileDataLoaderTest, OutOfBoundsLoadFails) { + // Create a temp file; contents don't matter. + uint8_t data[256] = {}; + TempFile tf(data, sizeof(data)); + + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str(), alignment()); + ASSERT_EQ(fdl.error(), Error::Ok); + + // Loading beyond the end of the data should fail. + { + Result fb = fdl->load( + /*offset=*/0, + /*size=*/sizeof(data) + 1, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + EXPECT_NE(fb.error(), Error::Ok); + } + + // Loading zero bytes still fails if it's past the end of the data. + { + Result fb = fdl->load( + /*offset=*/sizeof(data) + 1, + /*size=*/0, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + EXPECT_NE(fb.error(), Error::Ok); + } +} + +TEST_P(ExynosFileDataLoaderTest, FromMissingFileFails) { + // Wrapping a file that doesn't exist should fail. + Result fdl = ExynosFileDataLoader::from( + "/tmp/FILE_DOES_NOT_EXIST_EXECUTORCH_EXYNOS_LOADER_TEST"); + EXPECT_NE(fdl.error(), Error::Ok); +} + +TEST_P(ExynosFileDataLoaderTest, FromEmptyFilePathFails) { + // Nullptr should fail + Result fdl = ExynosFileDataLoader::from(nullptr); + EXPECT_NE(fdl.error(), Error::Ok); +} + +TEST_P(ExynosFileDataLoaderTest, BadAlignmentFails) { + // Create a temp file; contents don't matter. + uint8_t data[256] = {}; + TempFile tf(data, sizeof(data)); + + // Creating a loader with default alignment works fine. + { + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str()); + ASSERT_EQ(fdl.error(), Error::Ok); + } + + // Bad alignments fail. + const std::vector bad_alignments = {0, 3, 5, 17}; + for (size_t bad_alignment : bad_alignments) { + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str(), bad_alignment); + ASSERT_EQ(fdl.error(), Error::InvalidArgument); + } +} + +// Tests that the move ctor works. +TEST_P(ExynosFileDataLoaderTest, MoveCtor) { + // Create a loader. + std::string contents = "FILE_CONTENTS"; + TempFile tf(contents); + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str(), alignment()); + ASSERT_EQ(fdl.error(), Error::Ok); + EXPECT_EQ(fdl->size().get(), contents.size()); + + // Move it into another instance. + ExynosFileDataLoader dl2(std::move(*fdl)); + + // Old loader should now be invalid. + EXPECT_EQ( + fdl->load( + 0, + 0, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)) + .error(), + Error::InvalidState); + EXPECT_EQ(fdl->size().error(), Error::InvalidState); + + // New loader should point to the file. + EXPECT_EQ(dl2.size().get(), contents.size()); + Result fb = dl2.load( + /*offset=*/0, + contents.size(), + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + EXPECT_ALIGNED(fb->data(), alignment()); + ASSERT_EQ(fb->size(), contents.size()); + EXPECT_EQ(0, std::memcmp(fb->data(), contents.data(), fb->size())); +} + +// TODO: Allocation failure test +TEST_P(ExynosFileDataLoaderTest, EnnQueryFailure) { + uint8_t data[256] = {}; + TempFile tf(data, sizeof(data)); + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str(), alignment()); + ASSERT_EQ(fdl.error(), Error::Ok); + + // Allocate valid buffer + Result fb = fdl->load( + 0, 8, DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + // Test query with invalid pointer + EnnBufferPtr out; + void* invalid_ptr = reinterpret_cast(0xDEADBEEF); + bool found = SharedMemoryManager::getInstance()->query(&out, invalid_ptr, 8); + EXPECT_FALSE(found); + EXPECT_EQ(out, nullptr); +} + +TEST_P(ExynosFileDataLoaderTest, EnnFreeFailure) { + uint8_t data[256] = {}; + TempFile tf(data, sizeof(data)); + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str(), alignment()); + ASSERT_EQ(fdl.error(), Error::Ok); + + // Allocate and free with invalid ENN state + { + Result fb = fdl->load( + 0, 8, DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + EnnApi::getEnnApiInstance()->EnnDeinitialize(); + fb->Free(); // Free should fail gracefully + EnnApi::getEnnApiInstance()->EnnInitialize(); + } +} + +// Run all ExynosFileDataLoaderTests multiple times, varying the return value of +// `GetParam()` based on the `testing::Values` list. The tests will interpret +// the value as "alignment". +INSTANTIATE_TEST_SUITE_P( + VariedSegments, + ExynosFileDataLoaderTest, + testing::Values( + 1, + 4, + alignof(std::max_align_t), + 2 * alignof(std::max_align_t), + 128, + 1024)); \ No newline at end of file diff --git a/backends/samsung/test/models/test_mobilebert_finetuning.py b/backends/samsung/test/models/test_mobilebert_finetuning.py index ffb60219c39..0c695f14045 100644 --- a/backends/samsung/test/models/test_mobilebert_finetuning.py +++ b/backends/samsung/test/models/test_mobilebert_finetuning.py @@ -12,39 +12,17 @@ ) from executorch.backends.samsung.test.tester import SamsungTester from executorch.backends.samsung.test.utils.utils import TestConfig - from executorch.examples.samsung.scripts.mobilebert_finetune import MobileBertFinetune -from transformers import AutoTokenizer - - -def patch_mobilebert_finetuning(): - def _monkeypatch_load_tokenizer(self): - tokenizer = AutoTokenizer.from_pretrained( - do_lower_case=True, - ) - return tokenizer - - old_func = MobileBertFinetune.load_tokenizer - MobileBertFinetune.load_tokenizer = _monkeypatch_load_tokenizer - return old_func - - -def recover_mobilebert_finetuning(old_func): - MobileBertFinetune.load_tokenizer = old_func class Test_Milestone_MobileBertFinetune(unittest.TestCase): - @classmethod - def setUpClass(cls): - cls._old_func = patch_mobilebert_finetuning(cls.model_cache_dir) - - @classmethod - def tearDownClass(cls): - recover_mobilebert_finetuning(cls._old_func) - def test_mobilebert_finetuning_fp16(self): mobilebert_finetune = MobileBertFinetune() - model, _ = mobilebert_finetune.get_finetune_mobilebert(None) + # Smaller batches, because training keeps every intermediate value until + # the backward pass is done and a test runner cannot hold them all. The + # number of passes is left alone: it is what the example script uses, and + # the range this model ends up with does not fall off with fewer of them. + model, _ = mobilebert_finetune.get_finetune_mobilebert(None, batch_size=8) example_input = mobilebert_finetune.get_example_inputs() tester = SamsungTester( model, example_input, [gen_samsung_backend_compile_spec(TestConfig.chipset)] diff --git a/backends/samsung/test/tester/samsung_tester.py b/backends/samsung/test/tester/samsung_tester.py index 258aef191d0..001b7e3bb0d 100644 --- a/backends/samsung/test/tester/samsung_tester.py +++ b/backends/samsung/test/tester/samsung_tester.py @@ -11,8 +11,12 @@ import torch from executorch.backends.samsung.partition.enn_partitioner import EnnPartitioner from executorch.backends.samsung.quantizer.quantizer import EnnQuantizer, Precision +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) from executorch.backends.samsung.test.utils import RuntimeExecutor from executorch.backends.samsung.test.utils.quant_checkers import get_checker +from executorch.backends.samsung.test.utils.utils import TestConfig from executorch.backends.samsung.utils.export_utils import get_edge_compile_config from executorch.backends.test.harness import Tester as TesterBase from executorch.backends.test.harness.stages import StageType @@ -112,6 +116,7 @@ def run( transform_passes=self.transform_passes, partitioner=self.partitioners, compile_config=self.edge_compile_config, + generate_etrecord=generate_etrecord, ) @@ -145,6 +150,8 @@ def __init__( self.original_module = module self.exported_module = module self.example_inputs = example_inputs + if compile_specs is None: + compile_specs = [gen_samsung_backend_compile_spec(TestConfig.chipset)] self.compile_specs = compile_specs def quantize( @@ -167,9 +174,12 @@ def quantize( def to_edge_transform_and_lower( self, edge_compile_config: Optional[EdgeCompileConfig] = None, + generate_etrecord: bool = False, ): to_edge_transform_and_lower_stage = ToEdgeTransformAndLower( self.compile_specs, edge_compile_config ) - return super().to_edge_transform_and_lower(to_edge_transform_and_lower_stage) + return super().to_edge_transform_and_lower( + to_edge_transform_and_lower_stage, generate_etrecord + ) diff --git a/backends/samsung/test/utils/runtime_executor.py b/backends/samsung/test/utils/runtime_executor.py index 9bc274799d7..1117ffa3b84 100644 --- a/backends/samsung/test/utils/runtime_executor.py +++ b/backends/samsung/test/utils/runtime_executor.py @@ -26,7 +26,7 @@ def get_runner_path() -> Path: cwd=os.path.dirname(os.path.realpath(__file__)), text=True, ).strip() - return Path(git_root) / "build_samsung_android/backends/samsung/enn_executor_runner" + return Path(git_root) / "build_samsung_android/examples/samsung/enn_executor_runner" class EDBTestManager: @@ -157,7 +157,7 @@ def run_on_device(self) -> Tuple[torch.Tensor]: output_tensor = ( torch.from_numpy(output_array) .view(dtype=model_outputs[idx].dtype) - .view(*model_outputs[idx].shape) + .reshape(model_outputs[idx].shape) ) result.append(output_tensor) diff --git a/backends/samsung/utils/export_utils.py b/backends/samsung/utils/export_utils.py index 22f1833bd18..86512b75dec 100644 --- a/backends/samsung/utils/export_utils.py +++ b/backends/samsung/utils/export_utils.py @@ -37,6 +37,11 @@ def get_edge_compile_config(): exir_ops.edge.aten.layer_norm.default, exir_ops.edge.aten.matmul.default, exir_ops.edge.aten.hardsigmoid.default, + exir_ops.edge.aten.round.decimals, + exir_ops.edge.aten.median.dim, + exir_ops.edge.aten.median.default, + exir_ops.edge.aten.adaptive_max_pool2d.default, + exir_ops.edge.aten.adaptive_max_pool3d.default, ], ) diff --git a/backends/test/harness/BUCK b/backends/test/harness/BUCK index d432aa6dc3d..3331a5176dd 100644 --- a/backends/test/harness/BUCK +++ b/backends/test/harness/BUCK @@ -8,6 +8,7 @@ fbcode_target(_kind = runtime.python_library, srcs = native.glob(["*.py", "stages/*.py"]), visibility = ["PUBLIC"], deps = [ + "//executorch/backends/transforms:duplicate_dynamic_quant_chain", "//executorch/exir:graph_module", ], ) diff --git a/backends/test/suite/flow.py b/backends/test/suite/flow.py index 547c79326e6..7331adc4a49 100644 --- a/backends/test/suite/flow.py +++ b/backends/test/suite/flow.py @@ -227,4 +227,17 @@ def all_flows() -> dict[str, TestFlow]: except Exception as e: logger.info(f"Skipping MLX flow registration: {e}") + try: + from executorch.backends.test.suite.flows.samsung import ( + SAMSUNG_A8W8_TEST_FLOW, + SAMSUNG_TEST_FLOW, + ) + + flows += [ + SAMSUNG_TEST_FLOW, + SAMSUNG_A8W8_TEST_FLOW, + ] + except Exception as e: + logger.info(f"Skipping SAMSUNG flow registration: {e}") + return {f.name: f for f in flows if f is not None} diff --git a/backends/test/suite/flows/samsung.py b/backends/test/suite/flows/samsung.py new file mode 100644 index 00000000000..577fde77083 --- /dev/null +++ b/backends/test/suite/flows/samsung.py @@ -0,0 +1,43 @@ +import logging + +from executorch.backends.samsung.quantizer.quantizer import EnnQuantizer, Precision +from executorch.backends.samsung.test.tester.samsung_tester import ( + Quantize, + SamsungTester, +) +from executorch.backends.test.suite.flow import TestFlow + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + + +def _create_samsung_flow( + name: str, + quantize: bool = False, + quant_dtype: Precision | None = None, + is_per_channel: bool = True, + is_qat: bool = False, +) -> TestFlow: + if quantize and quant_dtype is None: + raise RuntimeError("Quant dtype must be provided when quantize is true.") + + def create_quantize_stage() -> Quantize: + quantizer = EnnQuantizer() + quantizer.setup_quant_params(quant_dtype, is_per_channel, is_qat) + return Quantize(quantizer=quantizer) + + return TestFlow( + name, + backend="samsung", + tester_factory=SamsungTester, + quantize=quantize, + quantize_stage_factory=create_quantize_stage if quantize else None, + supports_serialize=False, + ) + + +SAMSUNG_TEST_FLOW = _create_samsung_flow("samsung") + +SAMSUNG_A8W8_TEST_FLOW = _create_samsung_flow( + "samsung_a8w8", quantize=True, quant_dtype=Precision.A8W8 +) diff --git a/backends/test/suite/models/test_torchaudio.py b/backends/test/suite/models/test_torchaudio.py index 2287b226c37..b6879162297 100644 --- a/backends/test/suite/models/test_torchaudio.py +++ b/backends/test/suite/models/test_torchaudio.py @@ -62,7 +62,7 @@ def test_conformer(test_runner, dtype: torch.dtype, use_dynamic_shapes: bool): encoder_padding_mask, ) - test_runner.lower_and_run_model(model, inputs) + test_runner.lower_and_run_model(model, inputs, generate_random_test_inputs=False) @pytest.mark.parametrize("dtype", [torch.float32], ids=dtype_to_str) diff --git a/backends/transforms/canonicalize_view_copy_permute_pass.py b/backends/transforms/canonicalize_view_copy_permute_pass.py index 0a76f10011d..c29a246b04a 100644 --- a/backends/transforms/canonicalize_view_copy_permute_pass.py +++ b/backends/transforms/canonicalize_view_copy_permute_pass.py @@ -197,45 +197,83 @@ def _fuse_sequential_ops( any_changed = True continue - if index + 1 < len(updated_chain): - next_node = updated_chain[index + 1] - if ( - node.target == self._VIEW_TARGET - and next_node.target == self._VIEW_TARGET - ): - # Fuse conscutive views - self._set_node_op( - node, self._VIEW_TARGET, input_node, self._shape(next_node) - ) - self._remove_node( - graph_module, updated_chain, index + 1, replacement=node - ) - changed = True - any_changed = True - continue - - if self._is_permute(node) and self._is_permute(next_node): - # Fuse consecutive permutes - dims = self._permute_dims(node) - next_dims = self._permute_dims(next_node) - self._set_node_op( - node, - self._PERMUTE_TARGET, - input_node, - [dims[dim] for dim in next_dims], - ) - self._remove_node( - graph_module, updated_chain, index + 1, replacement=node - ) - changed = True - any_changed = True - continue + if self._fuse_pair(graph_module, updated_chain, index): + changed = True + any_changed = True + continue index += 1 if not changed: return updated_chain, any_changed + def _fuse_pair( + self, graph_module: GraphModule, chain: list[Node], index: int + ) -> bool: + """Fuse or reorder the adjacent pair at ``index``, if possible.""" + if index + 1 >= len(chain): + return False + + node, next_node = chain[index], chain[index + 1] + input_node = cast(Node, node.args[0]) + + if self._sink_singleton_view(chain, index): + return True + + if node.target == self._VIEW_TARGET and next_node.target == self._VIEW_TARGET: + # Fuse conscutive views + self._set_node_op( + node, self._VIEW_TARGET, input_node, self._shape(next_node) + ) + self._remove_node(graph_module, chain, index + 1, replacement=node) + return True + + if self._is_permute(node) and self._is_permute(next_node): + # Fuse consecutive permutes + dims = self._permute_dims(node) + next_dims = self._permute_dims(next_node) + self._set_node_op( + node, + self._PERMUTE_TARGET, + input_node, + [dims[dim] for dim in next_dims], + ) + self._remove_node(graph_module, chain, index + 1, replacement=node) + return True + + return False + + def _sink_singleton_view(self, chain: list[Node], index: int) -> bool: + """Rewrite ``view(S).permute(P)`` to ``permute(P').view(S')``. + + Only applies to a lone pair whose view just inserts unit dimensions. + Longer chains are reordered by the swap loop in ``call()``; a pair never + is. Permuting at the lower rank lets layout boundaries that reach the + same tensor from different ranks converge on one permute. + + """ + if len(chain) != 2: + return False + + view_node, permute_node = chain[index], chain[index + 1] + input_node = cast(Node, view_node.args[0]) + if ( + view_node.target != self._VIEW_TARGET + or not self._is_permute(permute_node) + or not self._only_inserts_singletons( + self._shape(input_node), self._shape(view_node) + ) + ): + return False + + swapped_args = self._view_permute_swap(view_node, permute_node) + if swapped_args is None: + return False + + self._set_node_op(view_node, self._PERMUTE_TARGET, input_node, swapped_args[0]) + self._set_node_op(permute_node, self._VIEW_TARGET, view_node, swapped_args[1]) + return True + def _maybe_swap_args( self, op1: Node, op2: Node ) -> tuple[Sequence[_Dim], Sequence[_Dim]] | None: @@ -319,6 +357,18 @@ def _inverse_permutation(permutation: Sequence[int]) -> list[int]: inverse[dim] = index return inverse + @classmethod + def _only_inserts_singletons( + cls, input_shape: Sequence[_Dim], output_shape: Sequence[_Dim] + ) -> bool: + """Whether a view only adds singleton dimensions to its input.""" + if len(output_shape) <= len(input_shape): + return False + kept = [dim for dim in output_shape if not _dim_equals(dim, 1)] + return cls._shapes_equal( + kept, [dim for dim in input_shape if not _dim_equals(dim, 1)] + ) + @classmethod def _is_singleton_permutation( cls, shape: Sequence[_Dim], permutation: Sequence[int] diff --git a/backends/transforms/fuse_duplicate_users_pass.py b/backends/transforms/fuse_duplicate_users_pass.py index b3989e76c94..47b6064b618 100644 --- a/backends/transforms/fuse_duplicate_users_pass.py +++ b/backends/transforms/fuse_duplicate_users_pass.py @@ -11,7 +11,76 @@ from executorch.exir.pass_base import ExportPass, PassResult from torch._ops import OpOverload from torch.fx import GraphModule, Node -from torch.fx.node import Argument, map_arg +from torch.fx.node import map_arg + + +DO_NOT_FUSE_DUPLICATE_META_KEY = "do_not_fuse_duplicate" + + +def _map_leaf_to_key(node: Node) -> str: + return node.name + + +def _to_hashable(value: Any) -> Hashable: + """Convert arbitrarily nested structures into hashable tuples.""" + + if isinstance(value, (list, tuple)): + return tuple(_to_hashable(v) for v in value) + if isinstance(value, dict): + normalized_items = [(k, _to_hashable(v)) for k, v in value.items()] + return tuple(sorted(normalized_items, key=lambda item: repr(item[0]))) + if isinstance(value, set): + hashable_values: List[Hashable] = [_to_hashable(v) for v in value] + return tuple(sorted(hashable_values, key=repr)) + if isinstance(value, slice): + return ( + "slice", + _to_hashable(value.start), + _to_hashable(value.stop), + _to_hashable(value.step), + ) + if isinstance(value, range): + return ("range", value.start, value.stop, value.step) + if isinstance(value, torch.Size): + return ("size", tuple(value)) + if isinstance(value, torch.dtype): + return ("dtype", str(value)) + if isinstance(value, torch.device): + return ("device", str(value)) + if isinstance(value, torch.memory_format): + return ("memory_format", str(value)) + if isinstance(value, torch.Tensor): + return ( + "tensor", + str(value.dtype), + tuple(value.size()), + value.device.type, + value.requires_grad, + ) + return value + + +def _get_target_key(target: Any) -> Hashable: + if isinstance(target, (EdgeOpOverload, OpOverload)): + return str(target) + return target + + +def build_node_signature( + node: Node, *, positional_arg_start: int = 0 +) -> Tuple[Hashable, ...] | None: + """Build a stable signature while ignoring leading positional operands.""" + try: + normalized_args = _to_hashable( + map_arg(node.args[positional_arg_start:], _map_leaf_to_key) + ) + normalized_kwargs = _to_hashable( + {k: map_arg(v, _map_leaf_to_key) for k, v in node.kwargs.items()} + ) + except TypeError: + return None + + return (node.op, _get_target_key(node.target), normalized_args, normalized_kwargs) class FuseDuplicateUsersPass(ExportPass): @@ -106,7 +175,7 @@ def _get_candidate_groups(self, node_order, user_nodes): if user.target in self._excluded_targets: continue - target_key = self._get_target_key(user.target) + target_key = _get_target_key(user.target) target_signature = (user.op, target_key) users_by_target.setdefault(target_signature, []).append(user) @@ -120,62 +189,6 @@ def _get_candidate_groups(self, node_order, user_nodes): return candidate_groups def _build_user_signature(self, node: Node) -> Tuple[Hashable, ...] | None: - try: - normalized_args = self._to_hashable( - map_arg(node.args, self._map_leaf_to_key) - ) - normalized_kwargs = self._to_hashable( - {k: map_arg(v, self._map_leaf_to_key) for k, v in node.kwargs.items()} - ) - except TypeError: + if node.meta.get(DO_NOT_FUSE_DUPLICATE_META_KEY, False): return None - - target_key = self._get_target_key(node.target) - - return (node.op, target_key, normalized_args, normalized_kwargs) - - def _map_leaf_to_key(self, node: Node) -> Argument: - return node.name - - def _to_hashable(self, value: Any) -> Hashable: - """Convert arbitrarily nested structures into hashable tuples.""" - - if isinstance(value, (list, tuple)): - return tuple(self._to_hashable(v) for v in value) - if isinstance(value, dict): - normalized_items = [(k, self._to_hashable(v)) for k, v in value.items()] - return tuple(sorted(normalized_items, key=lambda item: repr(item[0]))) - if isinstance(value, set): - hashable_values: List[Hashable] = [self._to_hashable(v) for v in value] - return tuple(sorted(hashable_values, key=repr)) - if isinstance(value, slice): - return ( - "slice", - self._to_hashable(value.start), - self._to_hashable(value.stop), - self._to_hashable(value.step), - ) - if isinstance(value, range): - return ("range", value.start, value.stop, value.step) - if isinstance(value, torch.Size): - return ("size", tuple(value)) - if isinstance(value, torch.dtype): - return ("dtype", str(value)) - if isinstance(value, torch.device): - return ("device", str(value)) - if isinstance(value, torch.memory_format): - return ("memory_format", str(value)) - if isinstance(value, torch.Tensor): - return ( - "tensor", - str(value.dtype), - tuple(value.size()), - value.device.type, - value.requires_grad, - ) - return value - - def _get_target_key(self, target: Any) -> Hashable: - if isinstance(target, (EdgeOpOverload, OpOverload)): - return str(target) - return target + return build_node_signature(node) diff --git a/backends/transforms/merge_split_concat_chain.py b/backends/transforms/merge_split_concat_chain.py new file mode 100644 index 00000000000..b29ca68efab --- /dev/null +++ b/backends/transforms/merge_split_concat_chain.py @@ -0,0 +1,153 @@ +# 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. + +from __future__ import annotations + +import operator +from collections.abc import Callable, Sequence, Set + +import torch +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import PassResult +from torch.fx import GraphModule, Node +from torch.fx.passes.infra.pass_base import PassBase + + +_RAW_SPLIT_TARGETS: frozenset[Callable[..., object]] = frozenset( + { + torch.ops.aten.chunk.default, + torch.ops.aten.split.Tensor, + torch.ops.aten.split_with_sizes.default, + } +) +_EDGE_SPLIT_TARGETS: frozenset[Callable[..., object]] = frozenset( + {exir_ops.edge.aten.split_with_sizes_copy.default} +) + + +def _normalize_dim(dim: int, rank: int) -> int: + normalized = dim + rank if dim < 0 else dim + assert 0 <= normalized < rank + return normalized + + +def _ordered_split_inputs( + cat_node: Node, + split_targets: Set[Callable[..., object]], +) -> tuple[Node, list[Node]] | None: + cat_inputs = get_arg(cat_node, "tensors") + if not isinstance(cat_inputs, Sequence) or not cat_inputs: + return None + + getitem_nodes: list[Node] = [] + for inp in cat_inputs: + if not isinstance(inp, Node) or inp.target != operator.getitem: + return None + getitem_nodes.append(inp) + + split_node = getitem_nodes[0].args[0] + if not isinstance(split_node, Node) or split_node.target not in split_targets: + return None + + split_outputs = split_node.meta["val"] + if not isinstance(split_outputs, (tuple, list)) or len(getitem_nodes) != len( + split_outputs + ): + return None + for index, getitem_node in enumerate(getitem_nodes): + if getitem_node.args[0] != split_node or getitem_node.args[1] != index: + return None + return split_node, getitem_nodes + + +def _is_view_equivalent( + cat_node: Node, + split_node: Node, + getitem_nodes: Sequence[Node], +) -> bool: + split_input = get_arg(split_node, "input", Node) + input_val = split_input.meta["val"] + output_val = cat_node.meta["val"] + if not isinstance(input_val, torch.Tensor) or not isinstance( + output_val, torch.Tensor + ): + return False + if not input_val.is_contiguous() or input_val.numel() != output_val.numel(): + return False + + rank = input_val.ndim + split_dim = _normalize_dim(get_arg(split_node, "dim", int), rank) + cat_dim = _normalize_dim(get_arg(cat_node, "dim", int), rank) + + first_moved_dim, last_moved_dim = sorted((split_dim, cat_dim)) + # Moving split parts across axes preserves flattened storage order only + # when every crossed axis is singleton. + for getitem_node in getitem_nodes: + value = getitem_node.meta["val"] + if not isinstance(value, torch.Tensor) or any( + size != 1 for size in value.shape[first_moved_dim:last_moved_dim] + ): + return False + return True + + +class MergeSplitConcatChainPass(PassBase): + """Replace view-equivalent split/getitem/cat chains with a view. + + For example, splitting a contiguous ``[2, 1, 6, 4, 4]`` tensor into + three parts along dimension 2 and concatenating them along dimension 1 + is equivalent to a view with shape ``[2, 3, 2, 4, 4]``. + """ + + def _replace_cats( + self, + graph_module: GraphModule, + cat_target: Callable[..., object], + split_targets: Set[Callable[..., object]], + view_target: Callable[..., object], + ) -> bool: + modified = False + for cat_node in graph_module.graph.find_nodes( + op="call_function", target=cat_target + ): + match = _ordered_split_inputs(cat_node, split_targets) + if match is None: + continue + split_node, getitem_nodes = match + if not _is_view_equivalent(cat_node, split_node, getitem_nodes): + continue + + output_val = cat_node.meta["val"] + assert isinstance(output_val, torch.Tensor) + split_input = get_arg(split_node, "input", Node) + with graph_module.graph.inserting_before(cat_node): + replacement_view = graph_module.graph.call_function( + view_target, + (split_input, list(output_val.shape)), + ) + replacement_view.meta = cat_node.meta.copy() + cat_node.replace_all_uses_with(replacement_view) + modified = True + return modified + + def call(self, graph_module: GraphModule) -> PassResult: + modified = self._replace_cats( + graph_module, + torch.ops.aten.cat.default, + _RAW_SPLIT_TARGETS, + torch.ops.aten.view_copy.default, + ) + modified |= self._replace_cats( + graph_module, + exir_ops.edge.aten.cat.default, + _EDGE_SPLIT_TARGETS, + exir_ops.edge.aten.view_copy.default, + ) + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 1a3a6ada995..3f698f66562 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -61,11 +61,22 @@ class Subgraph: interleaves: dict[ torch.fx.Node, tuple[int, int, torch.fx.Node, torch.fx.Node] ] = field(default_factory=dict) + # Views whose target shapes must change when the surrounding layout + # transforms are removed. Boundary views are inside the region; sink + # views feed it from layout-invariant single-non-unit tensors. + view_shape_overrides: dict[torch.fx.Node, list[int]] = field( + default_factory=dict + ) + sink_edges_in: set[tuple[torch.fx.Node, torch.fx.Node]] = field( + default_factory=set + ) 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.amax.default, + exir_ops.edge.aten.amin.default, exir_ops.edge.aten.mul.Tensor, exir_ops.edge.aten.sub.Tensor, exir_ops.edge.aten.hardtanh.default, @@ -186,39 +197,62 @@ def _is_permutation_sink_view(self, node: torch.fx.Node) -> bool: non_unit = [d for d in shape if not (isinstance(d, int) and d == 1)] return len(non_unit) <= 1 - def _sink_users_are_layout_invariant(self, sink: torch.fx.Node) -> bool: - """Return whether dropping layout at ``sink`` is safe for its consumers.""" - frontier = [(user, sink) for user in sink.users] - visited: set[torch.fx.Node] = set() - while frontier: - node, producer = frontier.pop() - if node in visited: - continue - visited.add(node) + def _remapped_sink_shape( + self, sink: torch.fx.Node, start_permute: list[int] + ) -> list[int] | None: + """Return the sink shape in the layout before ``start_permute``. - if node.op == "output": - continue - if node.target == exir_ops.edge.aten.permute_copy.default: - # This explicit transform re-establishes the downstream layout, - # so consumers beyond it do not depend on the sink's layout. - continue - if self._is_permutation_sink_view(node): - continue + A sink input has at most one non-unit dimension, so changing the view's + shape does not change element order. The output rank must match the + region permutation so its broadcast axes can be remapped exactly. + """ + out_shape = self._concrete_shape(sink) + if out_shape is None or len(out_shape) != len(start_permute): + return None + # Reshape can relocate one contiguous non-unit run among singleton axes, + # but it cannot transpose multiple non-unit output axes. + if sum(dim != 1 for dim in out_shape) > 1: + return None + inverse = [start_permute.index(i) for i in range(len(start_permute))] + return [out_shape[index] for index in inverse] - tensor_inputs = [ - input_node - for input_node in node.all_input_nodes - if input_node.meta.get("val") is not None - ] - if any( - input_node is not producer and input_node.meta["val"].numel() != 1 - for input_node in tensor_inputs - ): - return False - if not self.is_node_permutable(node): - return False - frontier.extend((user, node) for user in node.users) - return True + def _singleton_view_boundary_shape( + self, + view: torch.fx.Node, + start_permute: list[int], + end_permute_node: torch.fx.Node, + ) -> list[int] | None: + """Compose a layout pair across a view that only inserts unit dims.""" + shapes = self._view_shapes(view) + end_dims = self.get_permutation(end_permute_node) + end_shape = self._concrete_shape(end_permute_node) + if shapes is None or end_dims is None or end_shape is None: + return None + in_shape, out_shape = shapes + if len(start_permute) != len(in_shape) or len(end_dims) != len(out_shape): + return None + + inserted = self._find_extra_ones(out_shape, in_shape) + if inserted is None: + return None + + # Label each old view-output axis by the corresponding axis before the + # incoming permutation. The outgoing permutation must restore those + # labels to identity order; inserted singleton axes carry no label. + labels: list[int | None] = list(start_permute) + for index in inserted: + labels.insert(index, None) + output_labels = [labels[index] for index in end_dims] + if [label for label in output_labels if label is not None] != list( + range(len(in_shape)) + ): + return None + + inverse = [start_permute.index(i) for i in range(len(start_permute))] + unpermuted_input = [in_shape[index] for index in inverse] + if self._find_extra_ones(end_shape, unpermuted_input) is None: + return None + return end_shape def _inserted_unit_dim(self, node: torch.fx.Node) -> int | None: """Position of the size-1 dim ``node`` inserts, else None. @@ -484,7 +518,19 @@ def visit( # noqa: C901 for user in users_source.users: if user.target in PERMUTE_COPY_TARGETS: user_perm = self.get_permutation(user) - if user_perm == downstream_end: + boundary_shape = None + if ( + triple is None + and self._is_squeeze_unsqueeze_view(node) + and len(node.users) == 1 + ): + boundary_shape = self._singleton_view_boundary_shape( + node, current_start_permute, user + ) + if boundary_shape is not None: + subgraph.view_shape_overrides[node] = boundary_shape + subgraph.edges_out.add((users_source, user)) + elif user_perm == downstream_end: subgraph.edges_out.add((users_source, user)) else: # Non-matching permute: keep it and fold the start permute into it @@ -501,10 +547,11 @@ def visit( # noqa: C901 elif user.op == "output": return False elif self._is_permutation_sink_view(user): - # The tensor's element order is invariant at this reshape, but - # its output shape can still carry broadcast-axis meaning. - if not self._sink_users_are_layout_invariant(user): - return False + # A sink with no other path into this region can terminate it: + # its single non-unit run has layout-invariant element order. + # If a later consumer is reached through another region branch, + # upstream traversal records and remaps the sink via + # ``sink_edges_in`` before any boundary is removed. continue elif not self.visit( user, subgraph, processed_nodes, downstream_end, downstream_start @@ -524,6 +571,15 @@ def visit( # noqa: C901 # stays wired directly. Notably this keeps lifted per-tensor # qparam placeholders as placeholders, which lowering requires. continue + elif inp in subgraph.nodes: + # Already part of the region; it is rewritten as a region node. + continue + elif self._is_permutation_sink_view(inp): + remapped_shape = self._remapped_sink_shape(inp, current_start_permute) + if remapped_shape is None or len(inp.users) != 1: + return False + subgraph.view_shape_overrides[inp] = remapped_shape + subgraph.sink_edges_in.add((inp, node)) elif self._is_constant(inp): const_rank = self._get_node_rank(inp) permute_rank = len(current_end_permute) @@ -603,6 +659,8 @@ def is_node_permutable(self, node: torch.fx.Node) -> bool: return False if node.target in self._permutable_ops: if node.target in ( + exir_ops.edge.aten.amax.default, + exir_ops.edge.aten.amin.default, exir_ops.edge.aten.mean.dim, exir_ops.edge.aten.sum.dim_IntList, ): @@ -693,16 +751,24 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 if node.target == exir_ops.edge.aten.cat.default: self.update_cat(node, node_start_perm) elif node.target in ( + exir_ops.edge.aten.amax.default, + exir_ops.edge.aten.amin.default, exir_ops.edge.aten.mean.dim, exir_ops.edge.aten.sum.dim_IntList, ): - self.update_mean_dim(node, node_start_perm) + self.update_reduction_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._PAD_OPS: self.update_pad(node, node_start_perm) elif node.target in self._VIEW_OPS: - self.update_view_copy(node, node_start_perm) + if node in subgraph.view_shape_overrides: + node.update_arg(1, subgraph.view_shape_overrides[node]) + else: + self.update_view_copy(node, node_start_perm) + + for sink, _ in subgraph.sink_edges_in: + sink.update_arg(1, subgraph.view_shape_overrides[sink]) for head, triple in subgraph.interleaves.items(): self.update_interleave( @@ -769,20 +835,19 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: """Return false if an earlier rewrite invalidated this candidate.""" for inp, out in subgraph.edges_in: - if inp.target not in PERMUTE_COPY_TARGETS or inp not in out.all_input_nodes: - return False - - # edges_out_to_update can rewrite a permute in place, leaving it wired. - if self.get_permutation(inp) != subgraph.node_start_permute.get( - out, subgraph.start_permute + if ( + inp.target not in PERMUTE_COPY_TARGETS + or inp not in out.all_input_nodes + # edges_out_to_update can rewrite a permute in place, leaving it wired. + or self.get_permutation(inp) + != subgraph.node_start_permute.get(out, subgraph.start_permute) ): return False - for inp, out in subgraph.edges_out: - if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users: - return False - - for inp, out, _ in subgraph.edges_out_to_update: + outgoing = list(subgraph.edges_out) + [ + (inp, out) for inp, out, _ in subgraph.edges_out_to_update + ] + for inp, out in outgoing: if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users: return False @@ -790,6 +855,10 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: if const_node not in user_node.all_input_nodes: return False + for sink, user_node in subgraph.sink_edges_in: + if sink not in user_node.all_input_nodes or len(sink.users) != 1: + return False + for head, (_, _, expand_node, view_node) in subgraph.interleaves.items(): if ( len(head.users) != 1 @@ -839,9 +908,19 @@ def update_cat(self, node: torch.fx.Node, start_permute: list[int]) -> None: 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: + def update_reduction_dim( + self, node: torch.fx.Node, start_permute: list[int] + ) -> None: dims = get_arg(node, "dim") - set_arg(node, "dim", [start_permute[d] for d in cast(list[int], dims)]) + rank = len(start_permute) + if isinstance(dims, int): + set_arg(node, "dim", start_permute[dims % rank]) + else: + set_arg( + node, + "dim", + [start_permute[d % rank] for d in cast(list[int], dims)], + ) def update_slice_copy(self, node: torch.fx.Node, start_permute: list[int]) -> None: dim = get_arg(node, "dim", int) diff --git a/backends/transforms/replace_channels_last_input_clones.py b/backends/transforms/replace_channels_last_input_clones.py new file mode 100644 index 00000000000..6d620de8ca4 --- /dev/null +++ b/backends/transforms/replace_channels_last_input_clones.py @@ -0,0 +1,133 @@ +# Copyright 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. + +from typing import Sequence + +import executorch.backends.transforms.channels_last_ops # noqa: F401 + +import torch + +from executorch.exir.dialects._ops import ops as exir_ops + +from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue +from torch.fx.passes.infra.pass_base import PassResult + +_DIM_ORDER_CHANGING_OPS: frozenset = frozenset( + { + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, + exir_ops.edge.dim_order_ops._clone_dim_order.default, + } +) + +# kwargs that carry no tensor semantics and require no value inspection. +_PASS_THROUGH_KWARGS: frozenset = frozenset({"dim_order", "non_blocking"}) + +# TensorOptions that the exporter may annotate explicitly on `_to_dim_order_copy` nodes even +# when no actual change occurs. These are allowed only when their value is identical to the +# source tensor's property (i.e. truly a no-op). Any value that would actually change the +# tensor (e.g. a different dtype) blocks replacement. Any kwarg not in either set is unknown +# and is also rejected. +_TENSOR_OPTION_KWARGS: frozenset = frozenset( + {"dtype", "layout", "device", "pin_memory"} +) + +_ALLOWED_KWARGS: frozenset = _PASS_THROUGH_KWARGS | _TENSOR_OPTION_KWARGS + +_TO_NHWC_PERMUTATION: list[int] = [0, 2, 3, 1] +_TO_NCHW_PERMUTATION: list[int] = [0, 3, 1, 2] + + +def _is_4d_contiguous(dim_order: Sequence[int]) -> bool: + return list(dim_order) == [0, 1, 2, 3] + + +def _is_4d_channels_last(dim_order: Sequence[int]) -> bool: + return list(dim_order) == [0, 2, 3, 1] + + +def _is_replaceable_input_boundary_clone(op, args, kwargs) -> bool: + """Return True if the op/args/kwargs describe a 4D channels-last-to-contiguous clone of a model input. + + These are the input boundary clones inserted by `EnforceContiguousDimOrder`: they consume a + channels-last placeholder and produce a contiguous tensor. Replacing them with a permute pair + allows them to be optimized out by other passes, leaving only the no-op `aten.permute_copy`. + """ + if op not in _DIM_ORDER_CHANGING_OPS: + return False + if not args or not hasattr(args[0], "node"): + return False + src = args[0].node + if not isinstance(src, torch.fx.Node) or src.op != "placeholder": + return False + val = src.meta.get("val") + if not isinstance(val, torch.Tensor): + return False + # Primary guard: reject any kwarg outside the known set. Unknown kwargs may carry + # semantics we cannot reason about, so we conservatively block replacement. + if not set(kwargs.keys()) <= _ALLOWED_KWARGS: + return False + # Secondary guard: each TensorOption that is present must be a no-op relative to the + # source tensor. The exporter annotates these even when no actual change occurs, but a + # differing value (e.g. a dtype cast) must not be silently dropped by the replacement. + if kwargs.get("dtype") not in (None, val.dtype): + return False + if kwargs.get("layout") not in (None, val.layout): + return False + if kwargs.get("device") not in (None, val.device): + return False + if kwargs.get("pin_memory") not in (None, False): + return False + return _is_4d_channels_last(val.dim_order()) and _is_4d_contiguous( + kwargs.get("dim_order", []) + ) + + +class ReplaceChannelsLastInputClones(ExportPass): + """Replace `_to_dim_order_copy` and `_clone_dim_order` with an equivalent sequence in the following pattern. This + approach allows the `channels_last.permute_copy` to be optimized out if there are subsequent channels last + operators in the model, leaving only the `aten.permute_copy`, which is effectively a no-op. As a result, the + input data doesn't have to be permuted in memory. + + │ [N, C, H, W] shape, (0, 2, 3, 1) dim order + │ data is stored channels last + │ [N, C, H, W] shape, (0, 2, 3, 1) dim order ┌─────────▼─────────┐ + │ data is stored channels last │ aten.permute_copy ◄──── [0, 2, 3, 1] permutation + ┌───────────▼────────────┐ └─────────┬─────────┘ + │ │ ────────────────► │ [N, H, W, C] shape, (0, 1, 2, 3) dim order + └───────────┬────────────┘ │ data is stored channels last + │ [N, C, H, W] shape, (0, 1, 2, 3) dim order ┌──────────────▼─────────────┐ + ▼ data is stored channels first │ channels_last.permute_copy ◄──── [0, 3, 1, 2] permutation + └──────────────┬─────────────┘ + │ [N, C, H, W] shape, (0, 1, 2, 3) dim order + ▼ data is stored channels first + """ + + _modified: bool + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + self._modified = False + result = super().call(graph_module) + return PassResult(result.graph_module, self._modified) + + def call_operator(self, op, args, kwargs, meta: NodeMetadata) -> ProxyValue: + if not _is_replaceable_input_boundary_clone(op, args, kwargs): + return super().call_operator(op, args, kwargs, meta) + + x = super().call_operator( + exir_ops.edge.aten.permute_copy.default, + (args[0], _TO_NHWC_PERMUTATION), + {}, + meta, + ) + x = super().call_operator( + exir_ops.edge.channels_last.permute_copy.default, + (x, _TO_NCHW_PERMUTATION), + {}, + meta, + ) + + self._modified = True + + return x diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index 1e456249d3f..eb282ce23e4 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -521,6 +521,30 @@ def define_common_targets(): ], ) + runtime.python_library( + name = "merge_split_concat_chain", + srcs = ["merge_split_concat_chain.py"], + visibility = ["PUBLIC"], + deps = [ + ":permute_pass_utils", + "//caffe2:torch", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_merge_split_concat_chain", + srcs = ["test/test_merge_split_concat_chain.py"], + deps = [ + ":merge_split_concat_chain", + "//caffe2:torch", + "//executorch/backends/test:graph_builder", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + runtime.python_library( name = "fuse_cascaded_transpose_or_permute_ops", srcs = ["fuse_cascaded_transpose_or_permute_ops.py"], @@ -716,3 +740,32 @@ def define_common_targets(): ":enforce_contiguous_dim_order", ], ) + + runtime.python_library( + name = "replace_channels_last_input_clones", + srcs = ["replace_channels_last_input_clones.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + ":channels_last_ops", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_replace_channels_last_input_clones", + srcs = [ + "test/test_replace_channels_last_input_clones.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir:lib", + ":channels_last_ops", + ":enforce_contiguous_dim_order", + ":replace_channels_last_input_clones", + "fbsource//third-party/pypi/pytest:pytest", + ], + ) diff --git a/backends/transforms/test/test_merge_split_concat_chain.py b/backends/transforms/test/test_merge_split_concat_chain.py new file mode 100644 index 00000000000..153bed74fa6 --- /dev/null +++ b/backends/transforms/test/test_merge_split_concat_chain.py @@ -0,0 +1,197 @@ +# 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. + +from __future__ import annotations + +import copy +import operator +import unittest +from collections.abc import Callable, Sequence +from typing import cast + +import torch +from executorch.backends.test.graph_builder import GraphBuilder +from executorch.backends.transforms.merge_split_concat_chain import ( + MergeSplitConcatChainPass, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import PassResult +from torch.fx import GraphModule + + +def _build_split_cat_graph( + input_value: torch.Tensor, + split_target: Callable[..., object], + split_spec: int | list[int], + split_dim: int, + output_count: int, + cat_target: Callable[..., object], + cat_dim: int, + order: Sequence[int] | None = None, +) -> GraphModule: + builder = GraphBuilder() + input_node = builder.placeholder("input", input_value) + split = builder.call_operator( + split_target, + (input_node, split_spec, split_dim), + ) + output_order = range(output_count) if order is None else order + split_outputs = [ + builder.call_operator(operator.getitem, (split, index)) + for index in output_order + ] + cat = builder.call_operator(cat_target, (split_outputs, cat_dim)) + builder.output([cat]) + return builder.get_graph_module() + + +_VIEW_EQUIVALENT_CASES = ( + ( + "split", + torch.ops.aten.split.Tensor, + 2, + (2, 1, 6, 4, 4), + 2, + 3, + torch.ops.aten.cat.default, + 1, + torch.ops.aten.view_copy.default, + ), + ( + "chunk", + torch.ops.aten.chunk.default, + 3, + (2, 1, 5, 4, 4), + 2, + 3, + torch.ops.aten.cat.default, + 2, + torch.ops.aten.view_copy.default, + ), + ( + "split_with_sizes", + torch.ops.aten.split_with_sizes.default, + [1, 1, 1, 1], + (1, 4, 3, 2), + 1, + 4, + torch.ops.aten.cat.default, + 0, + torch.ops.aten.view_copy.default, + ), + ( + "edge_split_with_sizes", + exir_ops.edge.aten.split_with_sizes_copy.default, + [1, 1, 1, 1], + (4, 1, 3, 2), + 0, + 4, + exir_ops.edge.aten.cat.default, + 1, + exir_ops.edge.aten.view_copy.default, + ), +) + +_UNSAFE_CASES = ( + ("reordered_outputs", [1, 0, 2, 3], 0, False), + ("non_singleton_crossing", [0, 1, 2, 3], 3, False), + ("non_contiguous_input", [0, 1, 2, 3], 0, True), +) + +_SPLIT_CAT_DIALECTS = ( + ( + "aten", + torch.ops.aten.split_with_sizes.default, + torch.ops.aten.cat.default, + ), + ( + "edge", + exir_ops.edge.aten.split_with_sizes_copy.default, + exir_ops.edge.aten.cat.default, + ), +) + + +class MergeSplitConcatChainPassTest(unittest.TestCase): + def test_merges_view_equivalent_chains(self) -> None: + for ( + name, + split_target, + split_spec, + input_shape, + split_dim, + output_count, + cat_target, + cat_dim, + view_target, + ) in _VIEW_EQUIVALENT_CASES: + with self.subTest(name=name): + input_value = torch.randn(input_shape) + graph_module = _build_split_cat_graph( + input_value, + split_target, + split_spec, + split_dim, + output_count, + cat_target, + cat_dim, + ) + reference = copy.deepcopy(graph_module) + + result = cast(PassResult, MergeSplitConcatChainPass()(graph_module)) + + self.assertTrue(result.modified) + torch.testing.assert_close( + reference(input_value), result.graph_module(input_value) + ) + for target, expected_count in ( + (split_target, 0), + (cat_target, 0), + (view_target, 1), + ): + self.assertEqual( + expected_count, + len( + result.graph_module.graph.find_nodes( + op="call_function", target=target + ) + ), + ) + + def test_does_not_merge_unsafe_chains(self) -> None: + for dialect, split_target, cat_target in _SPLIT_CAT_DIALECTS: + for name, order, cat_dim, non_contiguous in _UNSAFE_CASES: + with self.subTest(dialect=dialect, name=name): + input_value = torch.randn(1, 4, 3, 2) + if non_contiguous: + input_value = input_value.transpose(2, 3) + graph_module = _build_split_cat_graph( + input_value, + split_target, + [1, 1, 1, 1], + 1, + 4, + cat_target, + cat_dim, + order, + ) + reference = copy.deepcopy(graph_module) + + result = cast(PassResult, MergeSplitConcatChainPass()(graph_module)) + + self.assertFalse(result.modified) + torch.testing.assert_close( + reference(input_value), result.graph_module(input_value) + ) + self.assertEqual( + 1, + len( + result.graph_module.graph.find_nodes( + op="call_function", + target=cat_target, + ) + ), + ) diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index 1f357deb171..da47fe8186e 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -1600,7 +1600,7 @@ def test_permutation_sink_view_splitting_the_non_unit_dim(self) -> None: "permutation_sink_view_splitting_the_non_unit_dim", ) - def test_permutation_sink_view_preserves_broadcast_layout(self) -> None: + def test_permutation_sink_view_terminal_broadcast_is_optimized(self) -> None: x_data = torch.randn(1, 4, 1, 1) direct_data = torch.randn(1, 8, 4) builder = GraphBuilder() @@ -1623,15 +1623,87 @@ def test_permutation_sink_view_preserves_broadcast_layout(self) -> None: gm_before = copy.deepcopy(original) result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) - self.assertFalse(result.modified) + self.assertTrue(result.modified) self.assertEqual( - count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 1 + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 ) validate_numerics( gm_before, result.graph_module, [x_data, direct_data], - "permutation_sink_view_preserves_broadcast_layout", + "permutation_sink_view_terminal_broadcast_is_optimized", + ) + + def test_split_sink_view_is_not_remapped_for_broadcast(self) -> None: + x_data = torch.randn(1, 4, 2) + sink_data = torch.randn(1, 1, 1, 8) + builder = GraphBuilder() + x = builder.placeholder("x", x_data) + sink_source = builder.placeholder("sink_source", sink_data) + permute_in = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + split_sink = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, + args=(sink_source, [1, 2, 4]), + ) + mul = builder.call_operator( + op=exir_ops.edge.aten.mul.Tensor, args=(permute_in, split_sink) + ) + permute_out = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(mul, [0, 2, 1]) + ) + builder.output([permute_out]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) + 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, sink_data], + "split_sink_view_is_not_remapped_for_broadcast", + ) + + def test_shared_singleton_view_boundaries_are_not_remapped(self) -> None: + x_data = torch.randn(1, 4, 8) + builder = GraphBuilder() + x = builder.placeholder("x", x_data) + permute_in = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + view = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, + args=(permute_in, [1, 8, 4, 1]), + ) + first_out = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(view, [0, 2, 3, 1]), + ) + second_out = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(view, [0, 3, 1, 2]), + ) + builder.output([first_out, second_out]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) + # The view feeds two differing outgoing permutes, so its shape cannot be + # composed across the boundary. Both outgoing permutes survive; only the + # incoming permute is folded into them. + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 2 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "shared_singleton_view_boundaries_are_not_remapped", ) def test_permutation_sink_view_preserves_cat_layout(self) -> None: @@ -1657,10 +1729,14 @@ def test_permutation_sink_view_preserves_cat_layout(self) -> None: gm_before = copy.deepcopy(original) result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) - self.assertFalse(result.modified) + self.assertTrue(result.modified) self.assertEqual( - count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 1 + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + (view_after,) = result.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.view_copy.default ) + self.assertEqual(view_after.args[1], [1, 1, 4]) validate_numerics( gm_before, result.graph_module, @@ -1693,10 +1769,14 @@ def test_permutation_sink_view_preserves_keyword_broadcast_layout(self) -> None: gm_before = copy.deepcopy(original) result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) - self.assertFalse(result.modified) + self.assertTrue(result.modified) self.assertEqual( - count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 1 + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + (view_after,) = result.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.view_copy.default ) + self.assertEqual(view_after.args[1], [1, 1, 4]) validate_numerics( gm_before, result.graph_module, diff --git a/backends/transforms/test/test_replace_channels_last_input_clones.py b/backends/transforms/test/test_replace_channels_last_input_clones.py new file mode 100644 index 00000000000..e6b30df204f --- /dev/null +++ b/backends/transforms/test/test_replace_channels_last_input_clones.py @@ -0,0 +1,402 @@ +# Copyright 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 executorch.backends.transforms.channels_last_ops # noqa: F401 +import pytest +import torch + +from executorch.backends.transforms.enforce_contiguous_dim_order import ( + EnforceContiguousDimOrder, +) +from executorch.backends.transforms.replace_channels_last_input_clones import ( + _is_4d_channels_last, + _is_4d_contiguous, + ReplaceChannelsLastInputClones, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import ExportedProgram +from torch.fx import GraphModule +from torch.fx.node import Target + +_CLONE_DIM_ORDER = exir_ops.edge.dim_order_ops._clone_dim_order.default +_TO_DIM_ORDER_COPY = exir_ops.edge.dim_order_ops._to_dim_order_copy.default +_ATEN_PERMUTE_COPY = exir_ops.edge.aten.permute_copy.default +_CHANNELS_LAST_PERMUTE_COPY = exir_ops.edge.channels_last.permute_copy.default + + +class SingleInputToDimOrderCopyModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.avg_pool = torch.nn.AvgPool2d(3) + + def forward(self, x): + # Expecting `x` to use the channels last memory format. + x = x.to(memory_format=torch.contiguous_format) + x = self.avg_pool(x) + return x + + +class MultiInputToDimOrderCopyModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.avg_pool = torch.nn.AvgPool2d(3) + + def forward(self, *inputs): + contiguous_inputs = [ + input_.to(memory_format=torch.contiguous_format) for input_ in inputs + ] + x = torch.concatenate(contiguous_inputs) + x = self.avg_pool(x) + return x + + +class ToDimOrderCopyAfterAddModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.avg_pool = torch.nn.AvgPool2d(3) + + def forward(self, x): + x = ( + x + x + ) # Make sure the `_to_dim_order_copy` is not consuming the model input. + x = x.to(memory_format=torch.contiguous_format) + x = self.avg_pool(x) + return x + + +class SingleInputModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.avg_pool = torch.nn.AvgPool2d(3) + + def forward(self, x): + x = self.avg_pool(x) + x = torch.relu(x) + return x + + +class MultiInputModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.single_input_module = SingleInputModule() + + def forward(self, *inputs): + x = torch.concatenate(inputs) + x = self.single_input_module(x) + return x + + +class IncompatibleDimOrderModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.avg_pool = torch.nn.AvgPool2d(3) + + def forward(self, x): + x = x.to(memory_format=torch.channels_last) # Incompatible dim order. + x = self.avg_pool(x) + return x + + +class DtypeCastAndLayoutChangeModule(torch.nn.Module): + """Module whose forward casts dtype AND changes memory format via `_to_dim_order_copy`. + The pass must not replace such a node, since doing so would silently drop the cast. + """ + + def __init__(self): + super().__init__() + self.avg_pool = torch.nn.AvgPool2d(3) + + def forward(self, x): + # Both a memory-format change and a dtype cast happen here. + x = x.to(dtype=torch.float64, memory_format=torch.contiguous_format) + x = x.to(dtype=torch.float32) # Cast back so avg_pool accepts it. + x = self.avg_pool(x) + return x + + +def _export_to_edge(module: torch.nn.Module, inputs: tuple) -> ExportedProgram: + ep = torch.export.export(module.eval(), inputs) + return to_edge(ep).exported_program() + + +def _find_nodes(gm: GraphModule, target: Target) -> list[torch.fx.Node]: + return [n for n in gm.graph.nodes if n.op == "call_function" and n.target == target] + + +def _count(gm: GraphModule, target: Target) -> int: + return len(_find_nodes(gm, target)) + + +def _run_pass(ep_or_result) -> tuple[GraphModule, bool]: + gm = ep_or_result.graph_module + result = ReplaceChannelsLastInputClones()(gm) + return result.graph_module, result.modified + + +def _assert_expected_result_pattern( + input_: torch.fx.Node, + aten_permute: torch.fx.Node, + channels_last_permute: torch.fx.Node, +): + assert input_.op == "placeholder" + assert input_.meta["val"].dim_order() == (0, 2, 3, 1) + assert aten_permute.args[0] == input_ + assert aten_permute.target == _ATEN_PERMUTE_COPY + assert aten_permute.args[1] == [0, 2, 3, 1] + assert aten_permute.meta["val"].dim_order() == (0, 1, 2, 3) + assert channels_last_permute.args[0] == aten_permute + assert channels_last_permute.target == _CHANNELS_LAST_PERMUTE_COPY + assert channels_last_permute.args[1] == [0, 3, 1, 2] + assert channels_last_permute.meta["val"].dim_order() == (0, 1, 2, 3) + + +@pytest.fixture(autouse=True) +def _reseed(): + torch.manual_seed(42) + yield + + +class TestReplaceChannelsLastInputDimOrderCopies: + """These tests use models with an explicit dim order change in their `forward()` method, which results in a + `_to_dim_order_copy` operator in edge dialect. + """ + + def test_single_input(self): + example_inputs = ( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last), + ) + + ep = _export_to_edge(SingleInputToDimOrderCopyModule(), example_inputs) + assert _count(ep.graph_module, _TO_DIM_ORDER_COPY) == 1 + output_before = ep.module()(*example_inputs) + + gm, modified = _run_pass(ep) + + assert modified + assert _count(gm, _TO_DIM_ORDER_COPY) == 0 + assert _count(gm, _CLONE_DIM_ORDER) == 0 + + nodes = list(gm.graph.nodes) + _assert_expected_result_pattern(*nodes[:3]) + + outputs_after = gm(*example_inputs)[0] + assert torch.allclose(output_before, outputs_after) + + def test_multiple_inputs(self): + num_inputs = 3 + example_inputs = tuple( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last) + for _ in range(num_inputs) + ) + + ep = _export_to_edge(MultiInputToDimOrderCopyModule(), example_inputs) + assert _count(ep.graph_module, _TO_DIM_ORDER_COPY) == num_inputs + output_before = ep.module()(*example_inputs) + + gm, modified = _run_pass(ep) + + assert modified + assert _count(gm, _TO_DIM_ORDER_COPY) == 0 + assert _count(gm, _CLONE_DIM_ORDER) == 0 + + nodes = list(gm.graph.nodes) + for i in range(num_inputs): + # Each input should have the expected pattern + input_ = nodes[i] + start_idx = num_inputs + i * (num_inputs - 1) + end_idx = start_idx + 2 + pattern = nodes[start_idx:end_idx] + _assert_expected_result_pattern(input_, *pattern) + + outputs_after = gm(*example_inputs)[0] + assert torch.allclose(output_before, outputs_after) + + def test__not_applied__not_consuming_model_input(self): + example_inputs = ( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last), + ) + + ep = _export_to_edge(ToDimOrderCopyAfterAddModule(), example_inputs) + assert _count(ep.graph_module, _TO_DIM_ORDER_COPY) == 1 + + gm, modified = _run_pass(ep) + + assert not modified + assert _count(gm, _TO_DIM_ORDER_COPY) == 1 + + def test__not_applied__incompatible_dim_order(self): + # This test uses a `to_dim_order_copy` which goes from contiguous to channels_last, which is not what the pass + # was made for. + example_inputs = (torch.randn(1, 3, 8, 8),) + + ep = _export_to_edge(IncompatibleDimOrderModule(), example_inputs) + assert _count(ep.graph_module, _TO_DIM_ORDER_COPY) == 1 + + gm, modified = _run_pass(ep) + + assert not modified + assert _count(gm, _TO_DIM_ORDER_COPY) == 1 + + def test__not_applied__dtype_cast(self): + """A `_to_dim_order_copy` that also changes dtype must not be replaced. + Replacing it with permutes would silently drop the cast, changing the graph's + semantics. + """ + example_inputs = ( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last), + ) + + ep = _export_to_edge(DtypeCastAndLayoutChangeModule(), example_inputs) + # The first `_to_dim_order_copy` carries a dtype change; the pass must leave it. + to_dim_order_copies_before = _count(ep.graph_module, _TO_DIM_ORDER_COPY) + + gm, modified = _run_pass(ep) + + assert not modified + assert _count(gm, _TO_DIM_ORDER_COPY) == to_dim_order_copies_before + + @pytest.mark.parametrize( + "extra_kwarg", + [ + # dtype differs from the source tensor's dtype -> cast must not be silently dropped. + {"dtype": torch.float64}, + # layout, device, and pin_memory are in _ALLOWED_KWARGS and are allowed through + # when their value matches the source tensor (which is always the case on CPU + # hardware). The whitelist guards against UNKNOWN kwargs; value checks guard + # against CHANGED TensorOptions. These cases verify that a differing dtype is + # caught; layout/device/pin_memory rejection can only be tested by providing a + # value that actually differs from the source tensor (not possible on CPU-only). + ], + ) + def test__not_applied__dtype_cast_kwarg(self, extra_kwarg): + # Verify that a `_to_dim_order_copy` carrying a dtype change is not replaced. + # The guard must not drop the cast silently, so replacement is blocked. + g = torch.fx.Graph() + ph = g.placeholder("x") + ph.meta["val"] = torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last) + kwargs = {"dim_order": [0, 1, 2, 3], **extra_kwarg} + clone = g.call_function(_TO_DIM_ORDER_COPY, args=(ph,), kwargs=kwargs) + clone.meta["val"] = torch.randn(1, 3, 8, 8) + g.output((clone,)) + gm = torch.fx.GraphModule({}, g) + + result = ReplaceChannelsLastInputClones()(gm) + + assert not result.modified + assert _count(result.graph_module, _TO_DIM_ORDER_COPY) == 1 + + +class TestReplaceChannelsLastInputCloneDimOrders: + """These tests use channels last example inputs for export and apply the `EnforceContiguousDimOrder` pass which + inserts a `clone_dim_order` operator right after the model inputs to make the dim order contiguous. This is + precisely the intended use-case for the `ReplaceChannelsLastInputClones`. + """ + + def test_single_input(self): + example_inputs = ( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last), + ) + + ep = _export_to_edge(SingleInputModule(), example_inputs) + assert _count(ep.graph_module, _CLONE_DIM_ORDER) == 0 + + # Turn the model contiguous and create the input `clone_dim_order` operator. + # SingleInputModule preserves channels-last format (avg_pool + relu), so ECDO inserts + # both an input boundary clone and an output boundary clone. + res1 = EnforceContiguousDimOrder()(ep.graph_module) + assert res1.modified + assert ( + _count(res1.graph_module, _CLONE_DIM_ORDER) == 2 + ) # 1 input + 1 output boundary + + output_before = ep.module()(*example_inputs) + gm, modified = _run_pass(res1) + + assert modified + assert _count(gm, _TO_DIM_ORDER_COPY) == 0 + # Input boundary clone is replaced by permutes; the output boundary clone remains. + assert _count(gm, _CLONE_DIM_ORDER) == 1 + + nodes = list(gm.graph.nodes) + _assert_expected_result_pattern(*nodes[:3]) + + outputs_after = gm(*example_inputs)[0] + assert torch.allclose(output_before, outputs_after) + + def test_multi_input(self): + num_inputs = 3 + example_inputs = tuple( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last) + for _ in range(num_inputs) + ) + + ep = _export_to_edge(MultiInputModule(), example_inputs) + assert _count(ep.graph_module, _CLONE_DIM_ORDER) == 0 + + # Turn the model contiguous and create the input `clone_dim_order` operator. + # MultiInputModule (avg_pool + relu) preserves channels-last, so ECDO inserts one + # input boundary clone per input PLUS one output boundary clone. + res1 = EnforceContiguousDimOrder()(ep.graph_module) + assert res1.modified + assert _count(res1.graph_module, _CLONE_DIM_ORDER) == num_inputs + 1 + + output_before = ep.module()(*example_inputs) + res2 = ReplaceChannelsLastInputClones()(res1.graph_module) + + assert res2.modified + assert _count(res2.graph_module, _TO_DIM_ORDER_COPY) == 0 + # Input boundary clones are replaced by permutes; the output boundary clone remains. + assert _count(res2.graph_module, _CLONE_DIM_ORDER) == 1 + + nodes = list(res2.graph_module.graph.nodes) + for i in range(num_inputs): + # Each input should have the expected pattern. + _assert_expected_result_pattern(*nodes[i * num_inputs : i * num_inputs + 3]) + + outputs_after = res2.graph_module(*example_inputs)[0] + assert torch.allclose(output_before, outputs_after) + + def test_idempotency(self): + """Running the pass twice must not alter the graph on the second run.""" + example_inputs = ( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last), + ) + + ep = _export_to_edge(SingleInputModule(), example_inputs) + res1 = EnforceContiguousDimOrder()(ep.graph_module) + assert res1.modified + + res2 = ReplaceChannelsLastInputClones()(res1.graph_module) + assert res2.modified + + # Second pass: input boundary clones are gone; only the output boundary clone remains. + res3 = ReplaceChannelsLastInputClones()(res2.graph_module) + assert not res3.modified + # Input boundary clones have been replaced; the output boundary clone is untouched. + assert _count(res3.graph_module, _CLONE_DIM_ORDER) == 1 + assert _count(res3.graph_module, _TO_DIM_ORDER_COPY) == 0 + assert _count(res3.graph_module, _ATEN_PERMUTE_COPY) == _count( + res2.graph_module, _ATEN_PERMUTE_COPY + ) + assert _count(res3.graph_module, _CHANNELS_LAST_PERMUTE_COPY) == _count( + res2.graph_module, _CHANNELS_LAST_PERMUTE_COPY + ) + + +class TestGuardPredicates: + """Unit tests for the `call_operator` guard conditions using hand-built graphs.""" + + def test_is_4d_channels_last(self): + assert _is_4d_channels_last([0, 2, 3, 1]) + assert not _is_4d_channels_last([0, 1, 2, 3]) + assert not _is_4d_channels_last([0, 2, 1]) + assert not _is_4d_channels_last([0, 2, 3, 4, 1]) + + def test_is_4d_contiguous(self): + assert _is_4d_contiguous([0, 1, 2, 3]) + assert not _is_4d_contiguous([0, 2, 3, 1]) + assert not _is_4d_contiguous([0, 1, 2]) + assert not _is_4d_contiguous([0, 1, 2, 3, 4]) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index eed287b3a08..add4d01a78e 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1412,6 +1412,7 @@ def register_index_select(): return OpFeatures( inputs_storage=utils.CHANNELS_PACKED_TEXTURE, inputs_dtypes=utils.FP_INT_BOOL_T, + supports_resize=True, ) @@ -1435,58 +1436,56 @@ def register_where(): # ============================================================================= -@update_features(exir_ops.edge.aten.index.Tensor) -def register_index_tensor(): - def _index_tensor_shapes(node: torch.fx.Node): - """(self_val, index_val) for the supported single-index form, else None.""" - self_arg = node.args[0] - indices = node.args[1] +def _index_tensor_shapes(node: torch.fx.Node): + """Return self, index, and axis for the supported form, else None.""" + self_arg = node.args[0] + indices = node.args[1] - if not isinstance(self_arg, torch.fx.Node): - return None - self_val = self_arg.meta.get("val", None) - if self_val is None: - return None + if not isinstance(self_arg, torch.fx.Node): + return None + self_val = self_arg.meta.get("val", None) + if self_val is None or not isinstance(indices, (list, tuple)): + return None - # Only support exactly one non-None index tensor, applied to dim 0. - if not isinstance(indices, (list, tuple)): - return None - non_none = [idx for idx in indices if idx is not None] - if len(non_none) != 1 or indices[0] is None: - return None - index_arg = non_none[0] - if not isinstance(index_arg, torch.fx.Node): - return None - index_val = index_arg.meta.get("val", None) - if index_val is None: - return None + non_none = [(dim, index) for dim, index in enumerate(indices) if index is not None] + if len(non_none) != 1: + return None + index_dim, index_arg = non_none[0] + if index_dim >= len(self_val.size()) or not isinstance(index_arg, torch.fx.Node): + return None + index_val = index_arg.meta.get("val", None) + if index_val is None: + return None - return self_val, index_val + return self_val, index_val, index_dim - def check_index_tensor_node(node: torch.fx.Node) -> bool: - shapes = _index_tensor_shapes(node) - if shapes is None: - return False - _, index_val = shapes - # The gather is expressed as "one index position per output slice", so - # the index must be 1-D. `self` may be any rank: the buffer shader - # copies self's trailing dims through unchanged. - return len(index_val.size()) == 1 - def pick_index_tensor_storage(node: torch.fx.Node): - shapes = _index_tensor_shapes(node) - # Only the buffer shader handles a higher-rank `self`; the texture - # variant still assumes the 1-D form (it reads self[idx, 0, 0, 0]). - if shapes is not None and len(shapes[0].size()) > 1: - return utils.CONTIGUOUS_BUFFER, utils.CONTIGUOUS_BUFFER - return utils.ANY_STORAGE, utils.ANY_STORAGE +def _check_index_tensor_node(node: torch.fx.Node) -> bool: + shapes = _index_tensor_shapes(node) + if shapes is None: + return False + _, index_val, _ = shapes + # The gather is expressed as "one index position per output slice", so + # the index must be 1-D. `self` may be any rank. + return len(index_val.size()) == 1 + + +def _pick_index_tensor_storage(node: torch.fx.Node): + shapes = _index_tensor_shapes(node) + # Only the buffer shader handles a higher-rank `self`. + if shapes is not None and len(shapes[0].size()) > 1: + return utils.CONTIGUOUS_BUFFER, utils.CONTIGUOUS_BUFFER + return utils.ANY_STORAGE, utils.ANY_STORAGE + +@update_features(exir_ops.edge.aten.index.Tensor) +def register_index_tensor(): return OpFeatures( inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_INT_T, supports_resize=True, - are_node_inputs_supported_fn=check_index_tensor_node, - pick_io_storage_fn=pick_index_tensor_storage, + are_node_inputs_supported_fn=_check_index_tensor_node, + pick_io_storage_fn=_pick_index_tensor_storage, ) @@ -1500,6 +1499,7 @@ def register_arange(): return OpFeatures( inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_INT_T, + supports_resize=True, ) @@ -1508,12 +1508,34 @@ def register_arange(): # ============================================================================= +def _check_pad_is_static(node: torch.fx.Node) -> bool: + """Only support constant_pad_nd when the pad amounts are static. + + A symbolic pad list is serialized as a VALUELIST rather than an INTLIST, and + Pad.cpp reads it with get_int_list(), which throws "Expected value to have + type IntList, got VALUELIST instead". + + Supporting it properly is more than swapping in + extract_int_or_symint_list(), the way Split.cpp, View.cpp and Expand.cpp + read their symbolic lists: add_constant_pad_nd_node() folds the amounts into + a per-dim offset and bakes that into a params buffer at BUILD time, so the + dispatch would still use stale offsets even if the list were read + symbolically. The buffer has to be refreshed on resize first. Decline the + node until then. + """ + pad = node.args[1] + if not isinstance(pad, (list, tuple)): + return False + return all(isinstance(p, int) for p in pad) + + @update_features(exir_ops.edge.aten.constant_pad_nd.default) def register_constant_pad_nd(): return OpFeatures( inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_INT_BOOL_T, supports_resize=True, + are_node_inputs_supported_fn=_check_pad_is_static, ) @@ -1735,6 +1757,22 @@ def register_embedding_q4gsw(): # ============================================================================= +def _check_batch_norm_is_4d(node: torch.fx.Node) -> bool: + """Only support batch norm on a 4d input. + + add_native_batch_norm_node() asserts + VK_CHECK_COND(in_sizes.size() == 4, "BatchNorm only support 4d tensor") on + both the input and the output, so partitioning a batch norm whose input is + not 4d yields a .pte that lowers cleanly and then aborts at execute time. + Any conv1d model reaches here with rank-3 activations. + """ + input_node = node.args[0] + if not isinstance(input_node, torch.fx.Node): + return False + val = input_node.meta.get("val") + return val is not None and val.dim() == 4 + + @update_features(exir_ops.edge.aten._native_batch_norm_legit_no_training.default) def register_native_batch_norm_legit_no_training(): return OpFeatures( @@ -1742,6 +1780,7 @@ def register_native_batch_norm_legit_no_training(): inputs_dtypes=utils.FP_T, supports_prepacking=True, supports_resize=True, + are_node_inputs_supported_fn=_check_batch_norm_is_4d, ) diff --git a/backends/vulkan/quantizer/vulkan_quantizer.py b/backends/vulkan/quantizer/vulkan_quantizer.py index 3d1e1eab0f2..c7cc80022f8 100644 --- a/backends/vulkan/quantizer/vulkan_quantizer.py +++ b/backends/vulkan/quantizer/vulkan_quantizer.py @@ -47,7 +47,9 @@ def get_symmetric_quantization_config( Return a QuantizationConfig for Vulkan quantizer. Args: - is_dynamic: If False, weight-only quantization. If True, dynamic quantization (activation + weight) + is_dynamic: If False, weight-only quantization. If True, dynamic + quantization (activation + weight), with the activation scale + computed per tensor at runtime weight_bits: Number of bits for weight quantization (4 or 8) act_bits: Number of bits for activation quantization (8) act_qmin: Minimum quantization value for activations (auto-calculated if None) @@ -87,7 +89,20 @@ def get_symmetric_quantization_config( act_quantization_spec = None output_activation_spec = None else: - # Dynamic quantization: per-token input quantization, no output quantization + # Dynamic quantization: a choose_qparams op computes one scale and + # zero point for the whole activation tensor at runtime, and the + # quantize/dequantize pair around the linear carries them. Per tensor, + # not per token, whatever the granularity of the surrounding graph. + # + # (The fused et_vk.linear_q8ta_q8csw kernel is a different path: it is + # matched when the input scale is a static scalar, not one chosen at + # runtime.) + # + # One scale for the whole tensor is a poor fit for transformer + # encoders, where a few outlier channels set it for everything else; on + # sentence-transformer models this costs an order of magnitude more + # accuracy than a per-token scheme. Prefer is_dynamic=False (weight + # only) when output fidelity matters. # Auto-calculate activation ranges if not provided if act_qmin is None or act_qmax is None: act_range = bits_to_range(act_bits) diff --git a/backends/vulkan/runtime/graph/ComputeGraph.cpp b/backends/vulkan/runtime/graph/ComputeGraph.cpp index f23d1f19c66..f84eac26d49 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.cpp +++ b/backends/vulkan/runtime/graph/ComputeGraph.cpp @@ -11,6 +11,8 @@ #include +#include + #include #include @@ -121,16 +123,16 @@ TmpTensor::~TmpTensor() { } int64_t TmpTensor::get_sobj_idx() { - int64_t sobj_idx; + int64_t idx; // If no available temporary shared objects, request a new one to be created if (graph_p->tmp_shared_object_idxs_.empty()) { - sobj_idx = graph_p->shared_objects_.size(); + idx = graph_p->shared_objects_.size(); } else { // Get the first available shared object idx - sobj_idx = graph_p->tmp_shared_object_idxs_.top(); + idx = graph_p->tmp_shared_object_idxs_.top(); graph_p->tmp_shared_object_idxs_.pop(); } - return sobj_idx; + return idx; } // @@ -240,27 +242,34 @@ utils::StorageType ComputeGraph::suggested_storage_type() { return utils::kTexture3D; } -bool ComputeGraph::was_value_updated(const ValueRef idx) const noexcept { - if (!is_valid_value_idx(idx)) { - return false; +bool ComputeGraph::was_value_list_updated(const ValueRef idx) const noexcept { + const auto& value_list = values_[static_cast(idx)].toConstValueList(); + for (const auto nested_idx : value_list) { + if (was_value_updated(nested_idx)) { + return true; + } } + return false; +} - // Check if this ValueRef itself was updated - if (updated_values_.find(idx) != updated_values_.end()) { - return true; +void ComputeGraph::mark_value_updated(const ValueRef idx) { + if (!is_valid_value_idx(idx)) { + return; } - - // If this is a ValueList, check each ValueRef in the list - if (val_is_value_list(idx)) { - const auto& value_list = values_.at(idx).toConstValueList(); - for (const auto& nested_idx : value_list) { - if (was_value_updated(nested_idx)) { - return true; - } - } + if (value_update_generations_.size() < values_.size()) { + value_update_generations_.resize(values_.size()); } + value_update_generations_[static_cast(idx)] = + current_update_generation_; +} - return false; +void ComputeGraph::advance_update_generation() noexcept { + current_update_generation_++; + if (current_update_generation_ == 0) { + std::fill( + value_update_generations_.begin(), value_update_generations_.end(), 0); + current_update_generation_ = 1; + } } utils::GPUMemoryLayout ComputeGraph::suggested_memory_layout( @@ -775,8 +784,7 @@ void ComputeGraph::set_symint(const ValueRef idx, const int32_t val) { int32_t cur_val = read_symint(idx); if (cur_val != val) { get_symint(idx)->set(val); - // Track that this ValueRef was updated - updated_values_.insert(idx); + mark_value_updated(idx); } } @@ -1047,6 +1055,8 @@ void ComputeGraph::maybe_cast_and_copy_from_staging( } void ComputeGraph::prepare() { + value_update_generations_.resize(values_.size()); + #define MERGE_FIELD(field) \ static_cast(std::ceil( \ std::max( \ @@ -1202,9 +1212,9 @@ void ComputeGraph::prepack() { shared_object.bind_users(this); } // Make sure all remaining tensors have allocations - for (int i = 0; i < values_.size(); i++) { - if (values_.at(i).isTensor()) { - create_dedicated_allocation_for(i); + for (int value_idx = 0; value_idx < values_.size(); value_idx++) { + if (values_.at(value_idx).isTensor()) { + create_dedicated_allocation_for(value_idx); } } } @@ -1258,8 +1268,7 @@ void ComputeGraph::execute() { execute_count_++; - // Clear the set of updated values at the end of inference - updated_values_.clear(); + advance_update_generation(); // Reset the re-encoding flag at the end of inference requires_reencode_ = false; @@ -1281,7 +1290,7 @@ void ComputeGraph::resize_input( const std::vector& new_sizes) { IOValueRef io_val = inputs_.at(idx); virtual_resize(io_val.value, new_sizes); - updated_values_.insert(io_val.staging); + mark_value_updated(io_val.staging); } void ComputeGraph::virtual_resize( @@ -1290,8 +1299,7 @@ void ComputeGraph::virtual_resize( std::vector cur_sizes = sizes_of(idx); if (cur_sizes != new_sizes) { get_tensor(idx)->virtual_resize(new_sizes); - // Track that this ValueRef was updated - updated_values_.insert(idx); + mark_value_updated(idx); } } diff --git a/backends/vulkan/runtime/graph/ComputeGraph.h b/backends/vulkan/runtime/graph/ComputeGraph.h index 1de890efb38..eb01e3abf5e 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.h +++ b/backends/vulkan/runtime/graph/ComputeGraph.h @@ -204,8 +204,8 @@ class ComputeGraph final { // List of command buffers deferred for submission std::vector deferred_cmd_list_; - // Set to track which ValueRefs were updated during inference - std::unordered_set updated_values_; + std::vector value_update_generations_; + uint32_t current_update_generation_ = 1; // Cache to prevent duplicate prepacking of the same weight tensor with the // same kernel. Key is (inputValueRef, kernel_name). @@ -712,6 +712,7 @@ class ComputeGraph final { private: void check_no_active_value_ptrs(); + bool was_value_list_updated(const ValueRef idx) const noexcept; public: /* @@ -1174,7 +1175,22 @@ class ComputeGraph final { // Check if a specific ValueRef (or ValueList) was updated, with recursive // handling - bool was_value_updated(const ValueRef idx) const noexcept; + inline bool was_value_updated(const ValueRef idx) const noexcept { + if (idx < 0) { + return false; + } + + const size_t value_idx = static_cast(idx); + if (value_idx >= values_.size()) { + return false; + } + if (value_idx < value_update_generations_.size() && + value_update_generations_[value_idx] == current_update_generation_) { + return true; + } + + return values_[value_idx].isValueList() && was_value_list_updated(idx); + } // Set the flag to indicate that re-encoding is required inline void set_requires_reencode() noexcept { @@ -1222,6 +1238,10 @@ class ComputeGraph final { void print_readable(); + private: + void mark_value_updated(const ValueRef idx); + void advance_update_generation() noexcept; + // // Friend classes // diff --git a/backends/vulkan/runtime/graph/containers/SharedObject.cpp b/backends/vulkan/runtime/graph/containers/SharedObject.cpp index 10ddd6f2ca3..7ea59d59e2e 100644 --- a/backends/vulkan/runtime/graph/containers/SharedObject.cpp +++ b/backends/vulkan/runtime/graph/containers/SharedObject.cpp @@ -10,6 +10,8 @@ #include +#include + namespace vkcompute { bool SharedObject::has_user(const ValueRef idx) const { diff --git a/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp b/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp index a1a089c88e5..a2bf4268b01 100644 --- a/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp +++ b/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp @@ -39,8 +39,11 @@ bool ExecuteNode::trigger_resize(ComputeGraph* graph) { } bool ExecuteNode::was_any_arg_updated(const ComputeGraph* const graph) const { - // Check all ValueRefs in ArgGroups + // Check input args. for (const auto& arg_group : args_) { + if (!(arg_group.access & vkapi::kRead)) { + continue; + } for (const auto& value_ref : arg_group.refs) { if (graph->was_value_updated(value_ref)) { return true; @@ -48,13 +51,25 @@ bool ExecuteNode::was_any_arg_updated(const ComputeGraph* const graph) const { } } - // Check all ValueRefs in resize_args + // Check resize args. for (const auto& value_ref : resize_args_) { if (graph->was_value_updated(value_ref)) { return true; } } + // Check output args. + for (const auto& arg_group : args_) { + if (arg_group.access & vkapi::kRead) { + continue; + } + for (const auto& value_ref : arg_group.refs) { + if (graph->was_value_updated(value_ref)) { + return true; + } + } + } + return false; } diff --git a/backends/vulkan/runtime/graph/ops/PrepackNode.h b/backends/vulkan/runtime/graph/ops/PrepackNode.h index 8a301ef1e0a..fee7cf6fa60 100644 --- a/backends/vulkan/runtime/graph/ops/PrepackNode.h +++ b/backends/vulkan/runtime/graph/ops/PrepackNode.h @@ -50,6 +50,10 @@ class PrepackNode final { node_id_ = node_id; } + inline const std::string& name() const { + return shader_.kernel_name; + } + protected: uint32_t node_id_; const vkapi::ShaderInfo shader_; diff --git a/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl index 2e9377533c8..9bffc1c4132 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl @@ -23,18 +23,28 @@ layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")} ${layout_declare_ubo(B, "BufferMetadata", "outp")} -${layout_declare_ubo(B, "float", "start")} -${layout_declare_ubo(B, "float", "step")} +${layout_declare_ubo(B, "uint", "start")} +${layout_declare_ubo(B, "uint", "step")} + +layout(push_constant) uniform restrict Block { + ivec2 params_are_int; +}; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "dispatch.glslh" +float decode_param(const uint value, const int is_int) { + return is_int != 0 ? float(int(value)) : uintBitsToFloat(value); +} + void main() { const uint out_bufi = linear_idx_from_gid(); if (out_of_bounds(out_bufi, outp)) { return; } - t_out[out_bufi] = T(start + out_bufi * step); + const float start_val = decode_param(start, params_are_int.x); + const float step_val = decode_param(step, params_are_int.y); + t_out[out_bufi] = T(start_val + out_bufi * step_val); } diff --git a/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl b/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl index 0a5636b300f..73c2b5e5dd6 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl @@ -23,14 +23,22 @@ layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")} ${layout_declare_ubo(B, "TextureMetadata", "outp")} -${layout_declare_ubo(B, "float", "start")} -${layout_declare_ubo(B, "float", "step")} +${layout_declare_ubo(B, "uint", "start")} +${layout_declare_ubo(B, "uint", "step")} + +layout(push_constant) uniform restrict Block { + ivec2 params_are_int; +}; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; ${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")} const int packed_dim = get_packed_dim(out_layout); +float decode_param(const uint value, const int is_int) { + return is_int != 0 ? float(int(value)) : uintBitsToFloat(value); +} + void main() { const ivec3 out_pos = ivec3(gl_GlobalInvocationID); @@ -44,11 +52,13 @@ void main() { // arange output is 1D, so the W dimension holds the element index. // Compute the value for each element in the texel along the packed dim. VEC4_T outtex = VEC4_T(0); + const float start_val = decode_param(start, params_are_int.x); + const float step_val = decode_param(step, params_are_int.y); int limit = min( 4, safe_idx(outp.sizes, packed_dim) - out_tidx.data[packed_dim]); for (int comp = 0; comp < limit; comp++) { int elem_idx = out_tidx.data[0]; // W index is the linear element index - outtex[comp] = VEC4_T(start + elem_idx * step).x; + outtex[comp] = VEC4_T(start_val + elem_idx * step_val).x; out_tidx.data[packed_dim]++; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml index 1f217acb127..d1f3600cfc9 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml @@ -30,9 +30,15 @@ binary_op_buffer: OPERATOR: floor(X / Y) - NAME: binary_minimum_buffer OPERATOR: min(X, Y) - - NAME: binary_eq_int32_buffer + # Named to match the dispatcher, which builds "binary_" + op + storage + + # dtype. A variant called binary_eq_int32_buffer generates shaders nothing + # ever asks for, and aten.eq on an int32 tensor aborts at dispatch with + # "Could not find ShaderInfo with name binary_eq_buffer_int32". + - NAME: binary_eq_buffer OPERATOR: X == Y - DTYPE: int32 + generate_variant_forall: + DTYPE: + - VALUE: int32 - NAME: binary_eq_buffer OPERATOR: abs(X - Y) < 1e-5 generate_variant_forall: diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.yaml b/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.yaml index 289466e7845..07e3905af02 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.yaml @@ -32,9 +32,13 @@ binary_op_texture: MASK_PADDING: 1 - NAME: binary_minimum_texture3d OPERATOR: min(X, Y) - - NAME: binary_eq_int32_texture3d + # See binary_op_buffer.yaml: the name has to match what the dispatcher + # builds, otherwise the generated shader is unreachable. + - NAME: binary_eq_texture3d OPERATOR: equal(X, Y) - DTYPE: int32 + generate_variant_forall: + DTYPE: + - VALUE: int32 - NAME: binary_eq_texture3d OPERATOR: lessThan(abs(X - Y), VEC4_T(1e-5)) generate_variant_forall: diff --git a/backends/vulkan/runtime/graph/ops/glsl/index_select.glsl b/backends/vulkan/runtime/graph/ops/glsl/index_select.glsl index 4500d43b932..8b44daa9621 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/index_select.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/index_select.glsl @@ -20,7 +20,8 @@ ${layout_declare_tensor(0, "w", "t_out", DTYPE, STORAGE)} ${layout_declare_tensor(1, "r", "t_in", DTYPE, STORAGE)} ${layout_declare_tensor(2, "r", "t_idx", "int", STORAGE)} ${layout_declare_ubo(3, "ivec4", "sizes")} -${layout_declare_ubo(4, "int", "gpu_dim", "int", "stride")} +${layout_declare_ubo(4, "ivec4", "in_sizes")} +${layout_declare_ubo(5, "int", "gpu_dim")} layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; @@ -33,6 +34,13 @@ void main() { return; } + // Selecting along the batch dim steps over the z axis in units of channel + // texels, because the batch and channel dims share that axis. The width and + // height dims each have an axis to themselves, so they step by one. The + // stride is read from in_sizes rather than baked in at build time so that it + // follows a resize that changes the channel count. + const int stride = gpu_dim == 2 ? ((in_sizes.z + 3) / 4) : 1; + const int out_idx = out_pos[gpu_dim] / stride; const int within_stride = out_pos[gpu_dim] % stride; const int in_idx = texelFetch(t_idx, ivec3(out_idx, 0, 0), 0).x; diff --git a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl index db61e0859f2..b2497f98ccc 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl @@ -9,6 +9,7 @@ #version 450 core ${define_required_extensions("buffer", DTYPE)} +${define_required_extensions(INDEX_STORAGE, "int")} #define PRECISION ${PRECISION} @@ -22,19 +23,68 @@ layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")} ${layout_declare_tensor(B, "r", "t_self", DTYPE, "buffer")} -${layout_declare_tensor(B, "r", "t_index", "int", "buffer")} +${layout_declare_tensor(B, "r", "t_index", "int", INDEX_STORAGE)} ${layout_declare_ubo(B, "BufferMetadata", "outp")} ${layout_declare_ubo(B, "BufferMetadata", "inp")} -${layout_declare_ubo(B, "BufferMetadata", "index")} +$if INDEX_STORAGE == "buffer": + ${layout_declare_ubo(B, "BufferMetadata", "index")} +$else: + ${layout_declare_ubo(B, "TextureMetadata", "index")} + +layout(push_constant) uniform restrict Block { + ivec2 index_params; +}; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "dispatch.glslh" -// Implements aten.index.Tensor for the case where self is 1D and there is -// exactly one index tensor. Each output element is: +// Implements aten.index.Tensor with exactly one index tensor. Each output +// element is: // output[...] = self[index[...]] +${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")} +${layout_declare_spec_const(C, "int", "inp_layout", "CONTIG_LAYOUT_INT")} +${layout_declare_spec_const(C, "int", "index_layout", "CONTIG_LAYOUT_INT")} + +int load_index(const TensorIndex out_tidx) { +$if INDEX_STORAGE == "buffer": + uint index_bufi = 0; + for (int d = 0; d < index_params.y; ++d) { + index_bufi += + stride_at(index, d) * idx_at(out_tidx, index_params.x + d); + } + return t_index[index_bufi]; +$else: + TensorIndex4D index_tidx = zero_tensor4d_idx(); + index_tidx.data.x = int(idx_at(out_tidx, index_params.x)); + if (index_params.y > 1) { + index_tidx.data.y = int(idx_at(out_tidx, index_params.x + 1)); + } + if (index_params.y > 2) { + index_tidx.data.z = int(idx_at(out_tidx, index_params.x + 2)); + } + if (index_params.y > 3) { + index_tidx.data.w = int(idx_at(out_tidx, index_params.x + 3)); + } + const TextureElementIndex index_elem = + tensor4d_idx_to_texture_element_idx_simple( + index, index_tidx, index_layout); + return texelFetch(t_index, index_elem.pos, 0)[index_elem.comp]; +} + +uint self_idx_at( + const TensorIndex out_tidx, + const int self_axis, + const uint index_value) { + if (self_axis == index_params.x) { + return index_value; + } + const int out_axis = self_axis < index_params.x + ? self_axis + : self_axis + index_params.y - 1; + return idx_at(out_tidx, out_axis); +} void main() { const uint out_bufi = linear_idx_from_gid(); @@ -45,22 +95,20 @@ void main() { // Convert output buffer index to tensor index TensorIndex out_tidx = linear_idx_to_tensor_idx(outp, out_bufi); - const uint self_rank = ndim(inp); - const uint index_rank = ndim(index); - // WHCN order places self's trailing axes before the index axes. - const uint index_axis_offset = self_rank - 1; - - uint index_bufi = 0; - for (uint d = 0; d < index_rank; ++d) { - index_bufi += - stride_at(index, d) * idx_at(out_tidx, index_axis_offset + d); - } - const int idx = t_index[index_bufi]; - - uint self_bufi = stride_at(inp, self_rank - 1) * uint(idx); - for (uint d = 0; d + 1 < self_rank; ++d) { - self_bufi += stride_at(inp, d) * idx_at(out_tidx, d); - } + const int idx = load_index(out_tidx); + + TensorIndex self_tidx; + initialize(self_tidx); + const int self_rank = int_ndim(inp); + if (self_rank > 0) self_tidx.data[0].x = self_idx_at(out_tidx, 0, uint(idx)); + if (self_rank > 1) self_tidx.data[0].y = self_idx_at(out_tidx, 1, uint(idx)); + if (self_rank > 2) self_tidx.data[0].z = self_idx_at(out_tidx, 2, uint(idx)); + if (self_rank > 3) self_tidx.data[0].w = self_idx_at(out_tidx, 3, uint(idx)); + if (self_rank > 4) self_tidx.data[1].x = self_idx_at(out_tidx, 4, uint(idx)); + if (self_rank > 5) self_tidx.data[1].y = self_idx_at(out_tidx, 5, uint(idx)); + if (self_rank > 6) self_tidx.data[1].z = self_idx_at(out_tidx, 6, uint(idx)); + if (self_rank > 7) self_tidx.data[1].w = self_idx_at(out_tidx, 7, uint(idx)); + const uint self_bufi = tensor_idx_to_linear_idx(inp, self_tidx); t_out[out_bufi] = t_self[self_bufi]; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.yaml index ef79704203f..f4f168dfb37 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.yaml @@ -8,7 +8,11 @@ index_tensor_buffer: parameter_names_with_default_values: DTYPE: float STORAGE: buffer + INDEX_STORAGE: buffer generate_variant_forall: + INDEX_STORAGE: + - VALUE: buffer + - VALUE: texture3d DTYPE: - VALUE: half - VALUE: float diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.glsl index 5682f044b1d..8e1a07024ed 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.glsl @@ -9,6 +9,7 @@ #version 450 core #define PRECISION ${PRECISION} +#define ADD_UNSIGNED_OFFSET ${ADD_UNSIGNED_OFFSET} ${define_active_storage_type(STORAGE)} @@ -69,7 +70,11 @@ void main() { weight_vals[col] = (t_int8_weight[word_idx >> 2][word_idx & 3] >> (byte_pos * 8)) & 0xFF; } } - packed_block[row] = pack_into_int32(weight_vals); + int packed = pack_into_int32(weight_vals); +#if ADD_UNSIGNED_OFFSET == 1 + packed = int(uint(packed) ^ 0x80808080u); +#endif + packed_block[row] = packed; } buf_idx += oc_stride; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.yaml b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.yaml index 9331de6e758..190eff31f4f 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.yaml @@ -7,9 +7,12 @@ pack_q8_conv2d_weights: parameter_names_with_default_values: STORAGE: buffer + ADD_UNSIGNED_OFFSET: 0 generate_variant_forall: STORAGE: - VALUE: buffer - VALUE: texture2d shader_variants: - NAME: pack_q8_conv2d_weights + - NAME: pack_q8_conv2d_weights_unsigned + ADD_UNSIGNED_OFFSET: 1 diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.glsl b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.glsl index f2c74b67283..7b17dbe4990 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.glsl @@ -9,6 +9,7 @@ #version 450 core #define PRECISION ${PRECISION} +#define ADD_UNSIGNED_OFFSET ${ADD_UNSIGNED_OFFSET} ${define_active_storage_type(STORAGE)} @@ -53,6 +54,10 @@ void main() { load_block_data_with_checks(block , k4, n, K4, N); } +#if ADD_UNSIGNED_OFFSET == 1 + block.data = ivec4(uvec4(block.data) ^ uvec4(0x80808080u)); +#endif + // The weight blocks are stored in a tranposed manner, such that weight blocks // are indexed like packed_weight[k4][n4]. This is to optimize memory // coalescing when computing tiled GEMM. diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.yaml b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.yaml index 13e6d43b2c5..d05ffb467bc 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.yaml @@ -7,8 +7,15 @@ pack_q8_linear_weight: parameter_names_with_default_values: STORAGE: buffer + ADD_UNSIGNED_OFFSET: 0 shader_variants: - NAME: pack_q8_linear_weight_buffer STORAGE: buffer - NAME: pack_q8_linear_weight_texture2d STORAGE: texture2d + - NAME: pack_q8_linear_weight_unsigned_buffer + STORAGE: buffer + ADD_UNSIGNED_OFFSET: 1 + - NAME: pack_q8_linear_weight_unsigned_texture2d + STORAGE: texture2d + ADD_UNSIGNED_OFFSET: 1 diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl index aeb98f7a41b..6e258cbf2ac 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl @@ -11,6 +11,10 @@ ${define_required_extensions("buffer", DTYPE)} #define USE_INT8_DOT_PRODUCT_EXT ${USE_INT8_DOT_PRODUCT_EXT} +#define USE_UNSIGNED_DOT_PRODUCT ${USE_UNSIGNED_DOT_PRODUCT} + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER #extension GL_EXT_control_flow_attributes : require $if USE_INT8_DOT_PRODUCT_EXT == 1: @@ -43,7 +47,7 @@ layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_packed_int8_output", "int", "buffer", is_scalar_array=True)} ${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_packed_int8_weight", "int", "texture2d", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_int8_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} ${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=False)} ${layout_declare_tensor(B, "r", "t_weight_scales", DTYPE, "buffer", is_scalar_array=False)} ${layout_declare_tensor(B, "r", "t_bias", DTYPE, "buffer", is_scalar_array=False)} @@ -59,6 +63,7 @@ layout(push_constant) uniform restrict Block { int output_zp; int K4_per_group; int OC4_per_group; + int stream_row_offset; }; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; @@ -69,6 +74,9 @@ ${layout_declare_spec_const(C, "int", "activation_type", "0")} // Layout specialization constants ${layout_declare_spec_const(C, "int", "outp_layout", "CONTIG_LAYOUT_INT")} ${layout_declare_spec_const(C, "int", "inp_layout", "CONTIG_LAYOUT_INT")} +// Row-tile input (im2col scratch) vs batched activation input. Uniform per +// dispatch; declared last so existing constant ids are unchanged. +${layout_declare_spec_const(C, "int", "use_flat_tile", "0")} int compute_outp_buffer_idx( const int w_block_idx, @@ -103,9 +111,29 @@ void main() { const int W4 = div_up_4(int(outp.sizes[0][0])); const int H = int(outp.sizes[0][1]); const int OC4 = div_up_4(int(outp.sizes[0][2])); - const int hn = int(gl_GlobalInvocationID.z); - const int n = hn / H; - const int oh = hn % H; + const int local_row_idx = int(gl_GlobalInvocationID.z); + int n; + int oh; + int input_n; + int input_h; + if (use_flat_tile == 1) { + if (local_row_idx >= int(inp.sizes[0][1])) { + return; + } + const int global_row_idx = stream_row_offset + local_row_idx; + if (global_row_idx >= int(outp.sizes[0][3]) * H) { + return; + } + n = global_row_idx / H; + oh = global_row_idx % H; + input_n = 0; + input_h = local_row_idx; + } else { + n = local_row_idx / H; + oh = local_row_idx % H; + input_n = n; + input_h = oh; + } // Bounds check in block space if (ow_block_idx >= W4 || @@ -126,10 +154,19 @@ void main() { const int inp_n_stride = int(inp.strides[0][3]); // Initialize int32 accumulator +#if USE_UNSIGNED_DOT_PRODUCT == 1 + uvec4 out_accum[TILE_M][TILE_N4]; + uvec4 input_sums = uvec4(0u); +#else ivec4 out_accum[TILE_M][TILE_N4]; +#endif [[unroll]] for (int m = 0; m < TILE_M; ++m) { [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { +#if USE_UNSIGNED_DOT_PRODUCT == 1 + out_accum[m][n4] = uvec4(0u); +#else out_accum[m][n4] = ivec4(0); +#endif } } @@ -139,8 +176,8 @@ void main() { // Compute initial input tile index with group offset // For grouped im2col, each group's K range starts at group_idx * K4_per_group // For non-grouped (groups=1), group_idx is always 0 so offset is 0 - int input_idx = n * inp_n_stride - + oh * inp_h_stride + int input_idx = input_n * inp_n_stride + + input_h * inp_h_stride + ow_block_idx * inp_w_stride + group_idx * K4_per_group; @@ -149,26 +186,52 @@ void main() { // Load the packed int8 input tile for the current width and K sub-block. // Each int contains 4 packed int8s (one per width position in the tile) ivec4 int8_input_tile = t_packed_int8_input[input_idx]; +#if USE_UNSIGNED_DOT_PRODUCT == 1 + const uvec4 uint8_input_tile = + uvec4(int8_input_tile) ^ uvec4(0x80808080u); +#endif // Load the int8 weight tile for the current K and output-channel sub-block. ivec4 int8_weight_tile[TILE_N4]; [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { +#ifdef WEIGHT_BUFFER + int8_weight_tile[n4] = t_packed_int8_weight[ + k4 * OC4 + oc_block_idx + n4]; +#else int8_weight_tile[n4] = texelFetch( t_packed_int8_weight, ivec2(oc_block_idx + n4, k4), 0); +#endif } +#if USE_UNSIGNED_DOT_PRODUCT == 1 + uvec4 uint8_weight_tile[TILE_N4]; + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + uint8_weight_tile[n4] = uvec4(int8_weight_tile[n4]); + } +#endif // Accumulate using int8 dot product // Input tile indexed as input[m] where m is the width index within tile // Weight tile indexed as weight[n4][n4i] where n4i is the channel index within block [[unroll]] for (int m = 0; m < TILE_M; ++m) { +#if USE_UNSIGNED_DOT_PRODUCT == 1 + input_sums[m] = dotPacked4x8AccSatEXT( + uint8_input_tile[m], 0x01010101u, input_sums[m]); +#endif [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { [[unroll]] for (int n4i = 0; n4i < 4; ++n4i) { +#if USE_UNSIGNED_DOT_PRODUCT == 1 + out_accum[m][n4][n4i] = dotPacked4x8AccSatEXT( + uint8_input_tile[m], + uint8_weight_tile[n4][n4i], + out_accum[m][n4][n4i]); +#else out_accum[m][n4][n4i] = dotPacked4x8AccSat( int8_input_tile[m], int8_weight_tile[n4][n4i], out_accum[m][n4][n4i]); +#endif } } } @@ -188,6 +251,14 @@ void main() { weight_sums[n4] = ivec4(t_weight_sums[oc_block_idx + n4]); } +#if USE_UNSIGNED_DOT_PRODUCT == 1 + ivec4 unsigned_weight_correction[TILE_N4]; + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + unsigned_weight_correction[n4] = + (128 + input_zp) * weight_sums[n4]; + } +#endif + // Initialize int8 output tile ivec4 int8_out_tile[TILE_M4][TILE_N4]; [[unroll]] for (int m4 = 0; m4 < TILE_M4; ++m4) { @@ -197,7 +268,9 @@ void main() { } // Compute int8 output tile from int32 accumulator +#if USE_UNSIGNED_DOT_PRODUCT == 0 ivec4 input_zp_vec = ivec4(-input_zp); +#endif if (apply_bias > 0) { // Load bias tile @@ -211,8 +284,14 @@ void main() { [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { const int m = mul_4(m4) + m4i; // Compute floating point output values +#if USE_UNSIGNED_DOT_PRODUCT == 1 + ivec4 accum_adjusted = ivec4(out_accum[m][n4]) + - ivec4(int(input_sums[m]) * 128) + - unsigned_weight_correction[n4]; +#else ivec4 accum_adjusted = input_zp_vec * weight_sums[n4] + out_accum[m][n4]; +#endif vec4 float_out_texel = fma(vec4(accum_adjusted), vec4(weight_scales[n4]) * input_scale, @@ -236,8 +315,14 @@ void main() { [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { const int m = mul_4(m4) + m4i; // Compute floating point output values +#if USE_UNSIGNED_DOT_PRODUCT == 1 + ivec4 accum_adjusted = ivec4(out_accum[m][n4]) + - ivec4(int(input_sums[m]) * 128) + - unsigned_weight_correction[n4]; +#else ivec4 accum_adjusted = input_zp_vec * weight_sums[n4] + out_accum[m][n4]; +#endif vec4 float_out_texel = vec4(accum_adjusted) * vec4(weight_scales[n4] * input_scale); // Apply ReLU if enabled diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml index 46670b8d2aa..f2979add927 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml @@ -8,10 +8,17 @@ q8ta_conv2d_pw: parameter_names_with_default_values: DTYPE: float USE_INT8_DOT_PRODUCT_EXT: 1 + USE_UNSIGNED_DOT_PRODUCT: 0 + WEIGHT_STORAGE: texture2d generate_variant_forall: DTYPE: - VALUE: float shader_variants: - NAME: q8ta_conv2d_pw + - NAME: q8ta_conv2d_pw_unsigned + USE_UNSIGNED_DOT_PRODUCT: 1 + - NAME: q8ta_conv2d_pw_unsigned_buffer + USE_UNSIGNED_DOT_PRODUCT: 1 + WEIGHT_STORAGE: buffer - NAME: q8ta_conv2d_pw_fallback USE_INT8_DOT_PRODUCT_EXT: 0 diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl index b0cc4866a03..a278ce394d1 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl @@ -33,12 +33,42 @@ ${layout_declare_spec_const(C, "int", "inp_layout", "CONTIG_LAYOUT_INT")} layout(push_constant) uniform restrict Block { int zp; + int stream_row_offset; }; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "dispatch.glslh" +int load_packed_input( + const int x, + const int y, + const int z4, + const int n, + const int input_W, + const int input_H, + const int input_Z4, + const int zp_packed) { + if (x < 0 || x >= input_W || y < 0 || y >= input_H || z4 < 0 || + z4 >= input_Z4) { + return zp_packed; + } + + const int x4 = div_4(x); + const int x_mod = mod_4(x); + if (get_outer_packed_dim_block_size(inp_layout) == 1) { + const int scalar_idx = n * int(inp.strides[0][3]) + + y * int(inp.strides[0][1]) + + x * int(inp.strides[0][0]) + z4 * int(inp.strides[0][2]); + return t_packed_int8_input[scalar_idx]; + } + + const int scalar_idx = mul_4( + n * int(inp.strides[0][3]) + y * int(inp.strides[0][1]) + + x4 * int(inp.strides[0][0]) + z4) + x_mod; + return t_packed_int8_input[scalar_idx]; +} + void main() { const int out_buf_idx = int(linear_idx_from_gid()); @@ -50,24 +80,31 @@ void main() { const int im2col_W4 = div_up_4(im2col_sizes.x); const int im2col_H = im2col_sizes.y; const int im2col_Z4 = div_up_4(im2col_sizes.z); - const int im2col_N = im2col_sizes.w; - // im2col block index from linear output buffer index + // im2col block index from linear output buffer index. The scratch holds one + // row tile, so local rows are rebased onto the global batch*height rows. const int c4_idx = out_buf_idx % im2col_Z4; const int row = out_buf_idx / im2col_Z4; const int w4_idx = row % im2col_W4; const int hn_idx = row / im2col_W4; + const int local_row_idx = hn_idx; const int h_idx = hn_idx % im2col_H; - const int n_idx = hn_idx / im2col_H; + const int output_H = + (input_sizes.y + 2 * conv2d_params.padding.y - + conv2d_params.dilation.y * (conv2d_params.kernel_size.y - 1) - 1) / + conv2d_params.stride.y + + 1; + const int global_row_idx = stream_row_offset + local_row_idx; + const int n_idx = global_row_idx / output_H; // out of bounds check - if (w4_idx >= im2col_W4 || h_idx >= im2col_H || - c4_idx >= im2col_Z4 || n_idx >= im2col_N) { + if (w4_idx >= im2col_W4 || local_row_idx >= im2col_H || + c4_idx >= im2col_Z4 || n_idx >= input_sizes.w) { return; } const int im2col_w = mul_4(w4_idx); - const int im2col_h = h_idx; + const int im2col_h = global_row_idx % output_H; const int im2col_k = mul_4(c4_idx); const int group_idx = im2col_k / conv2d_params.K_per_group; @@ -96,40 +133,51 @@ void main() { const int zp_packed = pack_into_int32(ivec4(zp)); const int z4 = div_4(input_z); - // Check if y and z are in bounds (constant for all 4 width elements) - const bool y_z_in_bounds = - (input_y >= 0 && input_y < input_H && z4 >= 0 && z4 < input_Z4); - - // Load 4 elements from input, one for each output width position. - // Each loaded int contains 4 packed int8 channel values. - ivec4 im2col_block; - for (int i = 0; i < 4; i++) { - const int x = input_x_base + i * conv2d_params.stride.x; - if (!y_z_in_bounds || x < 0 || x >= input_W) { - im2col_block[i] = zp_packed; - } else { - const int x4 = div_4(x); - const int x_mod = mod_4(x); - int scalar_idx; - if (get_outer_packed_dim_block_size(inp_layout) == 1) { - scalar_idx = n_idx * int(inp.strides[0][3]) - + input_y * int(inp.strides[0][1]) - + x * int(inp.strides[0][0]) - + z4 * int(inp.strides[0][2]); - } else { - scalar_idx = mul_4( - n_idx * int(inp.strides[0][3]) - + input_y * int(inp.strides[0][1]) - + x4 * int(inp.strides[0][0]) - + z4) + x_mod; - } - im2col_block[i] = t_packed_int8_input[scalar_idx]; - } - } - - // store_packed_int8_output_tile (with TILE_M4=1, TILE_N4=1) - const int buffer_idx = n_idx * int(im2col_outp.strides[0][3]) - + h_idx * int(im2col_outp.strides[0][1]) + const int stride_x = conv2d_params.stride.x; + + // Keep lane loads and their complete bounds checks static; some mobile + // drivers lose lanes otherwise. + const ivec4 im2col_block = ivec4( + load_packed_input( + input_x_base, + input_y, + z4, + n_idx, + input_W, + input_H, + input_Z4, + zp_packed), + load_packed_input( + input_x_base + stride_x, + input_y, + z4, + n_idx, + input_W, + input_H, + input_Z4, + zp_packed), + load_packed_input( + input_x_base + 2 * stride_x, + input_y, + z4, + n_idx, + input_W, + input_H, + input_Z4, + zp_packed), + load_packed_input( + input_x_base + 3 * stride_x, + input_y, + z4, + n_idx, + input_W, + input_H, + input_Z4, + zp_packed)); + + // store_packed_int8_output_tile (with TILE_M4=1, TILE_N4=1). The scratch + // has a single batch, so every tile writes from row 0 of the same buffer. + const int buffer_idx = h_idx * int(im2col_outp.strides[0][1]) + w4_idx * int(im2col_outp.strides[0][0]) + c4_idx; diff --git a/backends/vulkan/runtime/graph/ops/glsl/reduce.glsl b/backends/vulkan/runtime/graph/ops/glsl/reduce.glsl index 209440cec6a..029e3b16756 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/reduce.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/reduce.glsl @@ -32,12 +32,17 @@ layout(constant_id = 5) const int group_dim = 1; // A more verbose name would be NWORKERS_PER_GROUP. This describes the number of // threads that will co-operate to compute one reduction output. There may be // multiple groups computing distinct reduction outputs within one work group. -#define NWORKERS 4 +// Supplied by the dispatch so it can scale with the length of the reduction. +// A global average pool reduces a whole HxW plane into one value, and four +// workers left the GPU essentially idle for it. +layout(constant_id = 6) const int NWORKERS = 4; // Sets an upper limit on the total size of a work group based on how many // elements are allocated in the shared memory array below. Each thread in the // work group will write into its assigned element in the shared array. -#define MAX_NTHREADS 16 +// Upper bound on NWORKERS * NGROUPS, and the size of the shared array below. +// 256 vec4 is 4 KiB of shared memory, well inside the guaranteed 16 KiB. +#define MAX_NTHREADS 256 shared vec4 shared_vecs[MAX_NTHREADS]; @@ -86,19 +91,30 @@ int tid_to_smi(const ivec2 tid) { * This case is simpler because each element of a texel belongs to a separate * reduction "group", meaning we don't have to perform reduction along a texel. */ -void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) { +void reduce_nonpacked_dim( + const ivec2 tid, + ivec3 scan_pos, + const bool in_bounds) { // shared memory index of this thread const int smi = tid_to_smi(tid); - scan_pos[reduce_dim] = 0; - vec4 accum = INIT_ACCUM(load_texel(tin, scan_pos)); - - scan_pos[reduce_dim] = tid.x; - // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... of - // the reduction row - for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim); - i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) { - accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + // Out of bounds invocations cannot return early: barrier() below has to be + // reached by every invocation in the work group, and skipping it is undefined + // behaviour that hangs some GPUs. They still take a shared memory slot, but + // it is one that no in-bounds group aggregates over, so what they leave in it + // is never read. + vec4 accum = vec4(0); + if (in_bounds) { + scan_pos[reduce_dim] = 0; + accum = INIT_ACCUM(load_texel(tin, scan_pos)); + + scan_pos[reduce_dim] = tid.x; + // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... + // of the reduction row + for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim); + i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) { + accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + } } // Write partial output to shared memory and synchronize work group shared_vecs[smi] = accum; @@ -106,7 +122,7 @@ void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) { // Since the reduction row is reduced to only one element, only the "main" // thread in the group needs aggregate the partial outputs - if (tid.x == 0) { + if (in_bounds && tid.x == 0) { // Iterate over the partial outputs to obtain the overall output int group_i = tid.y * NWORKERS; accum = shared_vecs[group_i++]; @@ -141,7 +157,10 @@ void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) { * elements in texels (which occur when the size of the packed dim is not a * multiple of 4) so that they do not influence the output of reduction. */ -void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) { +void reduce_packed_dim( + const ivec2 tid, + ivec3 scan_pos, + const bool in_bounds) { // shared memory index of this thread const int smi = tid_to_smi(tid); @@ -151,23 +170,32 @@ void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) { // handled specially if it has padding elements. const int reduce_len = safe_idx(tin_sizes, packed_dim) - nspill; - scan_pos[reduce_dim] = 0; - vec4 accum = INIT_ACCUM(vec4(load_texel(tin, scan_pos).x)); - - // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... of - // the reduction row - scan_pos[reduce_dim] = tid.x; - for (int i = tid.x * 4; i < reduce_len; - i += NWORKERS * 4, scan_pos[reduce_dim] += NWORKERS) { - accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); - } - // For the last texel in the dim, if there are padding elements then each - // element of the texel needs to be processed individually such that the - // padding elements are ignored - if (scan_pos[reduce_dim] == safe_idx(tin_limits, reduce_dim) - 1 && nspill > 0) { - const vec4 intex = load_texel(tin, scan_pos); - for (int i = 0; i < nspill; i++) { - accum.x = UPDATE_ACCUM(accum.x, intex[i]); + // Out of bounds invocations cannot return early: barrier() below has to be + // reached by every invocation in the work group, and skipping it is undefined + // behaviour that hangs some GPUs. They still take a shared memory slot, but + // it is one that no in-bounds group aggregates over, so what they leave in it + // is never read. + vec4 accum = vec4(0); + if (in_bounds) { + scan_pos[reduce_dim] = 0; + accum = INIT_ACCUM(vec4(load_texel(tin, scan_pos).x)); + + // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... + // of the reduction row + scan_pos[reduce_dim] = tid.x; + for (int i = tid.x * 4; i < reduce_len; + i += NWORKERS * 4, scan_pos[reduce_dim] += NWORKERS) { + accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + } + // For the last texel in the dim, if there are padding elements then each + // element of the texel needs to be processed individually such that the + // padding elements are ignored + if (scan_pos[reduce_dim] == safe_idx(tin_limits, reduce_dim) - 1 && + nspill > 0) { + const vec4 intex = load_texel(tin, scan_pos); + for (int i = 0; i < nspill; i++) { + accum.x = UPDATE_ACCUM(accum.x, intex[i]); + } } } // Write partial output to shared memory and synchronize work group @@ -176,7 +204,7 @@ void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) { // Since the reduction row is reduced to only one element, only the "main" // thread in the group needs aggregate the partial outputs - if (tid.x == 0) { + if (in_bounds && tid.x == 0) { // Iterate over the partial maximums to obtain the overall maximum int group_i = tid.y * NWORKERS; accum = shared_vecs[group_i++]; @@ -203,13 +231,13 @@ void main() { gl_LocalInvocationID[reduce_dim], gl_LocalInvocationID[group_dim]); - if (any(greaterThanEqual(scan_pos, tin_limits))) { - return; - } + const bool in_bounds = all(lessThan(scan_pos, tin_limits)); + // reduce_dim and packed_dim are specialization constants, so this branch is + // uniform across the work group and safe to take around a barrier. if (reduce_dim != packed_dim) { - reduce_nonpacked_dim(tid, scan_pos); + reduce_nonpacked_dim(tid, scan_pos, in_bounds); } else { - reduce_packed_dim(tid, scan_pos); + reduce_packed_dim(tid, scan_pos, in_bounds); } } diff --git a/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl b/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl index bd55025f534..58e7c6d0b3d 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl @@ -33,12 +33,17 @@ layout(constant_id = 6) const int group_dim = 2; // A more verbose name would be NWORKERS_PER_GROUP. This describes the number of // threads that will co-operate to compute one reduction output. There may be // multiple groups computing distinct reduction outputs within one work group. -#define NWORKERS 4 +// Supplied by the dispatch so it can scale with the length of the reduction. +// A global average pool reduces a whole HxW plane into one value, and four +// workers left the GPU essentially idle for it. +layout(constant_id = 7) const int NWORKERS = 4; // Sets an upper limit on the total size of a work group based on how many // elements are allocated in the shared memory array below. Each thread in the // work group will write into its assigned element in the shared array. -#define MAX_NTHREADS 16 +// Upper bound on NWORKERS * NGROUPS, and the size of the shared array below. +// 256 vec4 is 4 KiB of shared memory, well inside the guaranteed 16 KiB. +#define MAX_NTHREADS 256 shared vec4 shared_vecs[MAX_NTHREADS]; @@ -59,23 +64,34 @@ int tid_to_smi(const ivec2 tid) { // with the accumulator. #define POSTPROCESS(accum) ${POSTPROCESS} -void reduce_2d_non_packed_dim(const ivec2 tid, ivec3 scan_pos) { +void reduce_2d_non_packed_dim( + const ivec2 tid, + ivec3 scan_pos, + const bool in_bounds) { // shared memory index of this thread const int smi = tid_to_smi(tid); - scan_pos[reduce_dim1] = 0; - scan_pos[reduce_dim2] = 0; - vec4 accum = INIT_ACCUM(load_texel(tin, scan_pos)); - - // First dimension reduction - scan_pos[reduce_dim1] = tid.x; - for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim1); - i += NWORKERS, scan_pos[reduce_dim1] += NWORKERS) { - - // Second dimension reduction + // Out of bounds invocations cannot return early: barrier() below has to be + // reached by every invocation in the work group, and skipping it is undefined + // behaviour that hangs some GPUs. They still take a shared memory slot, but + // it is one that no in-bounds group aggregates over, so what they leave in it + // is never read. + vec4 accum = vec4(0); + if (in_bounds) { + scan_pos[reduce_dim1] = 0; scan_pos[reduce_dim2] = 0; - for (int j = 0; j < safe_idx(tin_sizes, reduce_dim2); j++, scan_pos[reduce_dim2]++) { - accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + accum = INIT_ACCUM(load_texel(tin, scan_pos)); + + // First dimension reduction + scan_pos[reduce_dim1] = tid.x; + for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim1); + i += NWORKERS, scan_pos[reduce_dim1] += NWORKERS) { + // Second dimension reduction + scan_pos[reduce_dim2] = 0; + for (int j = 0; j < safe_idx(tin_sizes, reduce_dim2); + j++, scan_pos[reduce_dim2]++) { + accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + } } } @@ -84,7 +100,7 @@ void reduce_2d_non_packed_dim(const ivec2 tid, ivec3 scan_pos) { barrier(); // Main thread aggregates results - if (tid.x == 0) { + if (in_bounds && tid.x == 0) { // Iterate over the partial outputs to obtain the overall output int group_i = tid.y * NWORKERS; accum = shared_vecs[group_i++]; @@ -121,9 +137,7 @@ void main() { gl_LocalInvocationID[reduce_dim1], gl_LocalInvocationID[group_dim]); - if (any(greaterThanEqual(scan_pos, tin_limits))) { - return; - } + const bool in_bounds = all(lessThan(scan_pos, tin_limits)); - reduce_2d_non_packed_dim(tid, scan_pos); + reduce_2d_non_packed_dim(tid, scan_pos, in_bounds); } \ No newline at end of file diff --git a/backends/vulkan/runtime/graph/ops/glsl/unary_op.glsl b/backends/vulkan/runtime/graph/ops/glsl/unary_op.glsl index 45aa3ed7133..5ef891e7439 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/unary_op.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/unary_op.glsl @@ -23,16 +23,23 @@ ${define_active_storage_type(STORAGE)} layout(std430) buffer; -${layout_declare_tensor(0, "w", "t_out", DTYPE, STORAGE)} -${layout_declare_tensor(1, "r", "t_in", DTYPE, STORAGE)} +${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)} +${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)} + +$if DYNAMIC_PARAMS: + ${layout_declare_ubo(B, "uint", "minimum")} + ${layout_declare_ubo(B, "uint", "maximum")} layout(push_constant) uniform restrict Block { $if STORAGE == "buffer": int numel; $else: ivec4 out_limits; -float minimum; -float maximum; +$if DYNAMIC_PARAMS: + ivec2 bounds_are_int; +$else: + float minimum; + float maximum; }; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; @@ -40,6 +47,11 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "dispatch.glslh" #include "activations.h" +$if DYNAMIC_PARAMS: + float decode_bound(const uint value, const int is_int) { + return is_int != 0 ? float(int(value)) : uintBitsToFloat(value); + } + #ifdef USING_BUFFER void main() { @@ -48,7 +60,13 @@ void main() { return; } - float in_val = float(t_in[i]); +$if DYNAMIC_PARAMS: + const T in_val = T(t_in[i]); + const T minimum_val = T(decode_bound(minimum, bounds_are_int.x)); + const T maximum_val = T(decode_bound(maximum, bounds_are_int.y)); + t_out[i] = T(op(in_val, minimum_val, maximum_val)); +$else: + const float in_val = float(t_in[i]); t_out[i] = T(op(in_val, minimum, maximum)); } @@ -62,6 +80,11 @@ void main() { } VEC4_T in_texel = texelFetch(t_in, pos, 0); +$if DYNAMIC_PARAMS: + const VEC4_T minimum_val = VEC4_T(decode_bound(minimum, bounds_are_int.x)); + const VEC4_T maximum_val = VEC4_T(decode_bound(maximum, bounds_are_int.y)); + imageStore(t_out, pos, op(in_texel, minimum_val, maximum_val)); +$else: imageStore(t_out, pos, VEC4_T(op(in_texel, minimum, maximum))); } diff --git a/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml b/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml index 46d12806149..fc70b54076b 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml @@ -3,6 +3,7 @@ unary_op: OPERATOR: clamp(X, A, B) DTYPE: float STORAGE: texture3d + DYNAMIC_PARAMS: false generate_variant_forall: DTYPE: - VALUE: half @@ -18,12 +19,19 @@ unary_op: - NAME: clamp_int32 OPERATOR: clamp(X, A, B) DTYPE: int32 + - NAME: clamp_dynamic_int32 + OPERATOR: clamp(X, A, B) + DTYPE: int32 + DYNAMIC_PARAMS: true + - NAME: clamp_dynamic + OPERATOR: clamp(X, A, B) + DYNAMIC_PARAMS: true - NAME: cos OPERATOR: cos(X) - NAME: exp OPERATOR: exp(X) - NAME: gelu - OPERATOR: 0.5 * X * (1 + tanh(sqrt(2 / 3.141593) * (X + 0.044715 * X * X * X))) + OPERATOR: 0.5 * X * (1 + tanh(clamp(sqrt(2 / 3.141593) * (X + 0.044715 * X * X * X), -15.0, 15.0))) - NAME: neg OPERATOR: -X - NAME: sigmoid diff --git a/backends/vulkan/runtime/graph/ops/glsl/view_convert_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/view_convert_buffer.yaml index 200d58e1217..a770d986fd5 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/view_convert_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/view_convert_buffer.yaml @@ -18,6 +18,7 @@ view_convert_buffer: - parameter_values: [uint8, float] - parameter_values: [uint8, half] - parameter_values: [uint8, int32] + - parameter_values: [int32, uint8] - parameter_values: [float, int32] - parameter_values: [float, half] - parameter_values: [half, float] diff --git a/backends/vulkan/runtime/graph/ops/glsl/view_convert_texture.yaml b/backends/vulkan/runtime/graph/ops/glsl/view_convert_texture.yaml index 47e9c43ee24..227a18b6ccf 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/view_convert_texture.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/view_convert_texture.yaml @@ -18,6 +18,7 @@ view_convert_texture: - parameter_values: [uint8, float] - parameter_values: [uint8, half] - parameter_values: [uint8, int32] + - parameter_values: [int32, uint8] - parameter_values: [float, int32] - parameter_values: [float, half] - parameter_values: [half, float] diff --git a/backends/vulkan/runtime/graph/ops/impl/Arange.cpp b/backends/vulkan/runtime/graph/ops/impl/Arange.cpp index f635c9282f2..839b94f5e75 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Arange.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Arange.cpp @@ -15,6 +15,8 @@ #include +#include + namespace vkcompute { void resize_arange_node( @@ -23,18 +25,22 @@ void resize_arange_node( const std::vector& extra_args) { const ValueRef out = args.at(0).refs.at(0); - int start_val = 0; - int step_val = 1; + double start_val = 0.0; + double step_val = 1.0; if (!graph->val_is_none(extra_args.at(0))) { - start_val = graph->extract_scalar(extra_args.at(0)); + start_val = graph->extract_scalar(extra_args.at(0)); } - const int end_val = graph->extract_scalar(extra_args.at(1)); + const double end_val = graph->extract_scalar(extra_args.at(1)); if (!graph->val_is_none(extra_args.at(2))) { - step_val = graph->extract_scalar(extra_args.at(2)); + step_val = graph->extract_scalar(extra_args.at(2)); } + VK_CHECK_COND(step_val != 0.0, "arange: step must be nonzero"); + const double range_size = (end_val - start_val) / step_val; + VK_CHECK_COND( + range_size >= 0.0, "arange: bounds are inconsistent with step sign"); const std::vector out_sizes = { - utils::div_up(end_val - start_val, step_val)}; + static_cast(std::ceil(range_size))}; graph->virtual_resize(out, out_sizes); } @@ -55,39 +61,35 @@ void check_arange_input( } } +vkapi::BufferBindInfo get_arange_param_buffer( + ComputeGraph& graph, + const ValueRef value, + const float default_value) { + if (graph.val_is_symint(value)) { + return graph.get_or_create_int_param_buffer(value); + } + return graph.create_params_buffer( + graph.extract_scalar_or(value, default_value)); +} + void add_arange_node( ComputeGraph& graph, const ValueRef start, const ValueRef end, const ValueRef step, const ValueRef out) { - float start_val = 0.0f; - float step_val = 1.0f; - if (graph.val_is_none(end)) { VK_THROW("arange: end must be specified!"); } - if (!graph.val_is_none(start)) { - if (graph.val_is_int(start)) { - start_val = static_cast(graph.extract_scalar(start)); - } else { - start_val = graph.extract_scalar(start); - } - } - if (!graph.val_is_none(step)) { - if (graph.val_is_int(step)) { - step_val = static_cast(graph.extract_scalar(step)); - } else { - step_val = graph.extract_scalar(step); - } - } - std::string kernel_name("arange"); kernel_name.reserve(kShaderNameReserve); add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); add_dtype_suffix(kernel_name, graph.dtype_of(out)); + const utils::ivec2 params_are_int = { + graph.val_is_symint(start) ? 1 : 0, graph.val_is_symint(step) ? 1 : 0}; + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), @@ -97,10 +99,10 @@ void add_arange_node( {{out, vkapi::kWrite}}, // Shader params buffers {graph.meta_ubo(out), - graph.create_params_buffer(start_val), - graph.create_params_buffer(step_val)}, + get_arange_param_buffer(graph, start, 0.0f), + get_arange_param_buffer(graph, step, 1.0f)}, // Push Constants - {}, + {PushConstantDataInfo(¶ms_are_int, sizeof(params_are_int))}, // Specialization Constants {graph.hashed_layout_of(out)}, // Resize Args diff --git a/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp b/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp index cd1f9510bad..1b27a53628e 100644 --- a/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp @@ -34,23 +34,6 @@ void resize_choose_qparams_per_row( graph->virtual_resize(input_zeros, new_sizes); } -vkapi::ShaderInfo pick_choose_qparams_per_row_shader( - ComputeGraph* graph, - const std::vector& args, - const std::vector& resize_args) { - (void)resize_args; - - const ValueRef input = args.at(1).refs.at(0); - const ValueRef input_zps = args.at(0).refs.at(1); - - std::string kernel_name = "choose_qparams_per_row"; - add_storage_type_suffix(kernel_name, graph->storage_type_of(input)); - add_dtype_suffix(kernel_name, graph->dtype_of(input)); - add_zp_dtype_mode_suffix(kernel_name, graph->dtype_of(input_zps)); - - return VK_KERNEL_FROM_STR(kernel_name); -} - GlobalWorkGrid pick_choose_qparams_per_row_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, @@ -99,9 +82,13 @@ void add_choose_qparams_per_row_node( PushConstantDataInfo(&quant_max_val, sizeof(int32_t)), }; + std::string kernel_name = "choose_qparams_per_row"; + add_storage_type_suffix(kernel_name, graph.storage_type_of(input)); + add_dtype_suffix(kernel_name, graph.dtype_of(input)); + add_zp_dtype_mode_suffix(kernel_name, graph.dtype_of(input_zps)); graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, - pick_choose_qparams_per_row_shader, + VK_KERNEL_FROM_STR(kernel_name), pick_choose_qparams_per_row_gwg, pick_required_lwg, // Inputs and Outputs diff --git a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp index 4e8ff50fdac..b839ac8a61a 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp @@ -350,6 +350,38 @@ GlobalWorkGrid create_conv2d_gwg( } } +// Determines which convolution method a dispatch uses. +// +// Depthwise and transposed convolutions have shader names of their own, but +// the name alone cannot separate pointwise from sliding window: the sliding +// window shader is itself named "conv2d", and a pointwise convolution also +// takes that name when its weights are prepacked. Those two are therefore +// separated by the weight's spatial extent. Shared by the global and local +// workgroup size functions below so that the two cannot disagree about the +// same dispatch. +Conv2dMethod infer_conv2d_method_from_shader( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const ValueRef weight_data) { + const std::string& kernel_name = shader.kernel_name; + // Checked before the plain "conv2d" test below, which "conv2d_dw" and + // "conv2d_pw" would otherwise match too. + if (kernel_name.find("conv2d_dw") != std::string::npos) { + return Conv2dMethod::Depthwise; + } + if (kernel_name.find("conv2d_pw") != std::string::npos) { + return Conv2dMethod::Pointwise; + } + if (kernel_name.find("conv_transpose2d") != std::string::npos) { + return Conv2dMethod::Transposed; + } + const auto& weight_sizes = graph->get_tref(weight_data)->sizes; + if (weight_sizes.at(2) == 1 && weight_sizes.at(3) == 1) { + return Conv2dMethod::Pointwise; + } + return Conv2dMethod::SlidingWindow; +} + // Custom global workgroup size function for conv2d GlobalWorkGrid conv2d_gwg( ComputeGraph* graph, @@ -359,23 +391,8 @@ GlobalWorkGrid conv2d_gwg( const ValueRef out = args.at(0).refs.at(0); const ValueRef weight_data = resize_args.at(0); - // Determine method from shader name - Conv2dMethod method; - if (shader.kernel_name.find("conv2d_pw") != std::string::npos || - (shader.kernel_name.find("conv2d") != std::string::npos && - shader.kernel_name.find("conv_transpose2d") == std::string::npos)) { - // Check if it's pointwise by examining weight sizes - const auto& weight_sizes = graph->get_tref(weight_data)->sizes; - if (weight_sizes.at(2) == 1 && weight_sizes.at(3) == 1) { - method = Conv2dMethod::Pointwise; - } else { - method = Conv2dMethod::SlidingWindow; - } - } else if (shader.kernel_name.find("conv_transpose2d") != std::string::npos) { - method = Conv2dMethod::Transposed; - } else { - method = Conv2dMethod::SlidingWindow; - } + const Conv2dMethod method = + infer_conv2d_method_from_shader(graph, shader, weight_data); // Determine stride_equals_dilation from shader name bool stride_equals_dilation = @@ -404,17 +421,10 @@ LocalWorkGroup conv2d_lwg( const std::vector& args, const std::vector& resize_args) { (void)args; - (void)resize_args; - // Determine method from shader name - Conv2dMethod method; - if (shader.kernel_name.find("conv2d_pw") != std::string::npos || - (shader.kernel_name.find("conv2d") != std::string::npos && - shader.kernel_name.find("conv_transpose2d") == std::string::npos)) { - method = Conv2dMethod::Pointwise; - } else { - method = Conv2dMethod::SlidingWindow; - } + const ValueRef weight_data = resize_args.at(0); + const Conv2dMethod method = + infer_conv2d_method_from_shader(graph, shader, weight_data); if (method == Conv2dMethod::Pointwise) { uint32_t lwg_y = 1; diff --git a/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp b/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp index 61ba9349b45..80f7b505feb 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp @@ -112,7 +112,7 @@ void add_embedding_legacy_node( // Resize Args {}, // Resizing Logic - nullptr)); + resize_embedding_node)); } void embedding(ComputeGraph& graph, const std::vector& args) { diff --git a/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp b/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp index 20ab76813a7..224cc427451 100644 --- a/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp @@ -28,6 +28,35 @@ void check_index_select_args( VK_CHECK_COND(graph.packed_dim_of(out) == WHCN::kChannelsDim); } +// index_select replaces the selected dim with as many entries as the index +// tensor holds and leaves every other dim alone. +std::vector index_select_out_sizes( + ComputeGraph* graph, + const ValueRef in, + const ValueRef idx, + const DimIndex dim_idx) { + std::vector out_sizes = graph->sizes_of(in); + const int64_t ndim = static_cast(out_sizes.size()); + // dim_idx is a negative index counted from the innermost dim. + const int64_t dim = ndim + dim_idx; + VK_CHECK_COND(dim >= 0 && dim < ndim); + out_sizes.at(dim) = graph->numel_of(idx); + return out_sizes; +} + +void resize_index_select_channel_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + (void)resize_args; + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); + const ValueRef idx = args.at(1).refs.at(1); + + graph->virtual_resize( + out, index_select_out_sizes(graph, in, idx, kChannel4D)); +} + void add_index_select_channel_node( ComputeGraph& graph, ValueRef in, @@ -53,32 +82,43 @@ void add_index_select_channel_node( // Resize Args {}, // Resizing Logic - nullptr)); + resize_index_select_channel_node)); } struct IndexSelectParams final { int32_t gpu_dim; - int32_t stride; }; -IndexSelectParams create_index_select_params( - ComputeGraph& graph, - const int64_t dim_idx, - const ValueRef in) { +IndexSelectParams create_index_select_params(const int64_t dim_idx) { if (dim_idx == kWidth4D) { - return {0, 1}; + return {0}; } else if (dim_idx == kHeight4D) { - return {1, 1}; + return {1}; } else if (dim_idx == kBatch4D) { - const std::vector in_sizes = graph.sizes_of(in); - int64_t n_channels = dim_at(in_sizes, kChannel4D); - int64_t stride = utils::div_up_4(n_channels); - return {2, static_cast(stride)}; + // The batch axis shares the z axis with the channels, so the shader steps + // over one batch in units of channel texels. That stride is derived from + // the channel count, which a resize can change, so the shader reads it out + // of in_sizes rather than taking a value frozen at build time. + return {2}; } else { VK_THROW("Unexpected dim_idx!"); } } +void resize_index_select_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); + const ValueRef idx = args.at(1).refs.at(1); + + const DimIndex dim_idx = + static_cast(graph->extract_scalar(resize_args.at(0))); + + graph->virtual_resize(out, index_select_out_sizes(graph, in, idx, dim_idx)); +} + void add_index_select_node( ComputeGraph& graph, ValueRef in, @@ -87,7 +127,7 @@ void add_index_select_node( ValueRef out) { check_index_select_args(graph, in, idx, out); - IndexSelectParams params = create_index_select_params(graph, dim_idx, in); + IndexSelectParams params = create_index_select_params(dim_idx); std::string kernel_name = "index_select"; kernel_name.reserve(kShaderNameReserve); @@ -99,15 +139,17 @@ void add_index_select_node( default_pick_gwg, default_pick_lwg, {{out, vkapi::kWrite}, {{in, idx}, vkapi::kRead}}, - {graph.sizes_ubo(out), graph.create_params_buffer(params)}, + {graph.sizes_ubo(out), + graph.sizes_ubo(in), + graph.create_params_buffer(params)}, // Push Constants {}, // Specialization Constants {}, // Resize Args - {}, + {graph.get_or_add_value_for_int(dim_idx)}, // Resizing Logic - nullptr)); + resize_index_select_node)); } int64_t get_dim_idx(ComputeGraph& graph, ValueRef in, ValueRef dim_ref) { diff --git a/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp b/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp index ddd8e8994b1..f490f60f75c 100644 --- a/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp @@ -9,6 +9,7 @@ #include #include +#include #include @@ -18,19 +19,31 @@ void resize_index_tensor_node( ComputeGraph* graph, const std::vector& args, const std::vector& resize_args) { - (void)resize_args; const ValueRef out = args.at(0).refs.at(0); const ValueRef self = args.at(1).refs.at(0); const ValueRef index = args.at(1).refs.at(1); - // aten.index.Tensor with a single index tensor gathers along dim 0, so - // out.sizes = index.sizes ++ self.sizes[1:] - // Using the index's sizes alone is only correct when self is 1-D; for any - // higher-rank self it also changes the tensor's RANK, which virtual_resize - // rejects outright ("new sizes cannot modify the dimensionality"). + int64_t index_dim = -1; + { + const ValueListPtr indices = graph->get_value_list(resize_args.at(0)); + for (size_t dim = 0; dim < indices->size(); ++dim) { + if (!graph->val_is_none(indices->at(dim))) { + index_dim = utils::safe_downcast(dim); + break; + } + } + } + VK_CHECK_COND(index_dim >= 0, "index.Tensor: an index tensor is required"); + const std::vector self_sizes = graph->sizes_of(self); - std::vector out_sizes = graph->sizes_of(index); - out_sizes.insert(out_sizes.end(), self_sizes.begin() + 1, self_sizes.end()); + const std::vector index_sizes = graph->sizes_of(index); + std::vector out_sizes; + out_sizes.reserve(self_sizes.size() + index_sizes.size() - 1); + out_sizes.insert( + out_sizes.end(), self_sizes.begin(), self_sizes.begin() + index_dim); + out_sizes.insert(out_sizes.end(), index_sizes.begin(), index_sizes.end()); + out_sizes.insert( + out_sizes.end(), self_sizes.begin() + index_dim + 1, self_sizes.end()); graph->virtual_resize(out, out_sizes); } @@ -39,14 +52,24 @@ void add_index_tensor_node( ComputeGraph& graph, const ValueRef self, const ValueRef index, + const int64_t index_dim, + const ValueRef indices_list_ref, const ValueRef out) { std::string kernel_name = "index_tensor"; kernel_name.reserve(kShaderNameReserve); add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); + if (graph.is_buffer_storage(out)) { + add_storage_type_suffix(kernel_name, graph.storage_type_of(index)); + } add_dtype_suffix(kernel_name, graph.dtype_of(out)); vkapi::ParamsBindList param_ubos = { graph.meta_ubo(out), graph.meta_ubo(self), graph.meta_ubo(index)}; + const utils::ivec2 index_params = { + utils::safe_downcast(graph.dim_of(self) - 1 - index_dim), + utils::safe_downcast(graph.dim_of(index))}; + std::vector push_constants = { + PushConstantDataInfo(&index_params, sizeof(index_params))}; graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, @@ -58,11 +81,13 @@ void add_index_tensor_node( // Shader params buffers param_ubos, // Push Constants - {}, + push_constants, // Specialization Constants - {graph.hashed_layout_of(out), graph.hashed_layout_of(self)}, + {graph.hashed_layout_of(out), + graph.hashed_layout_of(self), + graph.hashed_layout_of(index)}, // Resize Args - {}, + {indices_list_ref}, // Resizing Logic resize_index_tensor_node)); } @@ -72,14 +97,27 @@ void index_tensor(ComputeGraph& graph, const std::vector& args) { ValueRef indices_list_ref = args[1]; ValueRef out = args[2]; - ValueListPtr indices_list = graph.get_value_list(indices_list_ref); + ValueRef index = -1; + int64_t index_dim = -1; + { + const ValueListPtr indices_list = graph.get_value_list(indices_list_ref); + for (size_t dim = 0; dim < indices_list->size(); ++dim) { + const ValueRef candidate = indices_list->at(dim); + if (graph.val_is_none(candidate)) { + continue; + } + VK_CHECK_COND( + index_dim < 0, "index.Tensor: only one index tensor is supported"); + index = candidate; + index_dim = utils::safe_downcast(dim); + } + } + VK_CHECK_COND(index_dim >= 0, "index.Tensor: an index tensor is required"); VK_CHECK_COND( - indices_list->size() == 1, - "index.Tensor: only one index tensor is supported"); - - ValueRef index = indices_list->at(0); + index_dim < graph.dim_of(self), + "index.Tensor: index dimension is invalid"); - add_index_tensor_node(graph, self, index, out); + add_index_tensor_node(graph, self, index, index_dim, indices_list_ref, out); } REGISTER_OPERATORS { diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp index 70171a04820..75f65d7a2a0 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp @@ -7,6 +7,7 @@ */ #include +#include #include @@ -42,52 +43,6 @@ bool q8ta_conv2d_check_4w4c_packed_dim_info(const api::PackedDimInfo& info) { info.outer_packed_dim_block_size == 4; } -namespace { - -uint64_t q8ta_conv2d_im2col_scratch_limit(ComputeGraph& graph) { - constexpr uint64_t kMaxBatchedIm2ColScratchBytes = 32ULL * 1024ULL * 1024ULL; - const uint64_t device_scratch_limit = - graph.context()->adapter_ptr()->max_buffer_numel(); - return device_scratch_limit < kMaxBatchedIm2ColScratchBytes - ? device_scratch_limit - : kMaxBatchedIm2ColScratchBytes; -} - -bool should_use_q8ta_conv2d_im2col( - ComputeGraph& graph, - const int64_t batch, - const int64_t groups, - const int64_t in_channels_per_group, - const int64_t flattened_kernel_size, - const int64_t out_height, - const int64_t out_width) { - const bool im2col_eligible = in_channels_per_group % 4 == 0; - if (!im2col_eligible) { - return false; - } - - const int64_t spatial_out = out_height * out_width; - if (batch > 1) { - constexpr int64_t kMinFlattenedKernelSize = 1024; - constexpr int64_t kMaxSpatialOutput = 64; - const uint64_t scratch_bytes = static_cast(batch) * - static_cast(flattened_kernel_size) * - static_cast(out_height) * - static_cast(utils::align_up_4(out_width)); - return groups == 1 && flattened_kernel_size >= kMinFlattenedKernelSize && - spatial_out <= kMaxSpatialOutput && - scratch_bytes <= q8ta_conv2d_im2col_scratch_limit(graph); - } - - if (graph.device_is_mali()) { - return true; - } - - return groups == 1 && (in_channels_per_group >= 32 || spatial_out <= 4096); -} - -} // namespace - // // Workgroup size selection functions // @@ -111,6 +66,7 @@ GlobalWorkGrid pick_q8ta_conv2d_gwg( (void)shader; (void)resize_args; + VK_CHECK_COND(graph != nullptr); const ValueRef output = args.at(0).refs.at(0); const uint32_t W = graph->size_at(-1, output); @@ -133,6 +89,7 @@ GlobalWorkGrid pick_q8ta_conv2d_gwg( * tensor dimensions. Uses experimentation results: * - {4, 2, 8} for medium tensors: +57% improvement on 81x81 * - {8, 1, 8} for very large tensors: best baseline performance + * - {2, 1, 32} or {4, 1, 16} for narrow output widths * - {64, 1, 1} for narrow channel dimensions: minimize inactive invocations */ LocalWorkGroup pick_q8ta_conv2d_lwg( @@ -144,14 +101,13 @@ LocalWorkGroup pick_q8ta_conv2d_lwg( (void)shader; (void)resize_args; + VK_CHECK_COND(graph != nullptr); const ValueRef output = args.at(0).refs.at(0); - - // Get actual tensor dimensions for adaptive sizing - const uint32_t H = graph->size_at(-2, output); + const uint32_t output_height = graph->size_at(-2, output); // For very large tensors (H >= 100 and large x/z), use {8, 1, 8} // This configuration performed best for 128x128 tensors in experiments - if (H >= 100 && gwg[0u] >= 24 && gwg[2u] >= 24) { + if (output_height >= 100 && gwg[0u] >= 24 && gwg[2u] >= 24) { return LocalWorkGroup(8u, 1u, 8u); } @@ -161,6 +117,16 @@ LocalWorkGroup pick_q8ta_conv2d_lwg( return LocalWorkGroup(4u, 2u, 8u); } + if (gwg[0u] == 2u && gwg[2u] >= 32u) { + return LocalWorkGroup(2u, 1u, 32u); + } + + // LWG x oversubscribes the 3 global groups here; safe only because the + // shader early-returns out-of-bounds invocations. + if (gwg[0u] == 3u && gwg[2u] >= 16u) { + return LocalWorkGroup(4u, 1u, 16u); + } + // For tensors with sufficient x and z dimensions, use square configuration if (gwg[0u] >= 6 && gwg[2u] >= 6) { return LocalWorkGroup(8u, 1u, 8u); @@ -520,27 +486,38 @@ void q8ta_conv2d(ComputeGraph& graph, const std::vector& args) { const ValueRef output = args.at(15); const int64_t groups = graph.extract_scalar(groups_ref); + // Valid models always carry groups >= 1; fail fast on corrupt input + // instead of dividing channel counts by zero downstream (both this + // dispatcher and q8ta_conv2d_general divide by groups). + VK_CHECK_COND(groups > 0, "q8ta_conv2d requires groups >= 1"); const int64_t in_channels = graph.size_at(-3, input); const int64_t in_channels_per_group = in_channels / groups; const int64_t batch = graph.size_at(-4, input); const int64_t H_out = graph.size_at(-2, output); const int64_t W_out = graph.size_at(-1, output); - int64_t flattened_kernel_size; + const int64_t out_channels = graph.size_at(-3, output); + int64_t kernel_height; + int64_t kernel_width; { const auto kernel_size = graph.get_int_list(kernel_size_ref); - flattened_kernel_size = utils::align_up_4( - in_channels_per_group * kernel_size->at(0) * kernel_size->at(1)); + kernel_height = kernel_size->at(0); + kernel_width = kernel_size->at(1); } - const bool use_im2col = should_use_q8ta_conv2d_im2col( - graph, + const bool use_im2col = should_use_q8ta_conv2d_im2col({ + graph.device_is_mali(), + graph.can_use_int8_dot_product(), + static_cast(graph.max_buffer_numel()), batch, groups, in_channels_per_group, - flattened_kernel_size, + out_channels, + kernel_height, + kernel_width, H_out, - W_out); + W_out, + }); if (use_im2col) { q8ta_conv2d_im2col(graph, args); diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h index 5d16cb3b78c..e48d1e32a08 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h @@ -14,6 +14,34 @@ namespace vkcompute { +inline constexpr int64_t kQ8taConv2dIm2ColScratchBudgetBytes = 16 * 1024 * 1024; +inline constexpr int64_t kQ8taConv2dMaxRowsPerTile = 65535; + +struct Q8taConv2dStreamPlan final { + int64_t aligned_out_width; + int64_t rows_per_tile; + int64_t num_tiles; + int64_t scratch_bytes; + bool feasible; +}; + +Q8taConv2dStreamPlan make_q8ta_conv2d_stream_plan( + int64_t batch, + int64_t flattened_kernel_size, + int64_t out_height, + int64_t out_width, + int64_t scratch_budget_bytes); + +// max_buffer_bytes is Adapter::max_buffer_numel(), which returns +// maxStorageBufferRange in bytes (not elements) — directly comparable with +// the byte-denominated scratch budget. +Q8taConv2dStreamPlan make_q8ta_conv2d_stream_plan_for_device( + int64_t batch, + int64_t flattened_kernel_size, + int64_t out_height, + int64_t out_width, + uint64_t max_buffer_bytes); + enum class ActivationType : uint32_t { kNone = 0, kRelu = 1, @@ -111,6 +139,7 @@ void add_q8ta_conv2d_node( void add_q8ta_conv2d_pw_node( ComputeGraph& graph, + const bool use_unsigned_dot, const ValueRef packed_int8_input, const ValueRef input_scale, const ValueRef input_zp, @@ -128,7 +157,23 @@ void add_q8ta_conv2d_pw_node( const ValueRef kernel_size = kDummyValueRef, const ValueRef stride = kDummyValueRef, const ValueRef padding = kDummyValueRef, - const ValueRef dilation = kDummyValueRef); + const ValueRef dilation = kDummyValueRef, + const bool is_im2col = false, + const ValueRef stream_row_offset_ref = kDummyValueRef, + const ValueRef max_im2col_rows_ref = kDummyValueRef); + +constexpr int64_t kMaxUnsignedDotAccumulatorBytes = 33025; + +bool can_use_unsigned_pw_dot( + const vkapi::Adapter& adapter, + int64_t k_per_group); + +void q8ta_conv2d_pw_impl( + ComputeGraph& graph, + bool use_unsigned_dot, + const std::vector& args); + +void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args); std::vector calculate_q8ta_im2col_sizes( ComputeGraph* graph, @@ -147,10 +192,21 @@ void add_q8ta_im2col_node( const ValueRef groups, const ValueRef packed_int8_output, const ValueRef packed_int8_im2col, - const int32_t zp); + const int32_t zp, + const ValueRef stream_row_offset_ref, + const ValueRef max_im2col_rows_ref = kDummyValueRef); void q8ta_conv2d_im2col(ComputeGraph& graph, const std::vector& args); +void q8ta_conv2d_im2col_impl( + ComputeGraph& graph, + bool use_unsigned_dot, + const std::vector& args); + +void q8ta_conv2d_general( + ComputeGraph& graph, + const std::vector& args); + // Transposed convolution void q8ta_conv2d_transposed( diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp index 182e8d684d2..a504d9ab3f6 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp @@ -29,6 +29,7 @@ GlobalWorkGrid pick_q8ta_conv2d_dw_gwg( (void)shader; (void)resize_args; + VK_CHECK_COND(graph != nullptr); const ValueRef output = args.at(0).refs.at(0); const uint32_t W = graph->size_at(-1, output); @@ -46,6 +47,15 @@ GlobalWorkGrid pick_q8ta_conv2d_dw_gwg( kTiledWorkGrid); } +/** + * Picks a local workgroup size for q8ta_conv2d_dw with adaptive sizing based + * on tensor dimensions. Uses experimentation results: + * - {2, 1, 32} or {4, 1, 16} for narrow output widths + * + * Unlike the regular conv picker, there is no medium-tensor branch shadowing + * gwg[0] == 4, so the second narrow branch matches 3..4 (the conv picker's + * {4, 2, 8} branch claims gwg[0] >= 4 first, leaving only == 3 reachable). + */ LocalWorkGroup pick_q8ta_conv2d_dw_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, @@ -57,6 +67,16 @@ LocalWorkGroup pick_q8ta_conv2d_dw_lwg( (void)args; (void)resize_args; + if (gwg[0u] == 2u && gwg[2u] >= 32u) { + return LocalWorkGroup(2u, 1u, 32u); + } + + // LWG x oversubscribes when gwg[0] is 3; safe only because the shader + // early-returns out-of-bounds invocations. + if (gwg[0u] >= 3u && gwg[0u] <= 4u && gwg[2u] >= 16u) { + return LocalWorkGroup(4u, 1u, 16u); + } + // Some inactive invocations are okay; set 6 as the threshold to use the // a square wg size. if (gwg[0u] >= 6 && gwg[2u] >= 6) { diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp index e93723c5125..062e5ea1d8d 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp @@ -16,8 +16,68 @@ #include #include +#include +#include + namespace vkcompute { +Q8taConv2dStreamPlan make_q8ta_conv2d_stream_plan( + const int64_t batch, + const int64_t flattened_kernel_size, + const int64_t out_height, + const int64_t out_width, + const int64_t scratch_budget_bytes) { + Q8taConv2dStreamPlan plan{}; + if (batch <= 0 || flattened_kernel_size <= 0 || out_height <= 0 || + out_width <= 0 || scratch_budget_bytes <= 0) { + return plan; + } + + constexpr int64_t kAlignment = 4; + if (out_width > std::numeric_limits::max() - (kAlignment - 1)) { + return plan; + } + plan.aligned_out_width = + (out_width + kAlignment - 1) / kAlignment * kAlignment; + if (flattened_kernel_size > + std::numeric_limits::max() / plan.aligned_out_width) { + return plan; + } + const int64_t bytes_per_row = flattened_kernel_size * plan.aligned_out_width; + if (bytes_per_row > scratch_budget_bytes || + batch > std::numeric_limits::max() / out_height) { + return plan; + } + + const int64_t total_rows = batch * out_height; + if (total_rows > std::numeric_limits::max()) { + return plan; + } + const int64_t max_rows_per_tile = std::min( + {total_rows, + scratch_budget_bytes / bytes_per_row, + kQ8taConv2dMaxRowsPerTile}); + plan.num_tiles = total_rows / max_rows_per_tile + + static_cast(total_rows % max_rows_per_tile != 0); + plan.rows_per_tile = total_rows / plan.num_tiles + + static_cast(total_rows % plan.num_tiles != 0); + plan.scratch_bytes = plan.rows_per_tile * bytes_per_row; + plan.feasible = true; + return plan; +} + +Q8taConv2dStreamPlan make_q8ta_conv2d_stream_plan_for_device( + const int64_t batch, + const int64_t flattened_kernel_size, + const int64_t out_height, + const int64_t out_width, + const uint64_t max_buffer_bytes) { + const int64_t scratch_budget = static_cast(std::min( + kQ8taConv2dIm2ColScratchBudgetBytes, max_buffer_bytes)); + return make_q8ta_conv2d_stream_plan( + batch, flattened_kernel_size, out_height, out_width, scratch_budget); +} + // // Shader dispatch utilities // @@ -28,21 +88,41 @@ GlobalWorkGrid pick_q8ta_im2col_gwg( const std::vector& args, const std::vector& resize_args) { (void)shader; - (void)resize_args; - + VK_CHECK_COND(graph != nullptr); const ValueRef im2col_output = args.at(0).refs.at(0); - - const uint32_t N = graph->size_at(-4, im2col_output); const uint32_t K = graph->size_at(-3, im2col_output); - const uint32_t H = graph->size_at(-2, im2col_output); + const uint32_t rows_per_tile = graph->size_at(-2, im2col_output); const uint32_t W = graph->size_at(-1, im2col_output); + const ValueRef input = resize_args.at(0); + const ValueRef kernel_size = resize_args.at(1); + const ValueRef stride = resize_args.at(2); + const ValueRef padding = resize_args.at(3); + const ValueRef dilation = resize_args.at(4); + const int64_t row_offset = graph->extract_scalar(resize_args.at(6)); + + const std::vector input_sizes = graph->sizes_of(input); + const int64_t batch = utils::val_at(-4, input_sizes); + const std::vector out_hw = calc_out_sizes_hw( + *graph, + input_sizes, + kernel_size, + /*kernel_size_only=*/true, + {stride, padding, dilation, dilation}, + /*transposed=*/false); + const int64_t total_rows = batch * out_hw.at(0); + if (row_offset >= total_rows) { + return graph->create_linear_gwg(0u); + } + const uint32_t live_rows = utils::safe_downcast( + std::min(rows_per_tile, total_rows - row_offset)); + const uint32_t K4 = utils::div_up_4(K); const uint32_t W4 = utils::div_up_4(W); // Each thread handles one 4x4 block in the output - return graph->create_linear_gwg( - utils::safe_downcast(static_cast(K4) * W4 * H * N)); + return graph->create_linear_gwg(utils::safe_downcast( + static_cast(K4) * W4 * live_rows)); } LocalWorkGroup pick_q8ta_im2col_lwg( @@ -102,19 +182,18 @@ std::vector calculate_q8ta_im2col_sizes( // Resize // -// resize_args = { input, kernel_size, stride, padding, dilation, groups } +// resize_args = { input, kernel_size, stride, padding, dilation, groups, +// row_offset, max_im2col_rows } // -// The im2col scratch tensor is [N, K, H_out, align_up_4(W_out)] where K (the -// flattened conv window, channel/kernel-derived) is shape-independent and -// H_out/W_out are the conv output spatial dims. The downstream PW GEMM that -// consumes this scratch is resized separately (it preserves H/W). Without this, -// the scratch freezes at the build-time upper bound and feeds garbage rows into -// the GEMM. Recompute H_out/W_out from the CURRENT input (NOT the conv output -// tensor, which may itself still be frozen at this point in the resize order). +// The scratch tensor is [1, K, rows_per_tile, align_up_4(W_out)]. K and +// rows_per_tile are fixed; only W_out tracks the current input shape. +// Batch/height growth past max im2col rows has no dispatches, +// so fail fast instead of leaving outputs stale. void resize_q8ta_im2col_node( ComputeGraph* graph, const std::vector& args, const std::vector& resize_args) { + VK_CHECK_COND(graph != nullptr); const ValueRef im2col_out = args.at(0).refs.at(0); const ValueRef in = resize_args.at(0); const ValueRef kernel_size = resize_args.at(1); @@ -122,11 +201,11 @@ void resize_q8ta_im2col_node( const ValueRef padding = resize_args.at(3); const ValueRef dilation = resize_args.at(4); const ValueRef groups = resize_args.at(5); + const int64_t max_im2col_rows = + graph->extract_scalar(resize_args.at(7)); const std::vector in_sizes = graph->sizes_of(in); - const int64_t batch = utils::val_at(-4, in_sizes); - - // Conv output H/W from the current input. + // Conv output width from the current input. const std::vector out_hw = calc_out_sizes_hw( *graph, in_sizes, @@ -134,7 +213,6 @@ void resize_q8ta_im2col_node( /*kernel_size_only=*/true, {stride, padding, dilation, dilation}, /*transposed=*/false); - const int64_t out_height = out_hw.at(0); const int64_t out_width = out_hw.at(1); // K (flattened conv window) is shape-independent — recompute from channels + @@ -149,7 +227,15 @@ void resize_q8ta_im2col_node( const int64_t K = flattened_kernel_len * groups_val; const int64_t W = utils::align_up_4(out_width); - graph->virtual_resize(im2col_out, {batch, K, out_height, W}); + const int64_t rows_per_tile = graph->size_at(-2, im2col_out); + + const int64_t batch = utils::val_at(-4, in_sizes); + const int64_t out_height = out_hw.at(0); + VK_CHECK_COND( + batch * out_height <= max_im2col_rows, + "q8ta im2col resize grew past max im2col rows"); + + graph->virtual_resize(im2col_out, {1, K, rows_per_tile, W}); } // @@ -166,7 +252,9 @@ void add_q8ta_im2col_node( const ValueRef groups, const ValueRef packed_int8_output, const ValueRef packed_int8_im2col, - const int32_t zp) { + const int32_t zp, + const ValueRef stream_row_offset_ref, + const ValueRef max_im2col_rows_ref) { // Validate packed dim info for input and output tensors VK_CHECK_COND(q8ta_conv2d_check_packed_dim_info( graph.packed_dim_info_of(packed_int8_input))); @@ -195,8 +283,14 @@ void add_q8ta_im2col_node( graph.buffer_meta_ubo(packed_int8_input), graph.create_params_buffer(conv_params)}; + VK_CHECK_COND(stream_row_offset_ref != kDummyValueRef); + VK_CHECK_COND(max_im2col_rows_ref != kDummyValueRef); + const int32_t stream_row_offset = utils::safe_downcast( + graph.extract_scalar(stream_row_offset_ref)); + std::vector push_constants = { PushConstantDataInfo(&zp, sizeof(zp)), + PushConstantDataInfo(&stream_row_offset, sizeof(stream_row_offset)), }; // Build spec constants: apply_bias + layout constants (for generic shader) @@ -212,6 +306,19 @@ void add_q8ta_im2col_node( // spec_constants.append(graph.hashed_layout_of(packed_int8_im2col)); // } + // resize_args = { input, kernel_size, stride, padding, dilation, groups, + // row_offset, max_im2col_rows }. The grid picker reads the + // row offset at index 6; append-only. + std::vector resize_args = { + packed_int8_input, + kernel_size, + stride, + padding, + dilation, + groups, + stream_row_offset_ref, + max_im2col_rows_ref}; + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), @@ -225,10 +332,7 @@ void add_q8ta_im2col_node( push_constants, // Specialization Constants spec_constants, - // Resize args: { input, kernel_size, stride, padding, dilation, groups } - {packed_int8_input, kernel_size, stride, padding, dilation, groups}, - // Resizing Logic: recompute the im2col scratch dims from the current - // input + resize_args, resize_q8ta_im2col_node)); } @@ -236,8 +340,9 @@ void add_q8ta_im2col_node( // High level operator impl // -void q8ta_conv2d_im2col( +void q8ta_conv2d_im2col_impl( ComputeGraph& graph, + const bool use_unsigned_dot, const std::vector& args) { int32_t idx = 0; const ValueRef packed_int8_input = args.at(idx++); @@ -257,11 +362,30 @@ void q8ta_conv2d_im2col( const ValueRef activation = args.at(idx++); const ValueRef packed_int8_output = args.at(idx++); + const std::vector full_im2col_sizes = calculate_q8ta_im2col_sizes( + &graph, packed_int8_input, packed_int8_output, kernel_size, groups); + const Q8taConv2dStreamPlan stream_plan = + make_q8ta_conv2d_stream_plan_for_device( + full_im2col_sizes.at(0), + full_im2col_sizes.at(1), + full_im2col_sizes.at(2), + full_im2col_sizes.at(3), + graph.max_buffer_numel()); + if (!stream_plan.feasible) { + q8ta_conv2d_general(graph, args); + return; + } + VK_CHECK_COND( + !use_unsigned_dot || + graph.size_at(-1, weight_data) <= + kMaxUnsignedDotAccumulatorBytes, + "Unsigned q8ta im2col convolution exceeds the accumulator bound"); + QuantizationConfig weight_quant_config(8, kPerChannel, {}); // Prepack weight using linear weight packing (for im2col approach) - ValueRef packed_weight = - prepack_quantized_linear_weight(graph, weight_quant_config, weight_data); + ValueRef packed_weight = prepack_quantized_linear_weight( + graph, weight_quant_config, weight_data, use_unsigned_dot); ValueRef packed_weight_sums = prepack_standard( graph, weight_sums_data, utils::kBuffer, utils::kWidthPacked); @@ -286,11 +410,15 @@ void q8ta_conv2d_im2col( uint32_t activation_type_val = static_cast( activation_type_from_string(graph.extract_string(activation))); - // Calculate im2col output sizes - std::vector im2col_sizes = calculate_q8ta_im2col_sizes( - &graph, packed_int8_input, packed_int8_output, kernel_size, groups); + // One fixed-size scratch buffer is reused across all row tiles; the full + // fit is a single tile. Interleaved write/read dispatches insert the + // barrier before the next tile overwrites it. + const std::vector im2col_sizes = { + 1, + full_im2col_sizes.at(1), + stream_plan.rows_per_tile, + stream_plan.aligned_out_width}; - // Create temporary tensor for im2col output (4W4C layout) TmpTensor packed_int8_im2col( &graph, im2col_sizes, @@ -299,45 +427,69 @@ void q8ta_conv2d_im2col( utils::kPackedInt8_4W4C); int32_t zp = graph.extract_scalar(input_zp); - - // Step 1: Perform im2col transformation - add_q8ta_im2col_node( - graph, - packed_int8_input, - kernel_size, - stride, - padding, - dilation, - groups, - packed_int8_output, - packed_int8_im2col, - zp); - - // Step 2: Perform pointwise convolution on the im2col result const int32_t groups_val = graph.extract_scalar(groups); - add_q8ta_conv2d_pw_node( - graph, - packed_int8_im2col, - input_scale, - input_zp, - packed_weight, - packed_weight_sums, - packed_weight_scales, - output_scale, - output_zp, - bias_data, - packed_bias, - activation_type_val, - packed_int8_output, - groups_val, - // Original activation + conv geometry so the PW output H/W is recomputed - // from the true conv result, not the width-padded im2col scratch. - packed_int8_input, - kernel_size, - stride, - padding, - dilation); + // Row tiles are fixed at build time: dynamic growth past max im2col rows + // has no dispatches, so each resize fails fast below instead of leaving + // outputs stale. Shrinkage only ever lowers the total below this bound. + const ValueRef max_im2col_rows_ref = graph.add_scalar( + stream_plan.num_tiles * stream_plan.rows_per_tile); + + for (int64_t tile = 0; tile < stream_plan.num_tiles; ++tile) { + const int64_t row_offset = tile * stream_plan.rows_per_tile; + // One row-offset scalar per tile feeds both nodes' push constants (via + // re-extraction) and resize args. + const ValueRef row_offset_ref = graph.add_scalar(row_offset); + + add_q8ta_im2col_node( + graph, + packed_int8_input, + kernel_size, + stride, + padding, + dilation, + groups, + packed_int8_output, + packed_int8_im2col, + zp, + row_offset_ref, + max_im2col_rows_ref); + + add_q8ta_conv2d_pw_node( + graph, + use_unsigned_dot, + packed_int8_im2col, + input_scale, + input_zp, + packed_weight, + packed_weight_sums, + packed_weight_scales, + output_scale, + output_zp, + bias_data, + packed_bias, + activation_type_val, + packed_int8_output, + groups_val, + packed_int8_input, + kernel_size, + stride, + padding, + dilation, + /*is_im2col=*/true, + row_offset_ref, + max_im2col_rows_ref); + } +} + +void q8ta_conv2d_im2col( + ComputeGraph& graph, + const std::vector& args) { + const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); + const ValueRef weight_data = args.at(3); + const int64_t k_per_group = graph.size_at(-1, weight_data); + const bool use_unsigned_dot = can_use_unsigned_pw_dot(*adapter, k_per_group); + q8ta_conv2d_im2col_impl(graph, use_unsigned_dot, args); } REGISTER_OPERATORS { diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp index ee234319e8c..454a115498b 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp @@ -14,19 +14,20 @@ #include #include +#include + namespace vkcompute { // // Shader dispatch utilities // -GlobalWorkGrid pick_q8ta_conv2d_pw_gwg( +GlobalWorkGrid pick_q8ta_conv2d_pw_gwg_impl( ComputeGraph* graph, - const vkapi::ShaderInfo& shader, const std::vector& args, - const std::vector& resize_args) { - (void)shader; - (void)resize_args; + const std::vector& resize_args, + const bool is_im2col_tile) { + VK_CHECK_COND(graph != nullptr); const ValueRef output = args.at(0).refs.at(0); @@ -35,23 +36,50 @@ GlobalWorkGrid pick_q8ta_conv2d_pw_gwg( const uint32_t C = graph->size_at(-3, output); const uint32_t N = graph->size_at(-4, output); - // Each thread covers a 4-width x 4-channel output block. - // Tile constants must match TILE_M4 / TILE_N4 in q8ta_conv2d_pw.glsl. - constexpr uint32_t TILE_N4 = 1; - constexpr uint32_t TILE_M4 = 1; - + // Each thread covers a 4-width x 4-channel output block, matching TILE_M4 / + // TILE_N4 (= 1) in q8ta_conv2d_pw.glsl. const uint32_t C4 = utils::div_up_4(C); const uint32_t W4 = utils::div_up_4(W); - // Global workgroup size: - // x = output channels / (TILE_N4 * 4) = C4 / TILE_N4 = C4 - // y = width / (TILE_M4 * 4) = W4 / TILE_M4 = W4 - // z = height * batch - return GlobalWorkGrid( - {utils::div_up(C4, TILE_N4), - utils::div_up(W4, TILE_M4), - utils::safe_downcast(static_cast(H) * N)}, - kTiledWorkGrid); + uint32_t z; + if (is_im2col_tile) { + // The bound input is one [1, K, rows, W] scratch tile; the grid covers + // the live tile rows after the tile offset. + const ValueRef input = args.at(1).refs.at(0); + const uint32_t rows_per_tile = graph->size_at(-2, input); + const int64_t row_offset = + graph->extract_scalar(resize_args.at(5)); + const int64_t total_rows = static_cast(N) * H; + if (row_offset >= total_rows) { + return GlobalWorkGrid({0u, 0u, 0u}, kTiledWorkGrid); + } + z = utils::safe_downcast( + std::min(rows_per_tile, total_rows - row_offset)); + } else { + // The bound input is the batched activation; the grid covers every + // output row. + z = utils::safe_downcast(static_cast(H) * N); + } + return GlobalWorkGrid({C4, W4, z}, kTiledWorkGrid); +} + +GlobalWorkGrid pick_q8ta_conv2d_pw_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)resize_args; + return pick_q8ta_conv2d_pw_gwg_impl(graph, args, resize_args, false); +} + +GlobalWorkGrid pick_q8ta_conv2d_pw_streaming_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + return pick_q8ta_conv2d_pw_gwg_impl(graph, args, resize_args, true); } LocalWorkGroup pick_q8ta_conv2d_pw_lwg( @@ -118,7 +146,8 @@ ValueRef prepack_quantized_conv2d_pw_weight( const QuantizationConfig& weight_quant_config, const ValueRef weight_data, const ValueRef input, - const ValueRef output) { + const ValueRef output, + const bool use_unsigned_dot) { VK_CHECK_COND(weight_quant_config.nbits == 8); VK_CHECK_COND(weight_quant_config.is_symmetric); @@ -149,7 +178,8 @@ ValueRef prepack_quantized_conv2d_pw_weight( std::vector packed_weight_sizes{output_height, output_width}; utils::StorageType storage_type = utils::kTexture2D; - uint32_t max_extent = graph.context()->adapter_ptr()->max_texture2d_dim(); + const uint32_t max_extent = + graph.context()->adapter_ptr()->max_texture2d_dim(); if (output_width > max_extent * 4 || output_height > max_extent) { storage_type = utils::kBuffer; } @@ -166,7 +196,8 @@ ValueRef prepack_quantized_conv2d_pw_weight( 1u}, kTiledWorkGrid); - std::string kernel_name = "pack_q8_conv2d_weights"; + std::string kernel_name = use_unsigned_dot ? "pack_q8_conv2d_weights_unsigned" + : "pack_q8_conv2d_weights"; add_storage_type_suffix(kernel_name, storage_type); graph.prepack_nodes().emplace_back(new PrepackNode( @@ -215,14 +246,16 @@ void resize_q8ta_conv2d_pw_node( graph->virtual_resize(out, new_sizes); } -// resize_args = { conv_input, kernel_size, stride, padding, dilation } +// resize_args = { conv_input, kernel_size, stride, padding, dilation, +// row_offset, max_im2col_rows }. The grid picker reads the row +// offset at index 5; append-only. // // im2col-path PW conv. Here the PW node's bound input is the im2col scratch -// tensor sized {K, H_out, align_up_4(W_out)} — its width is rounded up to a -// multiple of 4 for texel alignment, so it must NOT be used to size the output. -// Recompute the TRUE conv H_out/W_out from the ORIGINAL activation + conv -// geometry, exactly as resize_q8ta_conv2d_node does. N/C are shape-independent -// and stay as currently allocated. +// tensor sized {1, K, rows, align_up_4(W)} (one row tile) — its width is +// rounded up to a multiple of 4 for texel alignment, so it must NOT be used +// to size the output. +// Recompute the true conv N/H_out/W_out from the original activation + conv +// geometry, exactly as resize_q8ta_conv2d_node does. C is shape-independent. void resize_q8ta_conv2d_pw_im2col_node( ComputeGraph* graph, const std::vector& args, @@ -233,8 +266,13 @@ void resize_q8ta_conv2d_pw_im2col_node( const ValueRef stride = resize_args.at(2); const ValueRef padding = resize_args.at(3); const ValueRef dilation = resize_args.at(4); + // Row tiles are fixed at build time: fail fast on growth past max im2col + // rows instead of leaving outputs stale. + const int64_t max_im2col_rows = + graph->extract_scalar(resize_args.at(6)); const std::vector in_sizes = graph->sizes_of(conv_input); + const int64_t batch = utils::val_at(-4, in_sizes); const std::vector out_hw = calc_out_sizes_hw( *graph, @@ -244,8 +282,13 @@ void resize_q8ta_conv2d_pw_im2col_node( {stride, padding, dilation, dilation}, /*transposed=*/false); + VK_CHECK_COND( + batch * out_hw.at(0) <= max_im2col_rows, + "q8ta im2col resize grew past max im2col rows"); + std::vector new_sizes = graph->sizes_of(out); const size_t ndim = new_sizes.size(); + new_sizes.at(ndim - 4) = utils::val_at(-4, in_sizes); new_sizes.at(ndim - 2) = out_hw.at(0); new_sizes.at(ndim - 1) = out_hw.at(1); graph->virtual_resize(out, new_sizes); @@ -257,6 +300,7 @@ void resize_q8ta_conv2d_pw_im2col_node( void add_q8ta_conv2d_pw_node( ComputeGraph& graph, + const bool use_unsigned_dot, const ValueRef packed_int8_input, const ValueRef input_scale, const ValueRef input_zp, @@ -274,7 +318,10 @@ void add_q8ta_conv2d_pw_node( const ValueRef kernel_size, const ValueRef stride, const ValueRef padding, - const ValueRef dilation) { + const ValueRef dilation, + const bool is_im2col, + const ValueRef stream_row_offset_ref, + const ValueRef max_im2col_rows_ref) { VK_CHECK_COND(q8ta_conv2d_check_4w4c_packed_dim_info( graph.packed_dim_info_of(packed_int8_input))); VK_CHECK_COND(q8ta_conv2d_check_packed_dim_info( @@ -300,6 +347,16 @@ void add_q8ta_conv2d_pw_node( int32_t output_zp_val = graph.extract_scalar(output_zp); uint32_t apply_bias = graph.val_is_none(bias_data) ? 0u : 1u; + // The tile offset lives in the graph scalar; re-extract it here so the + // shader push constant and the grid picker read one value. Standalone + // dispatches have no tile and push 0. + int32_t stream_row_offset = 0; + if (is_im2col) { + VK_CHECK_COND(stream_row_offset_ref != kDummyValueRef); + VK_CHECK_COND(max_im2col_rows_ref != kDummyValueRef); + stream_row_offset = utils::safe_downcast( + graph.extract_scalar(stream_row_offset_ref)); + } std::vector push_constants = { PushConstantDataInfo(&input_scale_val, sizeof(input_scale_val)), PushConstantDataInfo(&input_zp_val, sizeof(input_zp_val)), @@ -307,12 +364,35 @@ void add_q8ta_conv2d_pw_node( PushConstantDataInfo(&output_zp_val, sizeof(output_zp_val)), PushConstantDataInfo(&K4_per_group, sizeof(K4_per_group)), PushConstantDataInfo(&OC4_per_group, sizeof(OC4_per_group)), + PushConstantDataInfo(&stream_row_offset, sizeof(stream_row_offset)), }; + // The im2col path consumes one flat scratch tile per dispatch; the + // standalone 1x1 path reads its batched activation input directly. The + // addressing is selected by spec constant, so both share one shader. + const uint32_t use_flat_tile = is_im2col ? 1u : 0u; + const bool use_hw_dot = graph.context()->adapter_ptr()->supports_int8_dot_product(); - std::string kernel_name = - use_hw_dot ? "q8ta_conv2d_pw" : "q8ta_conv2d_pw_fallback"; + std::string kernel_name; + if (use_unsigned_dot) { + VK_CHECK_COND( + use_hw_dot, + "Unsigned q8ta pointwise convolution requires integer dot product"); + kernel_name = "q8ta_conv2d_pw_unsigned"; + if (graph.storage_type_of(packed_weight) == utils::kBuffer) { + kernel_name += "_buffer"; + } + } else { + kernel_name = use_hw_dot ? "q8ta_conv2d_pw" : "q8ta_conv2d_pw_fallback"; + } + if (!use_unsigned_dot) { + // Signed PW kernels are only codegen'd for texture weights; a buffer + // weight here would fail kernel lookup at dispatch, so fail fast. + VK_CHECK_COND( + graph.storage_type_of(packed_weight) != utils::kBuffer, + "Signed q8ta pointwise convolution requires texture weights"); + } add_dtype_suffix(kernel_name, graph.dtype_of(packed_weight_scales)); vkapi::ParamsBindList param_buffers = { @@ -324,6 +404,8 @@ void add_q8ta_conv2d_pw_node( activation_type, graph.hashed_layout_of(packed_int8_output), graph.hashed_layout_of(packed_int8_input), + // Appended last to match the use_flat_tile declaration order. + use_flat_tile, }; // The im2col path passes the original activation + conv geometry so the @@ -333,18 +415,28 @@ void add_q8ta_conv2d_pw_node( // output matches directly. std::vector resize_args; ExecuteNode::ResizeFunction resize_fn; - if (conv_input == kDummyValueRef) { + if (!is_im2col) { resize_args = {packed_int8_input}; resize_fn = resize_q8ta_conv2d_pw_node; } else { - resize_args = {conv_input, kernel_size, stride, padding, dilation}; + resize_args = { + conv_input, + kernel_size, + stride, + padding, + dilation, + stream_row_offset_ref, + max_im2col_rows_ref}; resize_fn = resize_q8ta_conv2d_pw_im2col_node; } + const auto pick_gwg = + is_im2col ? pick_q8ta_conv2d_pw_streaming_gwg : pick_q8ta_conv2d_pw_gwg; + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_q8ta_conv2d_pw_gwg, + pick_gwg, pick_q8ta_conv2d_pw_lwg, {{packed_int8_output, vkapi::kWrite}, {{packed_int8_input, @@ -364,7 +456,18 @@ void add_q8ta_conv2d_pw_node( // High level operator impl // -void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args) { +bool can_use_unsigned_pw_dot( + const vkapi::Adapter& adapter, + const int64_t k_per_group) { + return adapter.accelerates_unsigned_packed4x8_dot() && + !adapter.accelerates_signed_packed4x8_dot() && + k_per_group <= kMaxUnsignedDotAccumulatorBytes; +} + +void q8ta_conv2d_pw_impl( + ComputeGraph& graph, + const bool use_unsigned_dot, + const std::vector& args) { int32_t idx = 0; const ValueRef packed_int8_input = args.at(idx++); const ValueRef input_scale = args.at(idx++); @@ -385,6 +488,12 @@ void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args) { const ValueRef activation_ref = args.at(idx++); const ValueRef packed_int8_output = args.at(idx++); + VK_CHECK_COND( + !use_unsigned_dot || + graph.size_at(-1, weight_data) <= + kMaxUnsignedDotAccumulatorBytes, + "Unsigned q8ta pointwise convolution exceeds the accumulator bound"); + uint32_t activation_type_val = static_cast( activation_type_from_string(graph.extract_string(activation_ref))); @@ -396,7 +505,8 @@ void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args) { weight_quant_config, weight_data, packed_int8_input, - packed_int8_output); + packed_int8_output, + use_unsigned_dot); ValueRef packed_weight_sums = prepack_standard( graph, weight_sums_data, utils::kBuffer, utils::kWidthPacked); @@ -422,6 +532,7 @@ void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args) { add_q8ta_conv2d_pw_node( graph, + use_unsigned_dot, packed_int8_input, input_scale, input_zp, @@ -436,6 +547,14 @@ void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args) { packed_int8_output); } +void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args) { + const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); + const ValueRef weight_data = args.at(3); + const int64_t k_per_group = graph.size_at(-1, weight_data); + const bool use_unsigned_dot = can_use_unsigned_pw_dot(*adapter, k_per_group); + q8ta_conv2d_pw_impl(graph, use_unsigned_dot, args); +} + REGISTER_OPERATORS { VK_REGISTER_OP(et_vk.q8ta_conv2d_pw.default, q8ta_conv2d_pw); } diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.cpp new file mode 100644 index 00000000000..780dfcbfbf3 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.cpp @@ -0,0 +1,147 @@ +/* + * 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. + */ + +#include + +#include +#include + +#include + +namespace vkcompute { + +namespace { + +// Computes the per-group im2col kernel columns (in-channels x kernel height +// x kernel width, aligned up to 4) for the stream-plan sizing. Returns false +// (leaving the output untouched) when any dimension is non-positive or an +// intermediate product would overflow int64; callers fail closed to the +// direct path. +bool calculate_aligned_kernel_size( + const Q8taConv2dRouteParams& params, + int64_t& aligned_kernel_size) { + if (params.in_channels_per_group <= 0 || params.kernel_height <= 0 || + params.kernel_width <= 0 || + params.in_channels_per_group > + std::numeric_limits::max() / params.kernel_height) { + return false; + } + const int64_t channel_kernel_height = + params.in_channels_per_group * params.kernel_height; + if (channel_kernel_height > + std::numeric_limits::max() / params.kernel_width) { + return false; + } + const int64_t unaligned_kernel_size = + channel_kernel_height * params.kernel_width; + if (unaligned_kernel_size > std::numeric_limits::max() - 3) { + return false; + } + aligned_kernel_size = utils::align_up_4(unaligned_kernel_size); + return true; +} + +} // namespace + +bool should_use_q8ta_conv2d_im2col(const Q8taConv2dRouteParams& params) { + if (params.batch <= 0 || params.groups <= 0 || + params.in_channels_per_group <= 0 || params.out_channels <= 0 || + params.kernel_height <= 0 || params.kernel_width <= 0 || + params.out_height <= 0 || params.out_width <= 0 || + params.out_height > + std::numeric_limits::max() / params.out_width) { + return false; + } + int64_t flattened_kernel_size; + if (!calculate_aligned_kernel_size(params, flattened_kernel_size)) { + return false; + } + const bool im2col_eligible = params.in_channels_per_group % 4 == 0; + if (!im2col_eligible) { + return false; + } + // Grouped im2col partitions output blocks per group in the PW GEMM + // (group_idx = oc_block / OC4_per_group), so each group must own whole + // packed-4 output blocks; otherwise one block straddles two groups and + // reads the wrong group's weights. + if (params.groups > 1 && + (params.out_channels % params.groups != 0 || + params.out_channels / params.groups % 4 != 0)) { + return false; + } + + const int64_t spatial_out = params.out_height * params.out_width; + if (params.batch > 1) { + constexpr int64_t kMinFlattenedKernelSize = 1024; + constexpr int64_t kMaxSpatialOutput = 64; + // Size the probe plan with the same budget the consumer uses, so + // num_tiles == 1 here means a single tile at execution too. + const Q8taConv2dStreamPlan full_plan = + make_q8ta_conv2d_stream_plan_for_device( + params.batch, + flattened_kernel_size, + params.out_height, + params.out_width, + params.max_buffer_bytes); + // Device-independent fast path: a large kernel over a tiny output makes + // the im2col materialization negligible next to the GEMM, and a single + // scratch tile means no streaming overhead, so this wins on every + // device without needing vendor-specific tuning. + const bool use_single_tile_batched_im2col = params.groups == 1 && + flattened_kernel_size >= kMinFlattenedKernelSize && + spatial_out <= kMaxSpatialOutput && full_plan.feasible && + full_plan.num_tiles == 1; + if (use_single_tile_batched_im2col) { + return true; + } + + if (!params.is_mali) { + return false; + } + + // Mali: route all eligible batched convolutions through bounded + // streaming im2col. The remaining guards are correctness bounds, not + // perf cliffs. + if (!params.supports_int8_dot_product || + // Conservative superset of the dispatch-time unsigned-path check + // (which compares the unaligned weight K): reject the aligned + // per-group K above the accumulator bound on every int8-dot path, + // failing closed for signed-path shapes near the bound as well. + flattened_kernel_size > kMaxUnsignedDotAccumulatorBytes) { + return false; + } + int64_t plan_kernel_size = flattened_kernel_size; + if (params.groups > 1) { + if (plan_kernel_size > + std::numeric_limits::max() / params.groups) { + return false; + } + plan_kernel_size *= params.groups; + } + const Q8taConv2dStreamPlan device_plan = + make_q8ta_conv2d_stream_plan_for_device( + params.batch, + plan_kernel_size, + params.out_height, + params.out_width, + params.max_buffer_bytes); + return device_plan.feasible; + } + + if (params.is_mali) { + return true; + } + + // Single-batch heuristic: im2col pays off when wide channels + // amortize the materialization over GEMM work, or when a small output + // keeps the materialized buffer cheap. Anything else stays direct. + return params.groups == 1 && + (params.in_channels_per_group >= 32 || spatial_out <= 4096); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.h b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.h new file mode 100644 index 00000000000..9f60b6fd372 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.h @@ -0,0 +1,33 @@ +/* + * 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 + +#include + +namespace vkcompute { + +struct Q8taConv2dRouteParams final { + bool is_mali; + bool supports_int8_dot_product; + // Adapter::max_buffer_numel(), which returns maxStorageBufferRange in bytes + // (not elements). + uint64_t max_buffer_bytes; + int64_t batch; + int64_t groups; + int64_t in_channels_per_group; + int64_t out_channels; + int64_t kernel_height; + int64_t kernel_width; + int64_t out_height; + int64_t out_width; +}; + +bool should_use_q8ta_conv2d_im2col(const Q8taConv2dRouteParams& params); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp index 97c939dcabf..7cf03d98a94 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp @@ -61,33 +61,6 @@ GlobalWorkGrid quantize_and_pack_4h4w_gwg( kTiledWorkGrid); } -vkapi::ShaderInfo pick_quantize_and_pack_4h4w_with_group_sums_shader( - ComputeGraph* graph, - const std::vector& args, - const std::vector& resize_args) { - const ValueRef packed_int_input = args.at(0).refs.at(0); - const ValueRef fp_input = args.at(1).refs.at(0); - const ValueRef packed_input_zps = args.at(1).refs.at(2); - const ValueRef group_size = resize_args.at(0); - - const int64_t group_size_val = graph->extract_scalar(group_size); - - std::string shader_name = "quantize_and_pack_4h4w_with_group_sums"; - if (group_size_val >= 128) { - shader_name += "_o2w32"; - } else { - shader_name += "_o4w16"; - } - - add_storage_type_suffix( - shader_name, graph->storage_type_of(packed_int_input)); - add_storage_type_suffix(shader_name, graph->storage_type_of(fp_input)); - add_dtype_suffix(shader_name, graph->dtype_of(fp_input)); - add_zp_dtype_mode_suffix(shader_name, graph->dtype_of(packed_input_zps)); - - return VK_KERNEL_FROM_STR(shader_name); -} - GlobalWorkGrid pick_quantize_and_pack_4h4w_with_group_sums_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, @@ -199,9 +172,20 @@ void add_quantize_and_pack_4h4w_with_group_sums_node( const int32_t group_size_val = graph.extract_scalar(group_size); const int32_t blocks_per_group = utils::div_up(group_size_val, int32_t(4)); + std::string shader_name = "quantize_and_pack_4h4w_with_group_sums"; + if (group_size_val >= 128) { + shader_name += "_o2w32"; + } else { + shader_name += "_o4w16"; + } + add_storage_type_suffix(shader_name, graph.storage_type_of(packed_int_input)); + add_storage_type_suffix(shader_name, graph.storage_type_of(fp_input)); + add_dtype_suffix(shader_name, graph.dtype_of(fp_input)); + add_zp_dtype_mode_suffix(shader_name, graph.dtype_of(packed_input_zps)); + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, - pick_quantize_and_pack_4h4w_with_group_sums_shader, + VK_KERNEL_FROM_STR(shader_name), pick_quantize_and_pack_4h4w_with_group_sums_gwg, pick_required_lwg, // Inputs and Outputs diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 25a8d0b89ef..1c394deca08 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -331,9 +331,11 @@ vkapi::ShaderInfo pick_linear_dqa_qw_shader( ValueRef prepack_quantized_linear_weight( ComputeGraph& graph, const QuantizationConfig& weight_quant_config, - const ValueRef qmat2_data) { + const ValueRef qmat2_data, + const bool use_unsigned_dot) { VK_CHECK_COND( weight_quant_config.nbits == 8 || weight_quant_config.nbits == 4); + VK_CHECK_COND(!use_unsigned_dot || weight_quant_config.nbits == 8); std::vector qmat2_orig_sizes = graph.sizes_of(qmat2_data); const int64_t ndim = graph.dim_of(qmat2_data); @@ -410,10 +412,13 @@ ValueRef prepack_quantized_linear_weight( if (output_width > max_extent * 4 || output_height > max_extent) { storage_type = utils::kBuffer; } - - std::string kernel_name = weight_quant_config.nbits == 4 - ? "pack_q4_linear_weight" - : "pack_q8_linear_weight"; + std::string kernel_name; + if (weight_quant_config.nbits == 4) { + kernel_name = "pack_q4_linear_weight"; + } else { + kernel_name = use_unsigned_dot ? "pack_q8_linear_weight_unsigned" + : "pack_q8_linear_weight"; + } add_storage_type_suffix(kernel_name, storage_type); // Check prepack cache before creating a new prepack node. This avoids @@ -939,6 +944,40 @@ void linear_q8csw(ComputeGraph& graph, const std::vector& args) { output); } +// aten._weight_int8pack_mm is what the AOT weight-only int8 fusion +// (FuseQuantizedOpsTransform) emits. It carries the same operands as +// et_vk.linear_q8csw minus the bias, so it runs through the same +// implementation. +void weight_int8pack_mm( + ComputeGraph& graph, + const std::vector& args) { + int32_t idx = 0; + const ValueRef fp_input = args.at(idx++); + const ValueRef weight_data = args.at(idx++); + const ValueRef weight_scales_data = args.at(idx++); + const ValueRef output = args.at(idx++); + + const int64_t K = graph.size_at(-1, fp_input); + + QuantizationConfig input_quant_config(32, kNoQuantization, {}); + QuantizationConfig weight_quant_config(8, kPerChannel, {K}); + + quantized_linear_impl( + graph, + input_quant_config, + weight_quant_config, + fp_input, + kDummyValueRef, // input scale + kDummyValueRef, // input zp + weight_data, + kDummyValueRef, // weight sums + weight_scales_data, + kDummyValueRef, // weight zeros + kDummyValueRef, // group size + kDummyValueRef, // bias + output); +} + void linear_dq8ca_q4gsw( ComputeGraph& graph, const std::vector& args) { @@ -977,6 +1016,7 @@ void linear_dq8ca_q4gsw( REGISTER_OPERATORS { VK_REGISTER_OP(et_vk.linear_q8ta_q8csw.default, linear_q8ta_q8csw); VK_REGISTER_OP(et_vk.linear_q8csw.default, linear_q8csw); + VK_REGISTER_OP(aten._weight_int8pack_mm.default, weight_int8pack_mm); VK_REGISTER_OP(et_vk.linear_dq8ca_q4gsw.default, linear_dq8ca_q4gsw); } diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h index 7cb0e172c4a..4ab87f8ac33 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h @@ -24,6 +24,7 @@ LocalWorkGroup quantized_linear_lwg( ValueRef prepack_quantized_linear_weight( ComputeGraph& graph, const QuantizationConfig& weight_quant_config, - const ValueRef qmat2_data); + const ValueRef qmat2_data, + const bool use_unsigned_dot = false); } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp index 5d9311e9761..c2e45b96418 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp @@ -16,34 +16,6 @@ namespace vkcompute { -// Custom global workgroup size function for linear_qcs8w -GlobalWorkGrid linear_qcs8w_gwg( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const std::vector& args, - const std::vector& resize_args) { - (void)shader; - (void)resize_args; - const ValueRef out = args.at(0).refs.at(0); - return graph->create_linear_gwg( - utils::safe_downcast(graph->numel_of(out))); -} - -// Custom local workgroup size function for linear_qcs8w -LocalWorkGroup linear_qcs8w_lwg( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const GlobalWorkGrid& gwg, - const std::vector& args, - const std::vector& resize_args) { - (void)graph; - (void)shader; - (void)gwg; - (void)args; - (void)resize_args; - return LocalWorkGroup(64u, 1u, 1u); -} - // Custom global workgroup size function for linear_qcsnw_tiled GlobalWorkGrid linear_qcsnw_tiled_gwg( ComputeGraph* graph, @@ -178,81 +150,6 @@ void resize_linear_qcsnw_node( graph->virtual_resize(out, new_out_sizes); } -void add_linear_qcs8w_node( - ComputeGraph& graph, - const ValueRef mat1, - const ValueRef q_mat2_data, - const ValueRef scales_data, - const ValueRef out) { - auto viewFn = VK_GET_OP_FN("aten.view_copy.default"); - ValueRef mat1_W_packed = mat1; - ValueRef out_W_packed = out; - // Create temporary tensors to store the width packed versions of mat1 and out - TmpTensor mat1_tmp( - &graph, graph.sizes_of(mat1), graph.dtype_of(mat1), utils::kWidthPacked); - TmpTensor out_tmp( - &graph, graph.sizes_of(out), graph.dtype_of(out), utils::kWidthPacked); - if (!graph.is_buffer_storage(out) && - graph.packed_dim_of(mat1) != WHCN::kWidthDim) { - // Ensure mat1 is width packed - mat1_W_packed = mat1_tmp; - viewFn(graph, {mat1, graph.add_none(), mat1_W_packed}); - // Ensure out is packed correctly - out_W_packed = out_tmp; - } - ValueRef q_mat2 = prepack_standard_hw_transposed( - graph, q_mat2_data, graph.storage_type_of(out), utils::kWidthPacked); - ValueRef scales = prepack_standard( - graph, scales_data, graph.storage_type_of(out), utils::kWidthPacked); - - std::string kernel_name = "linear_qcs8w"; - kernel_name.reserve(kShaderNameReserve); - add_packed_dim_suffix(kernel_name, graph.packed_dim_of(mat1_W_packed)); - add_packed_dim_suffix(kernel_name, graph.packed_dim_of(q_mat2)); - add_dtype_suffix(kernel_name, graph.dtype_of(out_W_packed)); - add_storage_type_suffix(kernel_name, graph.storage_type_of(out_W_packed)); - - std::vector pcs; - if (graph.is_buffer_storage(out_W_packed)) { - pcs = { - graph.sizes_pc_of(out_W_packed), - graph.strides_pc_of(out_W_packed), - graph.sizes_pc_of(mat1_W_packed), - graph.strides_pc_of(mat1), - graph.strides_pc_of(q_mat2), - graph.strides_pc_of(scales), - graph.numel_pc_of(out_W_packed)}; - } else { - pcs = { - graph.logical_limits_pc_of(out_W_packed), - graph.sizes_pc_of(mat1_W_packed), - graph.sizes_pc_of(q_mat2)}; - } - - graph.execute_nodes().emplace_back(new DynamicDispatchNode( - graph, - VK_KERNEL_FROM_STR(kernel_name), - linear_qcs8w_gwg, - linear_qcs8w_lwg, - // Inputs and Outputs - {{out_W_packed, vkapi::MemoryAccessType::WRITE}, - {{mat1_W_packed, q_mat2, scales}, vkapi::MemoryAccessType::READ}}, - // Shader params buffers - {}, - // Push Constants - pcs, - // Specialization Constants - {}, - // Resize Args - {}, - // Resizing Logic - resize_linear_qcsnw_node)); - if (!graph.is_buffer_storage(out) && - graph.packed_dim_of(out) != WHCN::kWidthDim) { - viewFn(graph, {out_W_packed, graph.add_none(), out}); - } -} - void add_linear_qcsnw_tiled_node( ComputeGraph& graph, const bool use_coop_algorithm, @@ -293,6 +190,8 @@ void add_linear_qcsnw_tiled_node( kernel_name = use_coop_algorithm ? "linear_qcs4w_coop" : "linear_qcs4w_tiled"; } else { + // Unreachable: linear_qcs4w is the only caller of this function and it + // always passes quant_nbits == 4. Kept so the 4-bit path is not disturbed. kernel_name = use_coop_algorithm ? "linear_qcs8w_coop" : "linear_qcs8w_tiled"; } @@ -389,18 +288,6 @@ bool can_use_coop_impl(ComputeGraph& graph, const ValueRef mat1) { return (graph.size_at(-2, mat1) == 1); } -void weight_int8pack_mm( - ComputeGraph& graph, - const std::vector& args) { - check_linear_qcsnw_args(graph, 8, args[0], args[1], args[2], args[3]); - if (can_use_tiled_impl(graph, args[0], args[1], args[2], args[3])) { - bool use_coop_algorithm = can_use_coop_impl(graph, args[0]); - return add_linear_qcsnw_tiled_node( - graph, use_coop_algorithm, 8, args[0], args[1], args[2], args[3]); - } - return add_linear_qcs8w_node(graph, args[0], args[1], args[2], args[3]); -} - void linear_qcs4w(ComputeGraph& graph, const std::vector& args) { check_linear_qcsnw_args(graph, 4, args[0], args[1], args[2], args[3]); @@ -411,7 +298,9 @@ void linear_qcs4w(ComputeGraph& graph, const std::vector& args) { } REGISTER_OPERATORS { - VK_REGISTER_OP(aten._weight_int8pack_mm.default, weight_int8pack_mm); + // aten._weight_int8pack_mm is registered in QuantizedLinear.cpp, on the + // maintained weight-only quantized linear implementation. The 8-bit path + // here dispatched to linear_qcs8w_* shaders that no longer exist. VK_REGISTER_OP(et_vk.linear_qcs4w.default, linear_qcs4w); } diff --git a/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp b/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp index f684906cf57..23c8f000bed 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp @@ -14,6 +14,8 @@ #include #include +#include + namespace vkcompute { using namespace utils; @@ -72,30 +74,111 @@ void resize_reduce_per_row_node( graph->virtual_resize(out, new_sizes); } -GlobalWorkGrid reduce_gwg( +// Number of threads that co-operate on one reduction output, and how many +// outputs one work group covers. +// +// The worker count used to be a flat 4 regardless of how much there was to +// reduce. A global average pool collapses a whole HxW plane, so four threads +// each walked thousands of elements while the dispatch ran 16 threads in total. +constexpr uint32_t kReduceMaxNThreads = 256u; +constexpr uint32_t kReduceNGroups = 4u; + +// The largest worker count this dispatch may ask for. Three ceilings apply and +// the smallest of them wins: +// +// - the shaders size shared_vecs at MAX_NTHREADS and every thread in the work +// group writes its own slot, so kReduceNGroups workers must fit; +// - the device bounds the invocations in one work group, and +// maxComputeWorkGroupInvocations is only guaranteed to be 128; +// - the device bounds each work group axis on its own, and the workers all sit +// on the reduction axis. +// +// Overrunning any of them aborts the dispatch in LocalWorkGroup::validate, so +// the shader capacity alone is not enough to go by. +uint32_t reduce_nworkers_cap( + ComputeGraph* graph, + const int32_t reduce_dim_whcn) { + const vkapi::Adapter* const adapter = graph->context()->adapter_ptr(); + uint32_t cap = kReduceMaxNThreads / kReduceNGroups; + cap = std::min( + cap, adapter->max_compute_workgroup_invocations() / kReduceNGroups); + cap = std::min(cap, adapter->max_compute_workgroup_size()[reduce_dim_whcn]); + return std::max(cap, 1u); +} + +uint32_t reduce_nworkers( + ComputeGraph* graph, + const ValueRef in, + const int32_t reduce_dim_whcn) { + const uint32_t cap = reduce_nworkers_cap(graph, reduce_dim_whcn); + const uint32_t extent = utils::safe_downcast( + graph->logical_limits_of(in)[reduce_dim_whcn]); + // 4 is what this used to be unconditionally; keep it as the floor so short + // reductions dispatch exactly as they did before. + uint32_t nworkers = std::min(4u, cap); + while (nworkers * 2u <= cap && nworkers < extent) { + nworkers *= 2u; + } + return nworkers; +} + +GlobalWorkGrid reduce_gwg_impl( ComputeGraph* graph, - const vkapi::ShaderInfo& shader, const std::vector& args, - const std::vector& resize_args) { - (void)shader; + const std::vector& resize_args, + const size_t reduce_dim_idx, + const size_t group_dim_idx, + const size_t nworkers_idx) { const ValueRef out = args.at(0).refs.at(0); const int32_t reduce_dim_whcn = - graph->extract_scalar(resize_args.at(1)); + graph->extract_scalar(resize_args.at(reduce_dim_idx)); const int64_t group_dim_whcn = - graph->extract_scalar(resize_args.at(2)); + graph->extract_scalar(resize_args.at(group_dim_idx)); utils::uvec3 extents = graph->logical_limits_of(out); extents[reduce_dim_whcn] = 1; - constexpr uint32_t max_nthreads = 16u; - constexpr uint32_t nworkers_per_group = 4u; - constexpr uint32_t ngroups = 4u; - VK_CHECK_COND(nworkers_per_group * ngroups <= max_nthreads); + // NWORKERS is baked into the shader as a specialization constant when the + // node is built, so it reflects the dynamic upper bound. This callback runs + // again after every resize and would see the smaller actual extent, giving a + // work group with fewer threads along the reduction dim than the shader's + // aggregation loop indexes over -- it would then fold in shared memory slots + // that no thread wrote. Read the count the node was built with instead. + // + // Launching more workers than there are elements is harmless: a worker whose + // loop body never runs contributes INIT_ACCUM, which is the identity for sum + // and mean and idempotent for amax and amin. + const uint32_t nworkers_per_group = utils::safe_downcast( + graph->extract_scalar(resize_args.at(nworkers_idx))); utils::uvec3 lwg_extents{1u, 1u, 1u}; lwg_extents[reduce_dim_whcn] = nworkers_per_group; - lwg_extents[group_dim_whcn] = ngroups; + lwg_extents[group_dim_whcn] = kReduceNGroups; return GlobalWorkGrid(extents, kTiledWorkGrid, LocalWorkGroup(lwg_extents)); } +// Resize args are {dim, reduce_dim_whcn, group_dim_whcn, nworkers}. +GlobalWorkGrid reduce_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + return reduce_gwg_impl(graph, args, resize_args, 1, 2, 3); +} + +// Resize args are {dims, reduce_dim1_whcn, reduce_dim2_whcn, group_dim_whcn, +// nworkers}, so the group dim sits one slot further along than in the 1d case. +// Sharing the 1d picker put the groups on reduce_dim2 and left the group axis +// one thread wide, so every group read tid.y == 0 and raced over the same +// shared memory slots. +GlobalWorkGrid reduce2d_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + return reduce_gwg_impl(graph, args, resize_args, 1, 3, 4); +} + void add_reduce_node( ComputeGraph& graph, const ValueRef in, @@ -137,6 +220,9 @@ void add_reduce_node( const ValueRef reduce_dim_whcn_ref = graph.get_or_add_value_for_int(reduce_dim); const ValueRef group_dim_whcn_ref = graph.get_or_add_value_for_int(group_dim); + const int32_t nworkers = + utils::safe_downcast(reduce_nworkers(&graph, in, reduce_dim)); + const ValueRef nworkers_ref = graph.get_or_add_value_for_int(nworkers); graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, @@ -149,10 +235,12 @@ void add_reduce_node( {graph.logical_limits_ubo(in), graph.sizes_ubo(in)}, // Push Constants {}, - // Specialization Constants - {graph.packed_dim_of(out), reduce_dim, group_dim}, + // Specialization Constants. NWORKERS must match the local work group + // extent reduce_gwg picks, so the count is computed once here and passed + // to reduce_gwg through the resize args. + {graph.packed_dim_of(out), reduce_dim, group_dim, nworkers}, // Resize Args - {dim_ref, reduce_dim_whcn_ref, group_dim_whcn_ref}, + {dim_ref, reduce_dim_whcn_ref, group_dim_whcn_ref, nworkers_ref}, // Resizing Logic resize_reduce_node)); } @@ -214,11 +302,14 @@ void add_reduce2d_node( const ValueRef reduce_dim2_whcn_ref = graph.get_or_add_value_for_int(reduce_dim2); const ValueRef group_dim_whcn_ref = graph.get_or_add_value_for_int(group_dim); + const int32_t nworkers = + utils::safe_downcast(reduce_nworkers(&graph, in, reduce_dim1)); + const ValueRef nworkers_ref = graph.get_or_add_value_for_int(nworkers); graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - reduce_gwg, + reduce2d_gwg, pick_required_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, @@ -226,13 +317,16 @@ void add_reduce2d_node( {graph.logical_limits_ubo(in), graph.sizes_ubo(in)}, // Push Constants {}, - // Specialization Constants - {graph.packed_dim_of(out), reduce_dim1, reduce_dim2, group_dim}, + // Specialization Constants. NWORKERS must match the local work group + // extent reduce_gwg picks, so the count is computed once here and passed + // to reduce_gwg through the resize args. + {graph.packed_dim_of(out), reduce_dim1, reduce_dim2, group_dim, nworkers}, // Resize Args {dims_ref, reduce_dim1_whcn_ref, reduce_dim2_whcn_ref, - group_dim_whcn_ref}, + group_dim_whcn_ref, + nworkers_ref}, // Resizing Logic resize_reduce2d_node)); } diff --git a/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp b/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp index fce7600a035..bec078564b4 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp @@ -14,6 +14,8 @@ #include #include +#include + namespace vkcompute { void add_squeeze_copy_dims_node( @@ -23,21 +25,27 @@ void add_squeeze_copy_dims_node( const ValueRef out) { const int64_t in_dim = graph.dim_of(in); const std::vector in_sizes = graph.sizes_of(in); - const std::vector out_sizes = graph.sizes_of(in); const std::vector dims = graph.extract_int_or_symint_list(dims_ref); std::vector squeeze_dims; - // Filter out edge cases that we don't need squeeze: - // 1. The size of squeeze dim is larger than 1. - // 2. Squeeze outter most dim - // For these cases, just pass input to output via clone. + // Filter out the edge case that we don't need to squeeze: the size of the + // squeeze dim is larger than 1. For that case, just pass input to output via + // clone. + // + // Note that the outermost dim must NOT be excluded here. Routing it to + // add_clone_node() leaves the output unresized at runtime, because + // resize_clone_node() only propagates sizes when input and output have the + // same dim count -- which is never true for a squeeze. Under dynamic shapes + // the output then keeps its upper-bound extents while consumers read it at + // the real size, silently producing wrong values. add_permute_node()'s + // resize function handles the rank-reducing case explicitly. for (int i = 0; i < dims.size(); ++i) { // adjust negative dims int64_t dim_val = dims.at(i); if (dim_val < 0) { dim_val += in_dim; } - if (dims.at(i) != 0 && in_sizes.at(dim_val) == 1) { + if (in_sizes.at(dim_val) == 1) { squeeze_dims.push_back(dim_val); } } diff --git a/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp b/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp index 6a50cb2f6a9..d17f57774f7 100644 --- a/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp @@ -69,6 +69,46 @@ void add_unary_op_node( resize_unary_op_node)); } +void add_dynamic_clamp_node( + ComputeGraph& graph, + const ValueRef in, + const ValueRef min, + const ValueRef max, + const ValueRef out) { + std::string kernel_name("clamp_dynamic"); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); + + const bool output_is_int = graph.dtype_of(out) == vkapi::kInt; + const utils::ivec2 bounds_are_int = { + output_is_int || graph.val_is_symint(min) ? 1 : 0, + output_is_int || graph.val_is_symint(max) ? 1 : 0}; + const vkapi::BufferBindInfo min_param = bounds_are_int[0] + ? graph.get_or_create_int_param_buffer( + min, std::numeric_limits::min()) + : graph.create_params_buffer(graph.extract_scalar_or( + min, -std::numeric_limits::infinity())); + const vkapi::BufferBindInfo max_param = bounds_are_int[1] + ? graph.get_or_create_int_param_buffer( + max, std::numeric_limits::max()) + : graph.create_params_buffer(graph.extract_scalar_or( + max, std::numeric_limits::infinity())); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + default_pick_gwg, + default_pick_lwg, + {{out, vkapi::kWrite}, {in, vkapi::kRead}}, + {min_param, max_param}, + {graph.is_buffer_storage(out) ? graph.numel_pc_of(out) + : graph.logical_limits_pc_of(out), + PushConstantDataInfo(&bounds_are_int, sizeof(bounds_are_int))}, + {}, + {}, + resize_unary_op_node)); +} + float get_val_or_inf(ComputeGraph& graph, const ValueRef& val, bool max) { if (!graph.val_is_none(val)) { return graph.extract_scalar(val); @@ -85,6 +125,10 @@ float get_val_or_inf(ComputeGraph& graph, const ValueRef& val, bool max) { #define DEFINE_CLAMP_FN(op_name) \ void op_name(ComputeGraph& graph, const std::vector& args) { \ + if (graph.val_is_symint(args[1]) || graph.val_is_symint(args[2])) { \ + return add_dynamic_clamp_node( \ + graph, args[0], args[1], args[2], args[3]); \ + } \ return add_unary_op_node( \ graph, \ args[0], \ diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp index df83bdec7e0..90b923194ad 100644 --- a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp @@ -18,10 +18,15 @@ std::vector calculate_broadcasted_output_size( const std::vector& sizes1, const std::vector& sizes2) { std::vector out_sizes(std::max(sizes1.size(), sizes2.size())); + // ndim must be signed. `-out_sizes.size()` is size_t arithmetic, so when both + // inputs are 0-dimensional it evaluates to 0 while `i` promotes to a huge + // unsigned value, the guard stays true, and `out_sizes.at(size() - 1)` throws + // std::out_of_range instead of the loop being skipped. + const int64_t ndim = static_cast(out_sizes.size()); // Match the sizes in reverse because sizes are in NCHW order - for (int i = -1; i >= -out_sizes.size(); --i) { - out_sizes.at(out_sizes.size() + i) = + for (int64_t i = -1; i >= -ndim; --i) { + out_sizes.at(static_cast(ndim + i)) = std::max(utils::val_at(i, sizes1), utils::val_at(i, sizes2)); } diff --git a/backends/vulkan/runtime/vk_api/Adapter.cpp b/backends/vulkan/runtime/vk_api/Adapter.cpp index 3d9acae8975..eea61a8fe3f 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.cpp +++ b/backends/vulkan/runtime/vk_api/Adapter.cpp @@ -350,6 +350,17 @@ Adapter::~Adapter() { } } +ScopedAdapterCapabilityOverride::ScopedAdapterCapabilityOverride( + Adapter* adapter, + AdapterCapabilityOverrides overrides) + : adapter_(adapter), previous_(adapter->capability_overrides_) { + adapter_->capability_overrides_ = overrides; +} + +ScopedAdapterCapabilityOverride::~ScopedAdapterCapabilityOverride() { + adapter_->capability_overrides_ = previous_; +} + Adapter::Queue Adapter::request_queue() { // Lock the mutex as multiple threads can request a queue at the same time std::lock_guard lock(queue_usage_mutex_); diff --git a/backends/vulkan/runtime/vk_api/Adapter.h b/backends/vulkan/runtime/vk_api/Adapter.h index 6eae09e8eb6..0c2784ae473 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.h +++ b/backends/vulkan/runtime/vk_api/Adapter.h @@ -17,6 +17,8 @@ #include +#include + #include namespace vkcompute { @@ -78,6 +80,8 @@ class Adapter final { VkQueue handle; }; + friend class ScopedAdapterCapabilityOverride; + private: // Use a mutex to manage queue usage info since // it can be accessed from multiple threads @@ -102,6 +106,9 @@ class Adapter final { // Miscellaneous bool linear_tiling_3d_enabled_; bool owns_device_; + // Test-only capability overrides; empty unless a ScopedCapabilityOverride + // is live. + AdapterCapabilityOverrides capability_overrides_; public: // Physical Device metadata @@ -231,7 +238,10 @@ class Adapter final { #endif /* VK_KHR_shader_float16_int8 */ } - inline bool supports_int8_dot_product() { + inline bool supports_int8_dot_product() const { + if (capability_overrides_.int8_dot_product.has_value()) { + return *capability_overrides_.int8_dot_product; + } #ifdef ETVK_FORCE_NO_EXTENSIONS return false; #endif @@ -243,6 +253,40 @@ class Adapter final { #endif /* VK_KHR_shader_integer_dot_product */ } + inline bool accelerates_signed_packed4x8_dot() const { + if (capability_overrides_.signed_packed4x8_dot.has_value()) { + return *capability_overrides_.signed_packed4x8_dot; + } +#ifdef ETVK_FORCE_NO_EXTENSIONS + return false; +#endif +#ifdef VK_KHR_shader_integer_dot_product + return supports_int8_dot_product() && + physical_device_.shader_int_dot_product_properties + .integerDotProductAccumulatingSaturating4x8BitPackedSignedAccelerated == + VK_TRUE; +#else + return false; +#endif /* VK_KHR_shader_integer_dot_product */ + } + + inline bool accelerates_unsigned_packed4x8_dot() const { + if (capability_overrides_.unsigned_packed4x8_dot.has_value()) { + return *capability_overrides_.unsigned_packed4x8_dot; + } +#ifdef ETVK_FORCE_NO_EXTENSIONS + return false; +#endif +#ifdef VK_KHR_shader_integer_dot_product + return supports_int8_dot_product() && + physical_device_.shader_int_dot_product_properties + .integerDotProductAccumulatingSaturating4x8BitPackedUnsignedAccelerated == + VK_TRUE; +#else + return false; +#endif /* VK_KHR_shader_integer_dot_product */ + } + inline bool supports_nv_cooperative_matrix2() { #ifdef ETVK_FORCE_NO_EXTENSIONS return false; diff --git a/backends/vulkan/runtime/vk_api/AdapterCapabilityOverrides.h b/backends/vulkan/runtime/vk_api/AdapterCapabilityOverrides.h new file mode 100644 index 00000000000..f4b6c6d9907 --- /dev/null +++ b/backends/vulkan/runtime/vk_api/AdapterCapabilityOverrides.h @@ -0,0 +1,61 @@ +/* + * 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 + +#include + +namespace vkcompute { +namespace vkapi { + +class Adapter; + +// Test-only device capability overrides. Each set field replaces the +// corresponding Adapter query, simulating a weaker device so fallback paths +// can be covered on capable hardware. Forcing a capability the device +// lacks is undefined behavior and typically fails at pipeline creation. +// Overrides apply to every graph built against the Adapter while set; +// hold a ScopedAdapterCapabilityOverride to bound the scope to one test. +// Not thread-safe: while set, every query on the Adapter observes the +// override, so hold it only around single-test build plus execute with no +// concurrent adapter use. +struct AdapterCapabilityOverrides { + std::optional int8_dot_product; + std::optional signed_packed4x8_dot; + std::optional unsigned_packed4x8_dot; + + static AdapterCapabilityOverrides without_dot_product_support() { + AdapterCapabilityOverrides overrides; + overrides.int8_dot_product = false; + overrides.signed_packed4x8_dot = false; + overrides.unsigned_packed4x8_dot = false; + return overrides; + } +}; + +class ScopedAdapterCapabilityOverride final { + public: + ScopedAdapterCapabilityOverride( + Adapter* adapter, + AdapterCapabilityOverrides overrides); + ~ScopedAdapterCapabilityOverride(); + ScopedAdapterCapabilityOverride(const ScopedAdapterCapabilityOverride&) = + delete; + ScopedAdapterCapabilityOverride& operator=( + const ScopedAdapterCapabilityOverride&) = delete; + ScopedAdapterCapabilityOverride(ScopedAdapterCapabilityOverride&&) = delete; + ScopedAdapterCapabilityOverride& operator=( + ScopedAdapterCapabilityOverride&&) = delete; + + private: + Adapter* adapter_; + AdapterCapabilityOverrides previous_; +}; + +} // namespace vkapi +} // namespace vkcompute diff --git a/backends/vulkan/serialization/vulkan_graph_builder.py b/backends/vulkan/serialization/vulkan_graph_builder.py index 46e01e701b1..e5dfde9d865 100644 --- a/backends/vulkan/serialization/vulkan_graph_builder.py +++ b/backends/vulkan/serialization/vulkan_graph_builder.py @@ -418,15 +418,26 @@ def get_or_create_value_for(self, arg: _Argument): raise RuntimeError(f"Cannot create value for arg of type {type(arg)}") def process_placeholder_node(self, node: Node) -> None: - # ignores any tensors that don't get used in any ops - if len(node.users) == 0: + # A non-param placeholder occupies a slot in the delegate call's + # argument list whether or not this graph goes on to use it, and + # VulkanBackend::execute matches `args` to graph inputs positionally. + # Dropping an unused one from input_ids desynchronises the two, and the + # runtime then rejects the call because it was handed more arguments + # than the graph declares inputs and outputs. That happens in practice + # when a placeholder's only consumers are folded away by the passes + # that run after partitioning, so the graph the partitioner tagged and + # the graph serialized here disagree about which inputs are live. + if is_param_node(self.program, node): + # Params are serialized into the blob rather than passed at call + # time, so an unused one costs nothing to skip. + if len(node.users) > 0: + self.create_node_value(node) return None ids = self.create_node_value(node) - if not is_param_node(self.program, node): - if isinstance(ids, int): - self.input_ids.append(ids) - else: - self.input_ids += ids + if isinstance(ids, int): + self.input_ids.append(ids) + else: + self.input_ids += ids def process_getitem_node(self, node: Node) -> None: # Find ValueList id from the collection node. diff --git a/backends/vulkan/serialization/vulkan_graph_serialize.py b/backends/vulkan/serialization/vulkan_graph_serialize.py index 96f944560a8..81de183021b 100644 --- a/backends/vulkan/serialization/vulkan_graph_serialize.py +++ b/backends/vulkan/serialization/vulkan_graph_serialize.py @@ -11,6 +11,7 @@ import importlib.resources as _resources import json import os +import re import tempfile from dataclasses import dataclass from typing import ClassVar, List @@ -27,8 +28,72 @@ from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile +# Python's json module spells the non-finite floats "Infinity" / "-Infinity" / +# "NaN"; flatc spells the infinities "inf" / "-inf" and rejects Python's +# spelling, so both directions need translating. A graph carries a non-finite +# scalar whenever the model does -- the -inf fill value of a transformer +# attention mask is the common case -- and without this the failure surfaces as +# a flatc byte offset into a temporary file rather than anything pointing at +# the graph. +# +# The rewrite runs over the serialized text rather than over the encoder's +# chunks: json only emits a float as a chunk of its own inside an object, and +# inside a list the chunk carries the delimiter with it ("[-Infinity"), so +# matching whole chunks silently missed every DoubleList. +_JSON_STRING_RE = re.compile(r'"(?:[^"\\]|\\.)*"') +_PY_NONFINITE_RE = re.compile(r"(? str: + """Apply ``sub`` to everything in ``text`` that is not a JSON string. + + String literals are copied through untouched, so a shader name or a string + value that happens to read "inf" is never rewritten. + """ + out = [] + last = 0 + for m in _JSON_STRING_RE.finditer(text): + out.append(sub(text[last : m.start()])) + out.append(m.group(0)) + last = m.end() + out.append(sub(text[last:])) + return "".join(out) + + +def _python_json_to_flatc_json(text: str) -> str: + """Rewrite json's ``Infinity`` tokens into the ``inf`` flatc accepts.""" + if "Infinity" not in text and "NaN" not in text: + return text + + def replace(m: "re.Match[str]") -> str: + if m.group(1) == "NaN": + raise ValueError( + "Cannot serialize a NaN float value into a Vulkan graph: " + "flatc rejects every spelling of NaN for a value inside a " + "union, and every float in the Vulkan schema is a member of " + "the VkValue union." + ) + return "inf" + + return _rewrite_outside_strings( + text, lambda segment: _PY_NONFINITE_RE.sub(replace, segment) + ) + + +def _flatc_json_to_python_json(text: str) -> str: + """Rewrite flatc's bare ``inf`` tokens so json.loads accepts them.""" + if "inf" not in text: + return text + return _rewrite_outside_strings( + text, lambda segment: _FLATC_INF_RE.sub("Infinity", segment) + ) + + def convert_to_flatbuffer(vk_graph: VkGraph) -> bytes: - vk_graph_json = json.dumps(vk_graph, cls=_DataclassEncoder) + vk_graph_json = _python_json_to_flatc_json( + json.dumps(vk_graph, cls=_DataclassEncoder) + ) with tempfile.TemporaryDirectory() as d: schema_path = os.path.join(d, "schema.fbs") @@ -63,7 +128,8 @@ def flatbuffer_to_vk_graph(flatbuffers: bytes) -> VkGraph: json_path = os.path.join(d, "schema.json") with open(json_path, "rb") as output_file: - return _json_to_dataclass(json.load(output_file), VkGraph) + raw = output_file.read().decode("utf-8") + return _json_to_dataclass(json.loads(_flatc_json_to_python_json(raw)), VkGraph) def extract_vk_flatbuffer(data: bytes) -> bytes: diff --git a/backends/vulkan/test/custom_ops/build_and_run.sh b/backends/vulkan/test/custom_ops/build_and_run.sh index b1195568b1b..2641970faf8 100755 --- a/backends/vulkan/test/custom_ops/build_and_run.sh +++ b/backends/vulkan/test/custom_ops/build_and_run.sh @@ -30,13 +30,13 @@ configure_and_build_main() { -B$CMAKE_OUT_DIR fi - cmake --build $CMAKE_OUT_DIR -j16 --target install + cmake --build $CMAKE_OUT_DIR -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install # -DCMAKE_CXX_FLAGS="-DVULKAN_DEBUG" \ } # Function to build main project only build_main() { - cmake --build $CMAKE_OUT_DIR -j16 --target install + cmake --build $CMAKE_OUT_DIR -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install } # Function to configure and build tests @@ -65,12 +65,12 @@ configure_and_build_tests() { -B$CMAKE_OUT_DIR/backends/vulkan/test/custom_ops fi - cmake --build $CMAKE_OUT_DIR/backends/vulkan/test/custom_ops -j16 --target all + cmake --build $CMAKE_OUT_DIR/backends/vulkan/test/custom_ops -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target all } build_tests() { - cmake --build $CMAKE_OUT_DIR/backends/vulkan/test/custom_ops -j16 --target all + cmake --build $CMAKE_OUT_DIR/backends/vulkan/test/custom_ops -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target all } # Function to rebuild both main and tests diff --git a/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp b/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp index 679ac33d11b..e4af2c56b9d 100644 --- a/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp @@ -10,10 +10,140 @@ #include #include +#include #include namespace vkcompute { +namespace { + +void assert_im2col_kernel_selection( + ComputeGraph& graph, + const bool expect_unsigned, + const bool expect_buffer_weights, + const bool expect_fallback_kernel = false) { + const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); + std::string expected_execute; + std::string expected_prepack; + if (expect_unsigned) { + expected_execute = expect_buffer_weights + ? "q8ta_conv2d_pw_unsigned_buffer_float" + : "q8ta_conv2d_pw_unsigned_float"; + expected_prepack = expect_buffer_weights + ? "pack_q8_linear_weight_unsigned_buffer" + : "pack_q8_linear_weight_unsigned_texture2d"; + } else { + VK_CHECK_COND(!expect_buffer_weights); + if (expect_fallback_kernel) { + expected_execute = "q8ta_conv2d_pw_fallback_float"; + } else { + expected_execute = adapter->supports_int8_dot_product() + ? "q8ta_conv2d_pw_float" + : "q8ta_conv2d_pw_fallback_float"; + } + expected_prepack = "pack_q8_linear_weight_texture2d"; + } + + int32_t execute_matches = 0; + std::string execute_names; + for (const auto& node : graph.execute_nodes()) { + const ExecuteNode* const node_ptr = node.get(); + VK_CHECK_COND(node_ptr != nullptr); + const std::string& node_name = node_ptr->name(); + execute_names += node_name + " "; + if (node_name.find("q8ta_conv2d_pw") == 0) { + VK_CHECK_COND( + node_name == expected_execute, + "Expected ", + expected_execute, + " but selected execute kernel ", + node_name); + ++execute_matches; + } + } + VK_CHECK_COND(execute_matches > 0, "Execute kernels: ", execute_names); + + int32_t prepack_matches = 0; + std::string prepack_names; + for (const auto& node : graph.prepack_nodes()) { + const PrepackNode* const node_ptr = node.get(); + VK_CHECK_COND(node_ptr != nullptr); + const std::string& node_name = node_ptr->name(); + prepack_names += node_name + " "; + if (node_name.find("pack_q8_linear_weight") == 0) { + VK_CHECK_COND( + node_name == expected_prepack, + "Expected ", + expected_prepack, + " but selected prepack kernel ", + node_name); + ++prepack_matches; + } + } + VK_CHECK_COND(prepack_matches > 0, "Prepack kernels: ", prepack_names); +} + +void assert_pw_kernel_selection( + ComputeGraph& graph, + const bool expect_unsigned, + const bool expect_buffer_weights) { + const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); + std::string expected_execute; + std::string expected_prepack; + if (expect_unsigned) { + expected_execute = "q8ta_conv2d_pw_unsigned_float"; + expected_prepack = expect_buffer_weights + ? "pack_q8_conv2d_weights_unsigned_buffer" + : "pack_q8_conv2d_weights_unsigned_texture2d"; + } else { + VK_CHECK_COND(!expect_buffer_weights); + expected_execute = adapter->supports_int8_dot_product() + ? "q8ta_conv2d_pw_float" + : "q8ta_conv2d_pw_fallback_float"; + expected_prepack = "pack_q8_conv2d_weights_texture2d"; + } + + int32_t execute_matches = 0; + std::string execute_names; + for (const auto& node : graph.execute_nodes()) { + const ExecuteNode* const node_ptr = node.get(); + VK_CHECK_COND(node_ptr != nullptr); + const std::string& node_name = node_ptr->name(); + execute_names += node_name + " "; + if (node_name.find("q8ta_conv2d_pw") == 0) { + VK_CHECK_COND( + node_name == expected_execute, + "Expected ", + expected_execute, + " but selected execute kernel ", + node_name); + ++execute_matches; + } + } + VK_CHECK_COND(execute_matches > 0, "Execute kernels: ", execute_names); + + int32_t prepack_matches = 0; + std::string prepack_names; + for (const auto& node : graph.prepack_nodes()) { + const PrepackNode* const node_ptr = node.get(); + VK_CHECK_COND(node_ptr != nullptr); + const std::string& node_name = node_ptr->name(); + prepack_names += node_name + " "; + if (node_name.find("pack_q8_conv2d_weights") == 0) { + VK_CHECK_COND( + node_name == expected_prepack, + "Expected ", + expected_prepack, + " but selected prepack kernel ", + node_name); + ++prepack_matches; + } + } + VK_CHECK_COND(prepack_matches > 0, "Prepack kernels: ", prepack_names); +} + +} // namespace + void test_q8ta_conv2d_dw( ComputeGraph& graph, const std::vector& args) { @@ -122,28 +252,34 @@ void test_q8ta_conv2d(ComputeGraph& graph, const std::vector& args) { const ValueRef dilation = args.at(idx++); const ValueRef groups = args.at(idx++); const ValueRef activation = args.at(idx++); - const ValueRef layout_int = args.at(idx++); + const ValueRef input_layout_int = args.at(idx++); + const ValueRef output_layout_int = args.at(idx++); const ValueRef impl_selector_str = args.at(idx++); const ValueRef fp_output = args.at(idx++); // Extract the layout parameter and cast to GPUMemoryLayout - int32_t layout_value = graph.extract_scalar(layout_int); - utils::GPUMemoryLayout layout = - static_cast(layout_value); + const auto input_layout = static_cast( + graph.extract_scalar(input_layout_int)); + const auto output_layout = static_cast( + graph.extract_scalar(output_layout_int)); // Extract the impl_selector string std::string impl_selector = graph.extract_string(impl_selector_str); // Create temporary packed int8 tensors for input and output TmpTensor packed_int8_input( - &graph, graph.sizes_of(fp_input), vkapi::kInt8x4, utils::kBuffer, layout); + &graph, + graph.sizes_of(fp_input), + vkapi::kInt8x4, + utils::kBuffer, + input_layout); TmpTensor packed_int8_output( &graph, graph.sizes_of(fp_output), vkapi::kInt8x4, utils::kBuffer, - layout); + output_layout); // Quantize floating point input to packed int8 add_q8ta_quantize_node( @@ -186,8 +322,40 @@ void test_q8ta_conv2d(ComputeGraph& graph, const std::vector& args) { groups, activation, packed_int8_output}; - if (impl_selector == "im2col") { - VK_GET_OP_FN("et_vk.q8ta_conv2d_im2col.default")(graph, conv_args); + if (impl_selector == "im2col_fallback") { + // Simulate a device without dot-product support so the fallback + // kernel is selected on any hardware. + vkapi::ScopedAdapterCapabilityOverride no_dot_support( + graph.context()->adapter_ptr(), + vkapi::AdapterCapabilityOverrides::without_dot_product_support()); + q8ta_conv2d_im2col_impl(graph, /*use_unsigned_dot=*/false, conv_args); + assert_im2col_kernel_selection( + graph, + /*expect_unsigned=*/false, + /*expect_buffer_weights=*/false, + /*expect_fallback_kernel=*/true); + } else if ( + impl_selector == "im2col" || impl_selector == "im2col_unsigned" || + impl_selector == "im2col_auto") { + const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); + bool expect_unsigned = impl_selector == "im2col_unsigned"; + if (impl_selector == "im2col_auto") { + VK_GET_OP_FN("et_vk.q8ta_conv2d_im2col.default")(graph, conv_args); + expect_unsigned = can_use_unsigned_pw_dot( + *adapter, graph.size_at(-1, weight_data)); + } else { + q8ta_conv2d_im2col_impl(graph, expect_unsigned, conv_args); + } + const int64_t packed_height = + utils::div_up_4(graph.size_at(-1, weight_data)); + const int64_t packed_width = + utils::div_up_4(graph.size_at(-2, weight_data)) * 4; + const int64_t max_texture_extent = adapter->max_texture2d_dim(); + const bool expect_buffer_weights = + packed_width > max_texture_extent * 4 || + packed_height > max_texture_extent; + assert_im2col_kernel_selection( + graph, expect_unsigned, expect_buffer_weights); } else if (impl_selector == "general") { VK_GET_OP_FN("et_vk.q8ta_conv2d_general.default")(graph, conv_args); } else { @@ -287,7 +455,29 @@ void test_q8ta_conv2d_pw( groups, activation, packed_int8_output}; - VK_GET_OP_FN("et_vk.q8ta_conv2d_pw.default")(graph, conv_args); + if (impl_selector == "pw_signed" || impl_selector == "pw_unsigned" || + impl_selector == "pw_auto") { + const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); + bool expect_unsigned = impl_selector == "pw_unsigned"; + if (impl_selector == "pw_auto") { + VK_GET_OP_FN("et_vk.q8ta_conv2d_pw.default")(graph, conv_args); + expect_unsigned = can_use_unsigned_pw_dot( + *adapter, graph.size_at(-1, weight_data)); + } else { + q8ta_conv2d_pw_impl(graph, expect_unsigned, conv_args); + } + const int64_t packed_height = + utils::div_up_4(graph.size_at(-1, weight_data)); + const int64_t packed_width = + utils::div_up_4(graph.size_at(-2, weight_data)) * 4; + const int64_t max_texture_extent = adapter->max_texture2d_dim(); + const bool expect_buffer_weights = + packed_width > max_texture_extent * 4 || + packed_height > max_texture_extent; + assert_pw_kernel_selection(graph, expect_unsigned, expect_buffer_weights); + } else { + VK_GET_OP_FN("et_vk.q8ta_conv2d_pw.default")(graph, conv_args); + } } // Dequantize packed int8 output to floating point diff --git a/backends/vulkan/test/custom_ops/q8ta_conv2d_route_test.cpp b/backends/vulkan/test/custom_ops/q8ta_conv2d_route_test.cpp new file mode 100644 index 00000000000..63d154b170f --- /dev/null +++ b/backends/vulkan/test/custom_ops/q8ta_conv2d_route_test.cpp @@ -0,0 +1,473 @@ +// 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. + +#include + +#include + +#include + +namespace vkcompute { +namespace { + +constexpr uint64_t kLargeDeviceBuffer = 1ULL << 32; + +Q8taConv2dRouteParams make_mali_grouped_params() { + return { + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/2, + /*in_channels_per_group=*/32, + /*out_channels=*/64, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/128, + /*out_width=*/128, + }; +} + +TEST(Q8taConv2dRouteTest, RoutesBatchedRegularMaliConvolutionToIm2Col) { + EXPECT_TRUE(should_use_q8ta_conv2d_im2col({ + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/64, + /*out_channels=*/128, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/20, + /*out_width=*/26, + })); + + EXPECT_TRUE(should_use_q8ta_conv2d_im2col({ + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/128, + /*out_channels=*/256, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/10, + /*out_width=*/13, + })); +} + +TEST(Q8taConv2dRouteTest, RoutesMeasuredSceneXGroupedConvolutionsToIm2Col) { + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(make_mali_grouped_params())); + + Q8taConv2dRouteParams params = make_mali_grouped_params(); + params.groups = 4; + params.in_channels_per_group = 32; + params.out_channels = 128; + params.kernel_height = 5; + params.kernel_width = 5; + params.out_height = 64; + params.out_width = 64; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + params = make_mali_grouped_params(); + params.out_height = 64; + params.out_width = 64; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + params = make_mali_grouped_params(); + params.groups = 3; + params.in_channels_per_group = 32; + params.out_channels = 96; + params.kernel_height = 4; + params.kernel_width = 4; + params.out_height = 64; + params.out_width = 64; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RoutesPreviouslyOutOfEnvelopeShapesToIm2Col) { + Q8taConv2dRouteParams params = make_mali_grouped_params(); + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + params.is_mali = false; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.is_mali = true; + params.supports_int8_dot_product = false; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.supports_int8_dot_product = true; + + // No group-count gate: regular and wide grouped shapes route alike, as + // long as each group owns whole packed-4 output blocks. + params.groups = 1; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.groups = 5; + params.out_channels = 80; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.groups = 2; + params.out_channels = 64; + + // No kernel squareness/size gate. + params.kernel_width = 5; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.kernel_height = 2; + params.kernel_width = 2; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.kernel_height = 6; + params.kernel_width = 6; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.kernel_height = 3; + params.kernel_width = 3; + + // Packing alignment is still required. + params.in_channels_per_group = 31; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.in_channels_per_group = 34; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.in_channels_per_group = 32; + + // Output channels indivisible by groups fail closed. + params.out_channels = 63; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.out_channels = 64; + // No spatial window gates. + params.out_height = 63; + params.out_width = 64; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.out_height = 128; + params.out_width = 129; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RejectsMisalignedGroupedOutputChannels) { + // The PW GEMM derives group_idx = oc_block / OC4_per_group, so a group + // with a non-multiple-of-4 channel count would straddle packed blocks. + Q8taConv2dRouteParams params = make_mali_grouped_params(); + params.groups = 3; + params.out_channels = 66; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.groups = 5; + params.out_channels = 65; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.out_channels = 80; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RejectsMisalignedGroupedSingleBatchMali) { + // The single-batch Mali branch dispatches the same grouped PW shader, so + // the packed-4 output-block requirement binds there too. + Q8taConv2dRouteParams params = make_mali_grouped_params(); + params.batch = 1; + params.groups = 2; + params.out_channels = 66; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.out_channels = 64; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RoutesSingleTileScratchToIm2ColOffTileToDirect) { + // The legacy full-scratch envelope (up to 32MiB) narrowed to a single + // 16MiB streaming tile. With K=1152 and 8x8 output, batch 227 needs + // 15.96MiB in one tile (routes im2col) while batch 228 needs 16.03MiB in + // two tiles (stays direct on non-Mali, though both fit the old budget). + auto make_params = [](int64_t batch) { + return Q8taConv2dRouteParams{ + /*is_mali=*/false, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + batch, + /*groups=*/1, + /*in_channels_per_group=*/128, + /*out_channels=*/128, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/8, + /*out_width=*/8, + }; + }; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(make_params(227))); + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(make_params(228))); +} + +TEST(Q8taConv2dRouteTest, RespectsMaliDeviceBufferLimitForGroupedShapes) { + Q8taConv2dRouteParams params = make_mali_grouped_params(); + params.groups = 4; + params.out_channels = 128; + params.kernel_height = 5; + params.kernel_width = 5; + params.out_height = 64; + params.out_width = 64; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + // One im2col row is 3200 kernel columns x 64 aligned width bytes. + constexpr uint64_t kBytesPerRow = 3200 * 64; + params.max_buffer_bytes = kBytesPerRow; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.max_buffer_bytes = kBytesPerRow - 1; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RejectsOverflowingMaliGroupedGeometry) { + Q8taConv2dRouteParams params = make_mali_grouped_params(); + params.groups = 4; + params.out_channels = 128; + // max()/36 + 2 is 0 (mod 4), so it passes the alignment gate and reaches + // the grouped overflow guard: align(C*3*3) > max()/4 with groups == 4. + params.in_channels_per_group = std::numeric_limits::max() / 36 + 2; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + + params = make_mali_grouped_params(); + params.out_height = std::numeric_limits::max(); + params.out_width = 2; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RejectsIneligibleBatchedMaliConvolutions) { + Q8taConv2dRouteParams params = { + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/64, + /*out_channels=*/128, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/20, + /*out_width=*/26, + }; + + // Grouped and pointwise batched shapes route alike on Mali. + params.groups = 2; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.groups = 1; + params.kernel_height = 1; + params.kernel_width = 1; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.kernel_height = 3; + params.kernel_width = 3; + + // Packing alignment is still required. + params.in_channels_per_group = 62; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.in_channels_per_group = 64; + + // Kernels past the unsigned-dot accumulator bound stay direct. + params.in_channels_per_group = 4096; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RoutesSmallAndLargeBatchedShapesToIm2Col) { + Q8taConv2dRouteParams params = { + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/32, + /*out_channels=*/64, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/20, + /*out_width=*/26, + }; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + // No kernel-size gate. + params.in_channels_per_group = 64; + params.kernel_height = 2; + params.kernel_width = 2; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.in_channels_per_group = 32; + params.kernel_height = 3; + params.kernel_width = 3; + + params.supports_int8_dot_product = false; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.supports_int8_dot_product = true; + params.in_channels_per_group = 30; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.in_channels_per_group = 32; + + // No output-channel or spatial window gates. + params.out_channels = 63; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.out_channels = 64; + params.out_height = 8; + params.out_width = 16; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.out_height = 8; + params.out_width = 15; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.out_height = 32; + params.out_width = 32; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.out_height = 1; + params.out_width = 521; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.out_height = 20; + params.out_width = 26; + + // No pointwise exclusion. + params.in_channels_per_group = 256; + params.kernel_height = 1; + params.kernel_width = 1; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RespectsMaliDeviceBufferLimit) { + EXPECT_FALSE(should_use_q8ta_conv2d_im2col({ + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + /*max_buffer_bytes=*/15 * 1024, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/64, + /*out_channels=*/128, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/20, + /*out_width=*/26, + })); + + // One im2col row is 576 kernel columns x 28 aligned width bytes; the + // device plan is feasible down to exactly that budget. + constexpr uint64_t kBytesPerRow = 576 * 28; + Q8taConv2dRouteParams params = { + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + /*max_buffer_bytes=*/kBytesPerRow, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/64, + /*out_channels=*/128, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/20, + /*out_width=*/26, + }; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.max_buffer_bytes = kBytesPerRow - 1; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, PreservesNonMaliBatchedHeuristic) { + Q8taConv2dRouteParams params = { + /*is_mali=*/false, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/128, + /*out_channels=*/256, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/10, + /*out_width=*/13, + }; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + + params.out_height = 8; + params.out_width = 8; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + params.is_mali = true; + params.supports_int8_dot_product = false; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + params.is_mali = false; + params.max_buffer_bytes = 1024; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, LegacyBatchedRoutePrecedesMaliExtensionGates) { + EXPECT_TRUE(should_use_q8ta_conv2d_im2col({ + /*is_mali=*/true, + /*supports_int8_dot_product=*/false, + kLargeDeviceBuffer, + /*batch=*/2, + /*groups=*/1, + /*in_channels_per_group=*/1024, + /*out_channels=*/32, + /*kernel_height=*/1, + /*kernel_width=*/1, + /*out_height=*/8, + /*out_width=*/8, + })); +} + +TEST(Q8taConv2dRouteTest, PreservesSingleBatchPolicy) { + Q8taConv2dRouteParams params = { + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/1, + /*groups=*/4, + /*in_channels_per_group=*/8, + /*out_channels=*/32, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/128, + /*out_width=*/128, + }; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + params.is_mali = false; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.groups = 1; + params.in_channels_per_group = 32; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RejectsInvalidAndOverflowingGeometry) { + Q8taConv2dRouteParams params = { + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/64, + /*out_channels=*/128, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/20, + /*out_width=*/26, + }; + + params.groups = 0; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.groups = 1; + params.batch = 0; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.batch = 60; + params.out_channels = 0; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.out_channels = 128; + params.out_height = std::numeric_limits::max(); + params.out_width = 2; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + + params.out_height = 20; + params.out_width = 26; + params.in_channels_per_group = std::numeric_limits::max(); + params.kernel_height = 2; + params.kernel_width = 1; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + + params.in_channels_per_group = 4; + params.kernel_height = std::numeric_limits::max() / 4; + params.kernel_width = 2; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + + params.in_channels_per_group = std::numeric_limits::max() - 2; + params.kernel_height = 1; + params.kernel_width = 1; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); +} + +} // namespace +} // namespace vkcompute diff --git a/backends/vulkan/test/custom_ops/q8ta_conv2d_stream_plan_test.cpp b/backends/vulkan/test/custom_ops/q8ta_conv2d_stream_plan_test.cpp new file mode 100644 index 00000000000..6f634a83c27 --- /dev/null +++ b/backends/vulkan/test/custom_ops/q8ta_conv2d_stream_plan_test.cpp @@ -0,0 +1,178 @@ +// 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. + +#include + +#include + +#include + +namespace vkcompute { +namespace { + +constexpr int64_t kEightMiB = 8 * 1024 * 1024; +constexpr int64_t kSixteenMiB = 16 * 1024 * 1024; + +TEST(Q8taConv2dStreamPlanTest, SplitsFirstSceneXConvolution) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/60, + /*flattened_kernel_size=*/576, + /*out_height=*/20, + /*out_width=*/26, + kEightMiB); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.aligned_out_width, 28); + EXPECT_EQ(plan.rows_per_tile, 400); + EXPECT_EQ(plan.num_tiles, 3); + EXPECT_EQ(plan.scratch_bytes, 6451200); +} + +TEST(Q8taConv2dStreamPlanTest, SplitsSecondSceneXConvolution) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/60, + /*flattened_kernel_size=*/1152, + /*out_height=*/10, + /*out_width=*/13, + kEightMiB); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.aligned_out_width, 16); + EXPECT_EQ(plan.rows_per_tile, 300); + EXPECT_EQ(plan.num_tiles, 2); + EXPECT_EQ(plan.scratch_bytes, 5529600); +} + +TEST(Q8taConv2dStreamPlanTest, UsesOneTileWhenFullScratchFits) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/2, + /*flattened_kernel_size=*/288, + /*out_height=*/7, + /*out_width=*/7, + kEightMiB); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.rows_per_tile, 14); + EXPECT_EQ(plan.num_tiles, 1); + EXPECT_EQ(plan.scratch_bytes, 32256); +} + +TEST(Q8taConv2dStreamPlanTest, SelectsFullFitBelowProductionBudget) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/10, + /*flattened_kernel_size=*/288, + /*out_height=*/30, + /*out_width=*/99, + kSixteenMiB); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.rows_per_tile, 300); + EXPECT_EQ(plan.num_tiles, 1); + EXPECT_EQ(plan.scratch_bytes, 8640000); +} + +TEST(Q8taConv2dStreamPlanTest, SelectsStreamingAboveProductionBudget) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/10, + /*flattened_kernel_size=*/576, + /*out_height=*/30, + /*out_width=*/99, + kSixteenMiB); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.rows_per_tile, 150); + EXPECT_EQ(plan.num_tiles, 2); + EXPECT_EQ(plan.scratch_bytes, 8640000); +} + +TEST(Q8taConv2dStreamPlanTest, RejectsOneRowLargerThanBudget) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/1, + /*flattened_kernel_size=*/16384, + /*out_height=*/1, + /*out_width=*/513, + kEightMiB); + + EXPECT_FALSE(plan.feasible); + EXPECT_EQ(plan.rows_per_tile, 0); + EXPECT_EQ(plan.num_tiles, 0); + EXPECT_EQ(plan.scratch_bytes, 0); +} + +TEST(Q8taConv2dStreamPlanTest, RejectsOutputWidthAlignmentOverflow) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/1, + /*flattened_kernel_size=*/1, + /*out_height=*/1, + /*out_width=*/std::numeric_limits::max(), + std::numeric_limits::max()); + + EXPECT_FALSE(plan.feasible); +} + +TEST(Q8taConv2dStreamPlanTest, HandlesTileCountCeilDivisionAtShaderLimit) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/std::numeric_limits::max(), + /*flattened_kernel_size=*/1, + /*out_height=*/1, + /*out_width=*/1, + /*scratch_budget_bytes=*/8); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.rows_per_tile, 2); + EXPECT_EQ(plan.num_tiles, 1073741824); + EXPECT_EQ(plan.scratch_bytes, 8); +} + +TEST(Q8taConv2dStreamPlanTest, RejectsRowOffsetBeyondShaderIndexRange) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/std::numeric_limits::max(), + /*flattened_kernel_size=*/1, + /*out_height=*/1, + /*out_width=*/1, + std::numeric_limits::max()); + + EXPECT_FALSE(plan.feasible); +} + +TEST(Q8taConv2dStreamPlanTest, CapsRowsAtGuaranteedWorkgroupCountZ) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/100000, + /*flattened_kernel_size=*/1, + /*out_height=*/1, + /*out_width=*/1, + std::numeric_limits::max()); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.num_tiles, 2); + EXPECT_EQ(plan.rows_per_tile, 50000); + EXPECT_EQ(plan.scratch_bytes, 200000); +} + +TEST(Q8taConv2dStreamPlanTest, DeviceBufferLimitMatchesActualPlan) { + constexpr int64_t kBytesPerRow = 576 * 28; + const auto rejected = make_q8ta_conv2d_stream_plan_for_device( + /*batch=*/60, + /*flattened_kernel_size=*/576, + /*out_height=*/20, + /*out_width=*/26, + /*max_buffer_bytes=*/kBytesPerRow - 1); + EXPECT_FALSE(rejected.feasible); + + const auto accepted = make_q8ta_conv2d_stream_plan_for_device( + /*batch=*/60, + /*flattened_kernel_size=*/576, + /*out_height=*/20, + /*out_width=*/26, + /*max_buffer_bytes=*/kBytesPerRow); + EXPECT_TRUE(accepted.feasible); + EXPECT_EQ(accepted.rows_per_tile, 1); + EXPECT_EQ(accepted.num_tiles, 1200); + EXPECT_EQ(accepted.scratch_bytes, kBytesPerRow); +} + +} // namespace +} // namespace vkcompute diff --git a/backends/vulkan/test/custom_ops/targets.bzl b/backends/vulkan/test/custom_ops/targets.bzl index 5d1045173a2..1895180495a 100644 --- a/backends/vulkan/test/custom_ops/targets.bzl +++ b/backends/vulkan/test/custom_ops/targets.bzl @@ -1,4 +1,4 @@ -load("@fbsource//tools/build_defs:platform_defs.bzl", "ANDROID") +load("@fbsource//tools/build_defs:platform_defs.bzl", "ANDROID", "CXX") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load( "@fbsource//xplat/executorch/backends/vulkan:targets.bzl", @@ -85,6 +85,39 @@ def define_common_targets(is_fbcode = False): link_whole = True, ) + runtime.cxx_test( + name = "utils_test", + srcs = [ + "utils_test.cpp", + ], + contacts = ["oncall+ai_infra_mobile_platform@xmail.facebook.com"], + platforms = [CXX], + deps = [ + ":prototyping_utils", + "//third-party/googletest:gtest_main", + ], + ) + + runtime.cxx_test( + name = "q8ta_conv2d_stream_plan_test", + srcs = ["q8ta_conv2d_stream_plan_test.cpp"], + platforms = get_platforms(), + deps = [ + "//third-party/googletest:gtest_main", + "//executorch/backends/vulkan:vulkan_graph_runtime", + ], + ) + + runtime.cxx_test( + name = "q8ta_conv2d_route_test", + srcs = ["q8ta_conv2d_route_test.cpp"], + platforms = get_platforms(), + deps = [ + "//third-party/googletest:gtest_main", + "//executorch/backends/vulkan:vulkan_graph_runtime", + ], + ) + define_custom_op_test_binary("test_add") define_custom_op_test_binary("test_q8csw_linear") define_custom_op_test_binary("test_q8csw_conv2d") diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp index b30212feb72..90fd4e1616b 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp @@ -4,10 +4,15 @@ // This source code is licensed under the BSD-style license found in the // LICENSE file in the root directory of this source tree. +#include +#include +#include #include +#include #include #include +#include #include #include @@ -22,14 +27,29 @@ using namespace executorch::vulkan::prototyping; using namespace vkcompute; static constexpr int64_t kRefDimSizeLimit = 100; +static constexpr int64_t kRefOperationLimit = 2 * 1024 * 1024; + +struct Im2colUnsignedTestOptions { + int32_t input_zero_point = 2; + int32_t output_zero_point = -1; + bool use_extreme_values = false; + bool use_accumulator_limit_values = false; + bool has_bias = true; + const char* activation = "none"; + float weight_scale = 1.0f / 256.0f; +}; // Utility function to create a test case from a Conv2dConfig -static TestCase create_test_case_from_config( +static TestCase create_test_case_from_config_with_layouts( const Conv2dConfig& config, vkapi::ScalarType input_dtype, utils::StorageType fp_storage_type, - utils::GPUMemoryLayout int8_memory_layout, - const std::string& impl_selector = "") { + utils::GPUMemoryLayout input_int8_memory_layout, + utils::GPUMemoryLayout output_int8_memory_layout, + const std::string& impl_selector = "", + const Im2colUnsignedTestOptions* im2col_options = nullptr, + const float input_scale_val = 0.008123f, + const DataGenType input_data_gen = DataGenType::RANDOM) { TestCase test_case; // Calculate output dimensions @@ -62,7 +82,10 @@ static TestCase create_test_case_from_config( std::to_string(config.stride.h) + " p" + std::to_string(config.padding.h) + " d" + std::to_string(config.dilation.h) + " g" + std::to_string(config.groups); - std::string storage_str = repr_str(utils::kBuffer, int8_memory_layout); + std::string storage_str = repr_str(utils::kBuffer, input_int8_memory_layout); + if (input_int8_memory_layout != output_int8_memory_layout) { + storage_str += "->" + repr_str(utils::kBuffer, output_int8_memory_layout); + } std::string suffix = impl_selector.empty() ? "" : "[" + impl_selector + "]"; std::string test_name = make_test_label( prefix, dtype_str, dtype_str, shape_str, storage_str, suffix); @@ -77,23 +100,36 @@ static TestCase create_test_case_from_config( input_dtype, fp_storage_type, fp_memory_layout, -#ifdef DEBUG_MODE - DataGenType::RANDOM -#else - DataGenType::RANDOM -#endif - ); + input_data_gen); if (debugging()) { print_valuespec_data(input_tensor, "input_tensor"); } - float input_scale_val = 0.008123; ValueSpec input_scale(input_scale_val); - int32_t input_zero_point_val = 2; + const int32_t input_zero_point_val = + im2col_options == nullptr ? 2 : im2col_options->input_zero_point; ValueSpec input_zero_point(input_zero_point_val); + if (im2col_options != nullptr && + im2col_options->use_accumulator_limit_values) { + input_tensor.ensure_data_generated(2401); + std::fill( + input_tensor.get_float_data().begin(), + input_tensor.get_float_data().end(), + (127.0f - input_zero_point_val) * input_scale_val); + } else if (im2col_options != nullptr && im2col_options->use_extreme_values) { + input_tensor.ensure_data_generated(2401); + constexpr std::array values = {-128, -1, 0, 1, 127}; + std::vector& input_data = input_tensor.get_float_data(); + for (size_t i = 0; i < input_data.size(); ++i) { + input_data.at(i) = (static_cast(values.at(i % values.size())) - + input_zero_point_val) * + input_scale_val; + } + } + // Quantized weight tensor (int8) - [C_out, C_in_per_group * K_h * K_w] // Memory layout: height, width, then channels - in_c is innermost (stride 1) // in the second dimension @@ -109,6 +145,22 @@ static TestCase create_test_case_from_config( DataGenType::RANDINT8); quantized_weight.set_constant(true); + if (im2col_options != nullptr && + im2col_options->use_accumulator_limit_values) { + quantized_weight.ensure_data_generated(2402); + std::fill( + quantized_weight.get_int8_data().begin(), + quantized_weight.get_int8_data().end(), + 127); + } else if (im2col_options != nullptr && im2col_options->use_extreme_values) { + quantized_weight.ensure_data_generated(2402); + constexpr std::array values = {-128, -1, 0, 1, 127}; + std::vector& weight_data = quantized_weight.get_int8_data(); + for (size_t i = 0; i < weight_data.size(); ++i) { + weight_data.at(i) = values.at((i * 3 + 1) % values.size()); + } + } + if (debugging()) { print_valuespec_data(quantized_weight, "weight_tensor"); } @@ -123,6 +175,13 @@ static TestCase create_test_case_from_config( utils::kWidthPacked, DataGenType::RANDOM_SCALES); weight_scales.set_constant(true); + if (im2col_options != nullptr) { + weight_scales.ensure_data_generated(2403); + std::fill( + weight_scales.get_float_data().begin(), + weight_scales.get_float_data().end(), + im2col_options->weight_scale); + } ValueSpec weight_sums( {aligned_out_channels}, // Per output channel @@ -144,12 +203,16 @@ static TestCase create_test_case_from_config( utils::kWidthPacked, DataGenType::ZEROS); bias.set_constant(true); + if (im2col_options != nullptr && !im2col_options->has_bias) { + bias.set_none(true); + } // Output quantization parameters float output_scale_val = 0.05314; ValueSpec output_scale(output_scale_val); - int32_t output_zero_point_val = -1; + const int32_t output_zero_point_val = + im2col_options == nullptr ? -1 : im2col_options->output_zero_point; ValueSpec output_zero_point(output_zero_point_val); // Stride and padding parameters @@ -188,12 +251,16 @@ static TestCase create_test_case_from_config( test_case.add_input_spec(groups); // Activation (none = no activation) - ValueSpec activation = ValueSpec::make_string("none"); + ValueSpec activation = ValueSpec::make_string( + im2col_options == nullptr ? "none" : im2col_options->activation); test_case.add_input_spec(activation); // Add memory layout parameter for the quantized tensors - ValueSpec layout_int(static_cast(int8_memory_layout)); - test_case.add_input_spec(layout_int); + ValueSpec input_layout_int(static_cast(input_int8_memory_layout)); + test_case.add_input_spec(input_layout_int); + + ValueSpec output_layout_int(static_cast(output_int8_memory_layout)); + test_case.add_input_spec(output_layout_int); // Add impl_selector string ValueSpec impl_selector_spec = ValueSpec::make_string(impl_selector); @@ -201,7 +268,9 @@ static TestCase create_test_case_from_config( test_case.add_output_spec(output); - test_case.set_abs_tolerance(output_scale_val + 1e-4f); + test_case.set_abs_tolerance( + im2col_options == nullptr ? output_scale_val + 1e-4f + : output_scale_val * 0.25f); // Filter out quantize/dequantize shaders from timing measurements test_case.set_shader_filter({ @@ -214,6 +283,61 @@ static TestCase create_test_case_from_config( return test_case; } +static TestCase create_test_case_from_config( + const Conv2dConfig& config, + vkapi::ScalarType input_dtype, + utils::StorageType fp_storage_type, + utils::GPUMemoryLayout int8_memory_layout, + const std::string& impl_selector = "", + const Im2colUnsignedTestOptions* im2col_options = nullptr, + const float input_scale_val = 0.008123f, + const DataGenType input_data_gen = DataGenType::RANDOM) { + return create_test_case_from_config_with_layouts( + config, + input_dtype, + fp_storage_type, + int8_memory_layout, + int8_memory_layout, + impl_selector, + im2col_options, + input_scale_val, + input_data_gen); +} + +static std::vector generate_narrow_workgroup_test_cases() { + std::vector test_cases; + std::vector configs = { + {OutInChannels(64, 32), + InputSize2D(9, 9), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1}, + {OutInChannels(128, 32), + InputSize2D(7, 7), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1}, + }; + + for (auto& config : configs) { + const bool is_performance = config.channels.out > kRefDimSizeLimit; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = make_test_case_name( + config, is_performance, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4C, + /*impl_selector=*/"general")); + } + return test_cases; +} + // Generate easy test cases for quantized conv2d operation (for debugging) std::vector generate_quantized_conv2d_easy_cases() { std::vector test_cases; @@ -271,10 +395,226 @@ std::vector generate_quantized_conv2d_easy_cases() { return test_cases; } +static std::vector generate_im2col_unsigned_test_cases( + const std::string& impl_selector); + +static std::vector generate_streaming_im2col_test_cases() { + std::vector test_cases; + + Conv2dConfig full_fit_config = { + OutInChannels(4, 32), + InputSize2D(30, 99), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1, + 10}; + full_fit_config.op_name = "conv2d_q8ta_q8csw_q8to"; + full_fit_config.test_case_name = make_test_case_name( + full_fit_config, false, utils::kTexture3D, utils::kBuffer); + + Conv2dConfig streaming_fallback_config = full_fit_config; + streaming_fallback_config.channels.in = 64; + streaming_fallback_config.test_case_name = make_test_case_name( + streaming_fallback_config, false, utils::kTexture3D, utils::kBuffer); + + // Forced-fallback cases need no int8 dot-product support, so they run on + // all devices; the kAuto cases below stay gated. + test_cases.push_back(create_test_case_from_config( + full_fit_config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + /*impl_selector=*/"im2col_fallback", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT)); + test_cases.push_back(create_test_case_from_config( + streaming_fallback_config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + /*impl_selector=*/"im2col_fallback", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT)); + + if (!vkcompute::api::context()->adapter_ptr()->supports_int8_dot_product()) { + return test_cases; + } + + for (const utils::GPUMemoryLayout layout : + {utils::kPackedInt8_4C1W, utils::kPackedInt8_4W4C}) { + test_cases.push_back(create_test_case_from_config( + full_fit_config, + vkapi::kFloat, + utils::kTexture3D, + layout, + /*impl_selector=*/"im2col_auto", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT)); + } + Conv2dConfig streaming_config = full_fit_config; + streaming_config.channels.in = 64; + streaming_config.test_case_name = make_test_case_name( + streaming_config, false, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + streaming_config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4C1W, + /*impl_selector=*/"im2col_auto", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT)); + Conv2dConfig grouped_config = { + OutInChannels(16, 32), + InputSize2D(30, 99), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 2, + 20}; + grouped_config.op_name = "conv2d_q8ta_q8csw_q8to"; + grouped_config.test_case_name = make_test_case_name( + grouped_config, false, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + grouped_config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + /*impl_selector=*/"im2col_auto", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT)); + + Conv2dConfig output_channel_tail_config = { + OutInChannels(10, 32), + InputSize2D(30, 99), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1, + 20}; + output_channel_tail_config.op_name = "conv2d_q8ta_q8csw_q8to"; + output_channel_tail_config.test_case_name = make_test_case_name( + output_channel_tail_config, false, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + output_channel_tail_config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + /*impl_selector=*/"im2col_auto", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT)); + return test_cases; +} + +// SceneX route tests. The kPackedInt8_4C input + kPackedInt8_4W4C output +// layout combination is only exercised here (default generators never pair +// them); run via --scenex-regular . Zero +// tolerances are exact by construction: both pipelines accumulate in int32 +// with identical requantize, so any mismatch is a real regression, not noise. +static TestCase create_scenex_test_case( + const Conv2dConfig& source_config, + const std::string& route) { + Conv2dConfig config = source_config; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = + make_test_case_name(config, false, utils::kTexture3D, utils::kBuffer); + + const std::string impl_selector = route == "auto" ? "" + : route == "direct" ? "general" + : "im2col"; + TestCase test_case = create_test_case_from_config_with_layouts( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4C, + utils::kPackedInt8_4W4C, + impl_selector, + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT); + test_case.set_abs_tolerance(0.0f); + test_case.set_rel_tolerance(0.0f); + return test_case; +} + +static TestCase generate_scenex_regular_test_case( + const std::string& route, + const int case_index) { + const std::vector configs = { + {OutInChannels(128, 64), + InputSize2D(40, 51), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1, + 60}, + {OutInChannels(256, 128), + InputSize2D(20, 26), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1, + 60}, + }; + return create_scenex_test_case(configs.at(case_index), route); +} + +static TestCase generate_scenex_grouped_test_case( + const std::string& route, + const int case_index) { + const std::vector configs = { + {OutInChannels(64, 64), + InputSize2D(128, 128), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 2, + 60}, + {OutInChannels(128, 128), + InputSize2D(128, 128), + KernelSize(5, 5), + Stride(2, 2), + Padding(2, 2), + Dilation(1, 1), + 4, + 60}, + {OutInChannels(64, 64), + InputSize2D(64, 64), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 2, + 60}, + }; + + return create_scenex_test_case(configs.at(case_index), route); +} + // Generate test cases for quantized conv2d operation static std::vector generate_quantized_conv2d_test_cases() { std::vector test_cases; - if (!vkcompute::api::context()->adapter_ptr()->supports_int8_dot_product()) { + api::Context* const context = vkcompute::api::context(); + if (!context->adapter_ptr()->supports_int8_dot_product()) { + for (const std::string& impl_selector : {"im2col", "im2col_auto"}) { + std::vector im2col_cases = + generate_im2col_unsigned_test_cases(impl_selector); + for (TestCase& test_case : im2col_cases) { + test_cases.push_back(std::move(test_case)); + } + } return test_cases; } @@ -542,6 +882,181 @@ static std::vector generate_quantized_conv2d_test_cases() { } } + for (const std::string& impl_selector : + {"im2col", "im2col_unsigned", "im2col_auto"}) { + std::vector im2col_cases = + generate_im2col_unsigned_test_cases(impl_selector); + for (TestCase& test_case : im2col_cases) { + test_cases.push_back(std::move(test_case)); + } + } + + auto narrow_workgroup_cases = generate_narrow_workgroup_test_cases(); + test_cases.insert( + test_cases.end(), + narrow_workgroup_cases.begin(), + narrow_workgroup_cases.end()); + + auto streaming_cases = generate_streaming_im2col_test_cases(); + test_cases.insert( + test_cases.end(), streaming_cases.begin(), streaming_cases.end()); + + return test_cases; +} + +static std::vector generate_im2col_unsigned_test_cases( + const std::string& impl_selector) { + api::Context* const context = vkcompute::api::context(); + if (impl_selector == "im2col_unsigned" && + !context->adapter_ptr()->supports_int8_dot_product()) { + return {}; + } + + std::vector> configs = { + {{OutInChannels(5, 4), + InputSize2D(5, 5), + KernelSize(3, 3), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1}, + {.input_zero_point = 2, + .output_zero_point = -1, + .use_extreme_values = true, + .use_accumulator_limit_values = false, + .has_bias = true, + .activation = "none"}}, + {{OutInChannels(8, 8), + InputSize2D(5, 5), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1}, + {.input_zero_point = -7, + .output_zero_point = 3, + .use_extreme_values = true, + .use_accumulator_limit_values = false, + .has_bias = false, + .activation = "none"}}, + {{OutInChannels(12, 8), + InputSize2D(7, 7), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1}, + {.input_zero_point = 127, + .output_zero_point = -5, + .use_extreme_values = true, + .use_accumulator_limit_values = false, + .has_bias = true, + .activation = "relu"}}, + {{OutInChannels(12, 8), + InputSize2D(9, 9), + KernelSize(3, 3), + Stride(1, 1), + Padding(2, 2), + Dilation(2, 2), + 1}, + {.input_zero_point = -128, + .output_zero_point = 5, + .use_extreme_values = true, + .use_accumulator_limit_values = false, + .has_bias = false, + .activation = "none"}}, + {{OutInChannels(8, 8), + InputSize2D(6, 7), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 2}, + {.input_zero_point = 11, + .output_zero_point = -3, + .use_extreme_values = true, + .use_accumulator_limit_values = false, + .has_bias = true, + .activation = "relu"}}, + {{OutInChannels(1, 4), + InputSize2D(90, 91), + KernelSize(90, 91), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1}, + {.input_zero_point = 0, + .output_zero_point = -1, + .use_extreme_values = false, + .use_accumulator_limit_values = true, + .has_bias = true, + .activation = "none", + .weight_scale = 1.0f / 1000000.0f}}, + }; + + std::vector test_cases; + test_cases.reserve(configs.size()); + for (auto& [config, options] : configs) { + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = + make_test_case_name(config, false, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + impl_selector, + &options)); + } + + if (impl_selector == "im2col_auto") { + Conv2dConfig config{ + OutInChannels(1, 4), + InputSize2D(91, 91), + KernelSize(91, 91), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1}; + Im2colUnsignedTestOptions options; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = + make_test_case_name(config, false, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + impl_selector, + &options)); + } + + const vkapi::Adapter& adapter = *context->adapter_ptr(); + if (impl_selector == "im2col_unsigned" || + (impl_selector == "im2col_auto" && can_use_unsigned_pw_dot(adapter, 4))) { + const int32_t buffer_output_channels = utils::safe_downcast( + static_cast(adapter.max_texture2d_dim()) * 4 + 1); + Conv2dConfig config{ + OutInChannels(buffer_output_channels, 4), + InputSize2D(1, 1), + KernelSize(1, 1), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1}; + Im2colUnsignedTestOptions options; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = + make_test_case_name(config, false, utils::kBuffer, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, + vkapi::kFloat, + utils::kBuffer, + utils::kPackedInt8_4W4C, + impl_selector, + &options)); + } + return test_cases; } @@ -565,9 +1080,10 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { const ValueSpec& dilation_spec = test_case.inputs()[idx++]; const ValueSpec& groups_spec = test_case.inputs()[idx++]; const ValueSpec& activation_spec = test_case.inputs()[idx++]; - (void)activation_spec; // Not used in reference implementation const ValueSpec& layout_spec = test_case.inputs()[idx++]; (void)layout_spec; // Not used in reference implementation + const ValueSpec& output_layout_spec = test_case.inputs()[idx++]; + (void)output_layout_spec; // Not used in reference implementation const ValueSpec& impl_selector_spec = test_case.inputs()[idx++]; (void)impl_selector_spec; // Not used in reference implementation @@ -606,10 +1122,12 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { int64_t dilation_w = dilation_data[1]; int64_t groups = groups_spec.get_int_value(); - // Skip for large tensors since computation time will be extremely slow - if (N > kRefDimSizeLimit || C_in > kRefDimSizeLimit || - H_in > kRefDimSizeLimit || W_in > kRefDimSizeLimit || - C_out > kRefDimSizeLimit) { + const int64_t reference_operations = + N * C_out * H_out * W_out * (C_in / groups) * K_h * K_w; + const bool has_large_dimension = N > kRefDimSizeLimit || + C_in > kRefDimSizeLimit || H_in > kRefDimSizeLimit || + W_in > kRefDimSizeLimit || C_out > kRefDimSizeLimit; + if (has_large_dimension && reference_operations > kRefOperationLimit) { throw std::invalid_argument( "One or more dimensions exceed the allowed limit for reference implementation."); std::cout @@ -643,7 +1161,7 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { auto& ref_data = output_spec.get_ref_float_data(); ref_data.resize(num_output_elements); - const int in_features = utils::align_up_4(C_in_per_group * K_h * K_w); + const int64_t in_features = utils::align_up_4(C_in_per_group * K_h * K_w); // Perform activation, weight, and output quantized conv2d operation for (int64_t n = 0; n < N; ++n) { @@ -722,8 +1240,12 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { float float_result = accum_adjusted * input_scale * weight_scales_data[out_c]; - // Add bias and store result - float_result += bias_data[out_c]; + if (!bias_spec.is_none()) { + float_result += bias_data[out_c]; + } + if (activation_spec.get_string_value() == "relu") { + float_result = std::max(float_result, 0.0f); + } // Quantize the output to int8 float quant_output_f = @@ -749,6 +1271,41 @@ static void reference_impl(TestCase& test_case) { conv2d_q8ta_q8csw_q8to_reference_impl(test_case); } +// The impl selector holds one of "", "general", or "im2col"; activation and +// other string inputs use disjoint values. Overwrite it by value so a +// reordered input list fails loudly instead of mutating the wrong spec. +// Note: for route == "direct" the measured run already forces "general", so +// this reference re-executes the identical implementation and only checks +// determinism; genuine cross-implementation correctness comes from the +// auto/im2col legs. +static void scenex_direct_reference(TestCase& test_case) { + TestCase direct_case = test_case; + bool found_selector = false; + for (auto it = direct_case.inputs().rbegin(); + it != direct_case.inputs().rend(); + ++it) { + if (it->is_string() && + (it->get_string_value().empty() || + it->get_string_value() == "general" || + it->get_string_value() == "im2col")) { + it->string_data = "general"; + found_selector = true; + break; + } + } + if (!found_selector) { + throw std::runtime_error("scenex reference: impl selector input not found"); + } + execute_test_case( + direct_case, + /*warmup_runs=*/1, + /*benchmark_runs=*/1, + /*chained_dispatches=*/1, + /*write_outputs=*/true); + test_case.outputs().at(0).get_ref_float_data() = + direct_case.outputs().at(0).get_float_data(); +} + // Custom FLOP calculator for quantized conv2d operation static int64_t quantized_conv2d_flop_calculator(const TestCase& test_case) { int kernel_idx = 9; // kernel_size is at index 9 for q8ta_q8csw_q8to @@ -779,7 +1336,164 @@ static int64_t quantized_conv2d_flop_calculator(const TestCase& test_case) { return flop; } +static void execute_streaming_dynamic_shrink_test() { + Conv2dConfig config = { + OutInChannels(4, 64), + InputSize2D(30, 99), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1, + 10}; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = + make_test_case_name(config, false, utils::kTexture3D, utils::kBuffer); + TestCase test_case = create_test_case_from_config( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + /*impl_selector=*/"im2col_auto", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT); + for (ValueSpec& input : test_case.inputs()) { + input.ensure_data_generated(/*seed=*/0); + } + + BenchmarkGraph benchmark_graph = setup_compute_graph( + test_case, test_case.operator_name(), /*op_invocations_per_execute=*/1); + ComputeGraph& graph = *benchmark_graph.graph; + graph.prepare(); + graph.prepack(); + + const std::vector upper_input_sizes = test_case.inputs().at(0).sizes; + const std::vector shrunk_input_sizes = {1, 64, 20, 51}; + const std::vector shrunk_output_sizes = {1, 4, 20, 51}; + TestCase shrunk_case = test_case; + shrunk_case.inputs().at(0).sizes = shrunk_input_sizes; + shrunk_case.inputs().at(0).resize_data(1 * 64 * 20 * 51); + shrunk_case.outputs().at(0).sizes = shrunk_output_sizes; + shrunk_case.outputs().at(0).resize_data(1 * 4 * 20 * 51); + reference_impl(shrunk_case); + + constexpr int kRepetitions = 4; + for (int repetition = 0; repetition < kRepetitions; ++repetition) { + graph.resize_input(0, upper_input_sizes); + graph.propagate_resize(); + graph.maybe_cast_and_copy_into_staging( + graph.inputs().at(0).staging, + test_case.inputs().at(0).get_data_ptr(), + test_case.inputs().at(0).numel(), + vkapi::kFloat); + graph.execute(); + + graph.resize_input(0, shrunk_input_sizes); + graph.propagate_resize(); + if (graph.sizes_of(graph.outputs().at(0).value) != shrunk_output_sizes) { + throw std::runtime_error("streaming im2col output did not shrink"); + } + graph.maybe_cast_and_copy_into_staging( + graph.inputs().at(0).staging, + shrunk_case.inputs().at(0).get_data_ptr(), + shrunk_case.inputs().at(0).numel(), + vkapi::kFloat); + graph.execute(); + graph.maybe_cast_and_copy_from_staging( + graph.outputs().at(0).staging, + shrunk_case.outputs().at(0).get_mutable_data_ptr(), + shrunk_case.outputs().at(0).numel(), + vkapi::kFloat); + if (!shrunk_case.outputs().at(0).validate_against_reference( + shrunk_case.get_abs_tolerance(), shrunk_case.get_rel_tolerance())) { + throw std::runtime_error("streaming im2col shrink output was stale"); + } + } + + const Q8taConv2dStreamPlan upper_plan = make_q8ta_conv2d_stream_plan( + /*batch=*/10, + /*flattened_kernel_size=*/576, + /*out_height=*/30, + /*out_width=*/99, + kQ8taConv2dIm2ColScratchBudgetBytes); + if (!upper_plan.feasible || upper_plan.num_tiles != 2 || + upper_plan.rows_per_tile <= 20) { + throw std::runtime_error("dynamic shrink test did not create dead tiles"); + } + const int64_t resized_scratch_bytes = + 576 * upper_plan.rows_per_tile * utils::align_up_4(51); + if (resized_scratch_bytes > kQ8taConv2dIm2ColScratchBudgetBytes) { + throw std::runtime_error("streaming im2col scratch exceeded its cap"); + } + std::cout << "Streaming im2col dynamic shrink PASSED" << std::endl; +} + +// Single usage string for the scenex route-test modes; argument errors +// return 2 like the other CLI errors in main. +static int print_scenex_usage(const char* mode, const char* cases) { + std::cerr << "Usage: " << mode << " <" << cases << ">" + << std::endl; + return 2; +} + int main(int argc, char* argv[]) { + const vkapi::Adapter& adapter = *vkcompute::api::context()->adapter_ptr(); + const bool prefers_unsigned_dot = + adapter.accelerates_unsigned_packed4x8_dot() && + !adapter.accelerates_signed_packed4x8_dot(); + VK_CHECK_COND( + can_use_unsigned_pw_dot(adapter, kMaxUnsignedDotAccumulatorBytes) == + prefers_unsigned_dot); + VK_CHECK_COND( + !can_use_unsigned_pw_dot(adapter, kMaxUnsignedDotAccumulatorBytes + 1)); + + std::string im2col_impl_selector; + bool narrow_workgroups_only = false; + bool streaming_im2col_only = false; + bool streaming_dynamic_shrink_only = false; + const bool scenex_regular = + argc == 4 && std::string(argv[1]) == "--scenex-regular"; + const bool scenex_grouped = + argc == 4 && std::string(argv[1]) == "--scenex-grouped"; + if (argc >= 2 && std::string(argv[1]) == "--scenex-regular" && + !scenex_regular) { + return print_scenex_usage("--scenex-regular", "0|1"); + } + if (argc >= 2 && std::string(argv[1]) == "--scenex-grouped" && + !scenex_grouped) { + return print_scenex_usage("--scenex-grouped", "0|1|2"); + } + if (!scenex_regular && !scenex_grouped) { + for (int i = 1; i < argc; ++i) { + const std::string arg(argv[i]); + if (arg == "--im2col-path=signed") { + im2col_impl_selector = "im2col"; + } else if (arg == "--im2col-path=unsigned") { + im2col_impl_selector = "im2col_unsigned"; + } else if (arg == "--im2col-path=auto") { + im2col_impl_selector = "im2col_auto"; + } else if (arg == "--narrow-workgroups-only") { + narrow_workgroups_only = true; + } else if (arg == "--streaming-im2col-only") { + streaming_im2col_only = true; + } else if (arg == "--streaming-dynamic-shrink-only") { + streaming_dynamic_shrink_only = true; + } else { + std::cerr << "Unknown argument: " << arg << std::endl; + return 2; + } + } + } + const int selected_modes = static_cast(!im2col_impl_selector.empty()) + + static_cast(narrow_workgroups_only) + + static_cast(streaming_im2col_only) + + static_cast(streaming_dynamic_shrink_only) + + static_cast(scenex_regular) + static_cast(scenex_grouped); + if (selected_modes > 1) { + std::cerr << "Test mode selectors are mutually exclusive" << std::endl; + return 2; + } set_debugging(false); set_print_output(false); #ifdef DEBUG_MODE @@ -796,19 +1510,84 @@ int main(int argc, char* argv[]) { print_separator(); ReferenceComputeFunc ref_fn = reference_impl; + int warmup_runs = 1; + int benchmark_runs = 1; - // Execute test cases using the new framework with custom FLOP calculator - auto results = execute_test_cases( + if (streaming_dynamic_shrink_only) { + execute_streaming_dynamic_shrink_test(); + return 0; + } #ifdef DEBUG_MODE - generate_quantized_conv2d_easy_cases, + std::function()> test_case_generator = + generate_quantized_conv2d_easy_cases; #else - generate_quantized_conv2d_test_cases, + std::function()> test_case_generator = + [im2col_impl_selector]() { + return im2col_impl_selector.empty() + ? generate_quantized_conv2d_test_cases() + : generate_im2col_unsigned_test_cases(im2col_impl_selector); + }; +#endif + if (narrow_workgroups_only) { + test_case_generator = generate_narrow_workgroup_test_cases; + } else if (streaming_im2col_only) { + test_case_generator = generate_streaming_im2col_test_cases; +#ifndef DEBUG_MODE + } else if (selected_modes == 0) { + // The default run also covers the unified tile path: multi-tile + // dispatches, grouped/streaming shapes, and the fallback kernel. + test_case_generator = [base_generator = test_case_generator]() { + auto cases = base_generator(); + const auto streaming_cases = generate_streaming_im2col_test_cases(); + cases.insert(cases.end(), streaming_cases.begin(), streaming_cases.end()); + return cases; + }; #endif + } else if (scenex_regular) { + const std::string route = argv[2]; + const std::string case_arg = argv[3]; + if ((route != "auto" && route != "direct" && route != "im2col") || + (case_arg != "0" && case_arg != "1")) { + return print_scenex_usage("--scenex-regular", "0|1"); + } + const int case_index = case_arg == "0" ? 0 : 1; + test_case_generator = [route, case_index]() { + return std::vector{ + generate_scenex_regular_test_case(route, case_index)}; + }; + ref_fn = scenex_direct_reference; + warmup_runs = 3; + benchmark_runs = 10; + } else if (scenex_grouped) { + const std::string route = argv[2]; + const std::string case_arg = argv[3]; + if ((route != "auto" && route != "direct" && route != "im2col") || + (case_arg != "0" && case_arg != "1" && case_arg != "2")) { + return print_scenex_usage("--scenex-grouped", "0|1|2"); + } + const int case_index = case_arg == "0" ? 0 : case_arg == "1" ? 1 : 2; + test_case_generator = [route, case_index]() { + return std::vector{ + generate_scenex_grouped_test_case(route, case_index)}; + }; + ref_fn = scenex_direct_reference; + warmup_runs = 3; + benchmark_runs = 10; + } + + auto results = execute_test_cases( + test_case_generator, quantized_conv2d_flop_calculator, "QuantizedConv2dQ8ToQ8To", - /*warmup_runs = */ 1, - /*benchmark_runs = */ 1, + warmup_runs, + benchmark_runs, ref_fn); +#ifndef DEBUG_MODE + if (selected_modes == 0) { + execute_streaming_dynamic_shrink_test(); + } +#endif + return 0; } diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp index 2dbb4909adb..c573b8b5d39 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp @@ -22,6 +22,7 @@ using namespace executorch::vulkan::prototyping; using namespace vkcompute; static constexpr int64_t kRefDimSizeLimit = 100; +static constexpr int64_t kRefOperationLimit = 2 * 1024 * 1024; // Utility function to create a test case from a Conv2dConfig for depthwise // convolution @@ -271,6 +272,43 @@ std::vector generate_quantized_conv2d_dw_easy_cases() { return test_cases; } +std::vector generate_quantized_conv2d_dw_narrow_workgroup_cases() { + std::vector test_cases; + std::vector configs = { + {OutInChannels(128, 128), + InputSize2D(7, 7), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 128}, + {OutInChannels(64, 64), + InputSize2D(9, 9), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 64}, + {OutInChannels(64, 64), + InputSize2D(13, 13), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 64}, + }; + + for (auto& config : configs) { + const bool is_performance = config.channels.out > kRefDimSizeLimit; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = make_test_case_name( + config, is_performance, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4C)); + } + return test_cases; +} + // Generate test cases for quantized depthwise conv2d operation std::vector generate_quantized_conv2d_dw_test_cases() { std::vector test_cases; @@ -439,6 +477,13 @@ std::vector generate_quantized_conv2d_dw_test_cases() { } } + auto narrow_workgroup_cases = + generate_quantized_conv2d_dw_narrow_workgroup_cases(); + test_cases.insert( + test_cases.end(), + narrow_workgroup_cases.begin(), + narrow_workgroup_cases.end()); + return test_cases; } @@ -498,10 +543,14 @@ void conv2d_q8ta_q8csw_q8to_dw_reference_impl(TestCase& test_case) { int64_t dilation_w = dilation_data[1]; int64_t groups = groups_spec.get_int_value(); - // Skip for large tensors since computation time will be extremely slow - if (N > kRefDimSizeLimit || C_in > kRefDimSizeLimit || - H_in > kRefDimSizeLimit || W_in > kRefDimSizeLimit || - C_out > kRefDimSizeLimit) { + // Skip large tensors only when the reference would be expensive: each + // output element costs K_h * K_w MACs (one input channel per output + // channel), so large-dim cases with few total operations still validate. + const int64_t reference_operations = N * C_out * H_out * W_out * K_h * K_w; + const bool has_large_dimension = N > kRefDimSizeLimit || + C_in > kRefDimSizeLimit || H_in > kRefDimSizeLimit || + W_in > kRefDimSizeLimit || C_out > kRefDimSizeLimit; + if (has_large_dimension && reference_operations > kRefOperationLimit) { throw std::invalid_argument( "One or more dimensions exceed the allowed limit for reference implementation."); } @@ -680,13 +729,27 @@ int main(int argc, char* argv[]) { ReferenceComputeFunc ref_fn = reference_impl; - // Execute test cases using the new framework with custom FLOP calculator - auto results = execute_test_cases( + bool narrow_workgroups_only = false; + for (int i = 1; i < argc; ++i) { + const std::string arg(argv[i]); + if (arg == "--narrow-workgroups-only") { + narrow_workgroups_only = true; + } else { + std::cerr << "Unknown argument: " << arg << std::endl; + return 2; + } + } #ifdef DEBUG_MODE - generate_quantized_conv2d_dw_easy_cases, + auto test_case_generator = generate_quantized_conv2d_dw_easy_cases; #else - generate_quantized_conv2d_dw_test_cases, + auto test_case_generator = generate_quantized_conv2d_dw_test_cases; #endif + if (narrow_workgroups_only) { + test_case_generator = generate_quantized_conv2d_dw_narrow_workgroup_cases; + } + + auto results = execute_test_cases( + test_case_generator, quantized_conv2d_dw_flop_calculator, "QuantizedDepthwiseInt8Conv2d", /*warmup_runs = */ 1, diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp index ee7d8c9e5bf..1003c0540ce 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp @@ -4,10 +4,12 @@ // This source code is licensed under the BSD-style license found in the // LICENSE file in the root directory of this source tree. +#include #include #include #include +#include #include #include @@ -23,13 +25,19 @@ using namespace vkcompute; static constexpr int64_t kRefDimSizeLimit = 100; +struct PointwiseTestOptions { + int32_t input_zero_point = 2; + bool has_bias = true; +}; + // Utility function to create a test case from a Conv2dConfig static TestCase create_test_case_from_config( const Conv2dConfig& config, vkapi::ScalarType input_dtype, utils::StorageType fp_storage_type, utils::GPUMemoryLayout int8_memory_layout, - const std::string& impl_selector = "") { + const std::string& impl_selector = "", + const PointwiseTestOptions& options = {}) { TestCase test_case; // Calculate output dimensions @@ -92,7 +100,7 @@ static TestCase create_test_case_from_config( float input_scale_val = 0.008123; ValueSpec input_scale(input_scale_val); - int32_t input_zero_point_val = 2; + const int32_t input_zero_point_val = options.input_zero_point; ValueSpec input_zero_point(input_zero_point_val); // Quantized weight tensor (int8) - [C_out, C_in_per_group * K_h * K_w] @@ -109,6 +117,15 @@ static TestCase create_test_case_from_config( utils::kWidthPacked, DataGenType::RANDINT8); quantized_weight.set_constant(true); + std::vector& weight_data = quantized_weight.get_int8_data(); + for (int64_t out_channel = 0; out_channel < config.channels.out; + ++out_channel) { + const auto padding_begin = + weight_data.begin() + out_channel * in_features + in_channels_per_group; + const auto padding_end = + weight_data.begin() + (out_channel + 1) * in_features; + std::fill(padding_begin, padding_end, 0); + } if (debugging()) { print_valuespec_data(quantized_weight, "weight_tensor"); @@ -145,6 +162,7 @@ static TestCase create_test_case_from_config( utils::kWidthPacked, DataGenType::ZEROS); bias.set_constant(true); + bias.set_none(!options.has_bias); // Output quantization parameters float output_scale_val = 0.05314; @@ -215,8 +233,12 @@ static TestCase create_test_case_from_config( return test_case; } -// Generate test cases for quantized pointwise conv2d operation -static std::vector generate_quantized_conv2d_pw_test_cases() { +// Generate test cases for quantized pointwise conv2d operation. When +// pw_selector is non-empty ("pw_signed", "pw_unsigned", "pw_auto"), every +// non-legacy case carries that selector so the test graph builder can force or +// verify the unsigned-dot routing instead of using the default automatic path. +static std::vector generate_quantized_conv2d_pw_test_cases( + const std::string& pw_selector = "") { std::vector test_cases; if (!vkcompute::api::context()->adapter_ptr()->supports_int8_dot_product()) { return test_cases; @@ -342,7 +364,11 @@ static std::vector generate_quantized_conv2d_pw_test_cases() { config.test_case_name = make_test_case_name( config, is_performance, fp_storage_type, utils::kBuffer); test_cases.push_back(create_test_case_from_config( - config, vkapi::kFloat, fp_storage_type, int8_memory_layout)); + config, + vkapi::kFloat, + fp_storage_type, + int8_memory_layout, + pw_selector)); // For 4W4C layout, also test the legacy implementation if (int8_memory_layout == utils::kPackedInt8_4W4C) { @@ -390,13 +416,40 @@ static std::vector generate_quantized_conv2d_pw_test_cases() { config.test_case_name = make_test_case_name( config, is_performance, utils::kTexture3D, utils::kBuffer); test_cases.push_back(create_test_case_from_config( - config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4C1W)); + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4C1W, + pw_selector)); if (config.batch == 2) { test_cases.push_back(create_test_case_from_config( - config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4W4C)); + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + pw_selector)); } } + Conv2dConfig edge_config{ + OutInChannels(13, 7), + InputSize2D(7, 5), + KernelSize(1, 1), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1}; + edge_config.op_name = "conv2d_q8ta_q8csw_q8to"; + edge_config.test_case_name = make_test_case_name( + edge_config, false, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + edge_config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + pw_selector, + {.input_zero_point = -128, .has_bias = false})); + return test_cases; } @@ -484,6 +537,7 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { auto& weight_data = weight_spec.get_int8_data(); auto& weight_scales_data = weight_scales_spec.get_float_data(); auto& bias_data = bias_spec.get_float_data(); + const bool has_bias = !bias_spec.is_none(); const float output_scale = output_scale_spec.get_float_value(); const int32_t output_zero_point = output_zeros_spec.get_int_value(); @@ -498,7 +552,7 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { auto& ref_data = output_spec.get_ref_float_data(); ref_data.resize(num_output_elements); - const int in_features = utils::align_up_4(C_in_per_group * K_h * K_w); + const int64_t in_features = utils::align_up_4(C_in_per_group * K_h * K_w); // Perform activation, weight, and output quantized conv2d operation for (int64_t n = 0; n < N; ++n) { @@ -578,7 +632,9 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { accum_adjusted * input_scale * weight_scales_data[out_c]; // Add bias and store result - float_result += bias_data[out_c]; + if (has_bias) { + float_result += bias_data[out_c]; + } // Quantize the output to int8 float quant_output_f = @@ -635,6 +691,30 @@ static int64_t quantized_conv2d_flop_calculator(const TestCase& test_case) { } int main(int argc, char* argv[]) { + const vkapi::Adapter& adapter = *vkcompute::api::context()->adapter_ptr(); + const bool prefers_unsigned_dot = + adapter.accelerates_unsigned_packed4x8_dot() && + !adapter.accelerates_signed_packed4x8_dot(); + VK_CHECK_COND( + can_use_unsigned_pw_dot(adapter, kMaxUnsignedDotAccumulatorBytes) == + prefers_unsigned_dot); + VK_CHECK_COND( + !can_use_unsigned_pw_dot(adapter, kMaxUnsignedDotAccumulatorBytes + 1)); + + std::string pw_impl_selector; + for (int i = 1; i < argc; ++i) { + const std::string arg(argv[i]); + if (arg == "--pw-path=signed") { + pw_impl_selector = "pw_signed"; + } else if (arg == "--pw-path=unsigned") { + pw_impl_selector = "pw_unsigned"; + } else if (arg == "--pw-path=auto") { + pw_impl_selector = "pw_auto"; + } else { + std::cerr << "Unknown argument: " << arg << std::endl; + return 2; + } + } set_debugging(false); set_print_output(false); #ifdef DEBUG_MODE @@ -653,12 +733,11 @@ int main(int argc, char* argv[]) { ReferenceComputeFunc ref_fn = reference_impl; // Execute test cases using the new framework with custom FLOP calculator + const auto test_case_generator = [pw_impl_selector]() { + return generate_quantized_conv2d_pw_test_cases(pw_impl_selector); + }; auto results = execute_test_cases( -#ifdef DEBUG_MODE - generate_quantized_conv2d_pw_test_cases, -#else - generate_quantized_conv2d_pw_test_cases, -#endif + test_case_generator, quantized_conv2d_flop_calculator, "QuantizedConv2dPW", /*warmup_runs = */ 1, diff --git a/backends/vulkan/test/custom_ops/utils.cpp b/backends/vulkan/test/custom_ops/utils.cpp index 60c3ddee30c..b00e6b175da 100644 --- a/backends/vulkan/test/custom_ops/utils.cpp +++ b/backends/vulkan/test/custom_ops/utils.cpp @@ -31,6 +31,56 @@ constexpr float kMinProbeTimeUs = 1.0f; } // namespace +// Benchmark graphs are immutable after input upload, so the operator nodes can +// be recorded repeatedly. Staging uploads/downloads are encoded once: +// repeating them would redo host-visible copies on every repetition and bias +// per-invocation timings. +RepeatedGraphExecutor::RepeatedGraphExecutor( + ComputeGraph& graph, + int repetitions, + OpNodeRange op_nodes) + : graph_(graph) { + VK_CHECK_COND(repetitions > 0); + VK_CHECK_COND(op_nodes.begin <= op_nodes.end); + VK_CHECK_COND(op_nodes.end <= graph_.execute_nodes().size()); + + api::Context* const context = graph_.context(); + context->flush(); + context->set_cmd(/*reusable=*/true); + context->cmd_reset_querypool(); + + for (size_t i = 0; i < op_nodes.begin; ++i) { + graph_.execute_nodes()[i]->encode(&graph_); + } + for (int i = 0; i < repetitions; ++i) { + for (size_t j = op_nodes.begin; j < op_nodes.end; ++j) { + graph_.execute_nodes()[j]->encode(&graph_); + } + } + for (size_t i = op_nodes.end; i < graph_.execute_nodes().size(); ++i) { + graph_.execute_nodes()[i]->encode(&graph_); + } + + command_ = + std::make_unique(std::move(context->extract_cmd())); +} + +void RepeatedGraphExecutor::execute() { + api::Context* const context = graph_.context(); + command_->end(); + + // Intentionally bypasses Context::submit_cmd_to_gpu(): its submit-count + // bookkeeping and threshold-split path only serve submit_compute_job + // recording, which benchmarks don't use. + vkapi::VulkanFence fence = context->fences().get_fence(); + context->adapter_ptr()->submit_cmd( + context->queue(), + command_->get_submit_handle(/*final_use=*/false), + fence.get_submit_handle()); + fence.wait(); + context->fences().return_fence(fence); +} + int get_seed() { static int seed = 42; return seed++; @@ -172,11 +222,30 @@ void set_debugging(bool enable_debugging) { } // ValueSpec implementation -void ValueSpec::generate_tensor_data(int seed) { +void ValueSpec::ensure_unique_data() const { + if (data_.use_count() != 1) { + data_ = std::make_shared(*data_); + } +} + +void ValueSpec::ensure_unique_reference_data() const { + if (reference_data_.use_count() != 1) { + reference_data_ = std::make_shared(*reference_data_); + } +} + +void ValueSpec::generate_tensor_data(int seed) const { if (spec_type != SpecType::Tensor) { return; } + ensure_unique_data(); + auto& float_data = data_->float_data; + auto& int32_data = data_->int32_data; + auto& half_data = data_->half_data; + auto& int8_data = data_->int8_data; + auto& uint8_data = data_->uint8_data; + int64_t num_elements = numel(); switch (dtype) { @@ -498,81 +567,95 @@ std::string ValueSpec::to_string() const { // Additional ValueSpec methods void ValueSpec::resize_data(size_t new_size) { + // Generate first so a deferred tensor keeps its data-gen pattern (resized, + // not pinned to zeros by the data_generated_ flag set below). + ensure_data_generated(); + ensure_unique_data(); switch (dtype) { case vkapi::kFloat: - float_data.resize(new_size); + data_->float_data.resize(new_size); break; case vkapi::kHalf: - half_data.resize(new_size); + data_->half_data.resize(new_size); break; case vkapi::kInt: - int32_data.resize(new_size); + data_->int32_data.resize(new_size); break; case vkapi::kChar: - int8_data.resize(new_size); + data_->int8_data.resize(new_size); break; case vkapi::kByte: - uint8_data.resize(new_size); + data_->uint8_data.resize(new_size); break; default: - float_data.resize(new_size); + data_->float_data.resize(new_size); break; } + data_generated_ = true; } void* ValueSpec::get_mutable_data_ptr() { + ensure_data_generated(); + ensure_unique_data(); switch (dtype) { case vkapi::kFloat: - return float_data.data(); + return data_->float_data.data(); case vkapi::kHalf: - return half_data.data(); + return data_->half_data.data(); case vkapi::kInt: - return int32_data.data(); + return data_->int32_data.data(); case vkapi::kChar: - return int8_data.data(); + return data_->int8_data.data(); case vkapi::kByte: - return uint8_data.data(); + return data_->uint8_data.data(); default: - return float_data.data(); + return data_->float_data.data(); } } float ValueSpec::get_element(size_t index) const { + ensure_data_generated(); if (index >= static_cast(numel())) { return 0.0f; } switch (dtype) { case vkapi::kFloat: - return index < float_data.size() ? float_data[index] : 0.0f; + return index < data_->float_data.size() ? data_->float_data[index] : 0.0f; case vkapi::kHalf: - return index < half_data.size() ? half_to_float(half_data[index]) : 0.0f; + return index < data_->half_data.size() + ? half_to_float(data_->half_data[index]) + : 0.0f; case vkapi::kInt: - return index < int32_data.size() ? static_cast(int32_data[index]) - : 0.0f; + return index < data_->int32_data.size() + ? static_cast(data_->int32_data[index]) + : 0.0f; case vkapi::kChar: - return index < int8_data.size() ? static_cast(int8_data[index]) - : 0.0f; + return index < data_->int8_data.size() + ? static_cast(data_->int8_data[index]) + : 0.0f; case vkapi::kByte: - return index < uint8_data.size() ? static_cast(uint8_data[index]) - : 0.0f; + return index < data_->uint8_data.size() + ? static_cast(data_->uint8_data[index]) + : 0.0f; default: return 0.0f; } } const void* ValueSpec::get_data_ptr() const { + ensure_data_generated(); switch (dtype) { case vkapi::kFloat: - return float_data.data(); + return data_->float_data.data(); case vkapi::kHalf: - return half_data.data(); + return data_->half_data.data(); case vkapi::kInt: - return int32_data.data(); + return data_->int32_data.data(); case vkapi::kChar: - return int8_data.data(); + return data_->int8_data.data(); case vkapi::kByte: - return uint8_data.data(); + return data_->uint8_data.data(); default: throw std::runtime_error("Unsupported data type for get_data_ptr"); } @@ -801,7 +884,7 @@ bool ValueSpec::validate_against_reference( } // Ensure data is generated for this ValueSpec -void ValueSpec::ensure_data_generated(int seed) { +void ValueSpec::ensure_data_generated(int seed) const { if (data_generated_) { return; } @@ -809,18 +892,22 @@ void ValueSpec::ensure_data_generated(int seed) { data_generated_ = true; } -// Copy input data from another ValueSpec -void ValueSpec::copy_data_from(const ValueSpec& other) { +void ValueSpec::share_data_from(const ValueSpec& other) { + if (!is_tensor() || !other.is_tensor()) { + return; + } + // Materialize the source first: sharing an ungenerated payload would let a + // later access materialize each spec independently under different seeds. + other.ensure_data_generated(); + data_ = other.data_; + data_generated_ = true; +} + +void ValueSpec::share_reference_from(const ValueSpec& other) { if (!is_tensor() || !other.is_tensor()) { return; } - // Copy raw data based on dtype - float_data = other.float_data; - int32_data = other.int32_data; - half_data = other.half_data; - int8_data = other.int8_data; - uint8_data = other.uint8_data; - data_generated_ = other.data_generated_; + reference_data_ = other.reference_data_; } // ReferenceKey implementation @@ -1396,17 +1483,23 @@ int64_t default_flop_calculator(const TestCase& test_case) { return total_elements; } -ComputeGraph setup_compute_graph( +BenchmarkGraph setup_compute_graph( TestCase& test_case, std::string op_name, int op_invocations_per_execute) { GraphConfig config; config.enable_querypool = true; + // Pool sizing takes max(execute, prepack) * factor, so scaling the factor + // also over-reserves the prepack side (encoded once). Accepted: precise + // execute-only sizing would need runtime changes. + config.descriptor_pool_safety_factor *= + std::max(1, op_invocations_per_execute); // Default-on (opt-out via TestCase::set_force_resize(false)): force every - // DynamicDispatchNode to run its resize function on each execute(), - // exercising the op's resize formula even when input shapes are unchanged. + // DynamicDispatchNode to run its resize function when execute_test_case + // runs propagate_resize() after prepack, exercising the op's resize formula + // even when input shapes are unchanged. config.force_resize = test_case.get_force_resize(); - ComputeGraph graph(config); + auto graph = std::make_unique(config); std::vector input_values; @@ -1415,17 +1508,17 @@ ComputeGraph setup_compute_graph( const ValueSpec& input_spec = test_case.inputs()[i]; if (input_spec.is_none()) { - input_values.push_back(graph.add_none()); + input_values.push_back(graph->add_none()); } else if (input_spec.is_float()) { ValueRef input_value = - graph.add_scalar(static_cast(input_spec.get_float_value())); + graph->add_scalar(static_cast(input_spec.get_float_value())); input_values.push_back(input_value); } else if (input_spec.is_int()) { ValueRef input_value = - graph.add_scalar(static_cast(input_spec.get_int_value())); + graph->add_scalar(static_cast(input_spec.get_int_value())); input_values.push_back(input_value); } else if (input_spec.is_bool()) { - ValueRef input_value = graph.add_scalar(input_spec.get_bool_value()); + ValueRef input_value = graph->add_scalar(input_spec.get_bool_value()); input_values.push_back(input_value); } else if (input_spec.is_int_list()) { // Convert int32_t list to int64_t list for ComputeGraph @@ -1435,20 +1528,20 @@ ComputeGraph setup_compute_graph( for (int32_t val : int32_list) { int64_list.push_back(static_cast(val)); } - ValueRef input_value = graph.add_scalar_list(std::move(int64_list)); + ValueRef input_value = graph->add_scalar_list(std::move(int64_list)); input_values.push_back(input_value); } else if (input_spec.is_string()) { std::string str_copy = input_spec.get_string_value(); - ValueRef input_value = graph.add_string(std::move(str_copy)); + ValueRef input_value = graph->add_string(std::move(str_copy)); input_values.push_back(input_value); } else if (input_spec.is_constant()) { - ValueRef input_value = graph.add_tensorref( + ValueRef input_value = graph->add_tensorref( input_spec.get_tensor_sizes(), input_spec.dtype, input_spec.get_data_ptr()); input_values.push_back(input_value); } else { - IOValueRef input_io = graph.add_input_tensor( + IOValueRef input_io = graph->add_input_tensor( input_spec.get_tensor_sizes(), input_spec.dtype, input_spec.storage_type, @@ -1468,7 +1561,7 @@ ComputeGraph setup_compute_graph( } // Create output tensor - ValueRef output_value = graph.add_tensor( + ValueRef output_value = graph->add_tensor( output_spec.get_tensor_sizes(), output_spec.dtype, output_spec.storage_type, @@ -1484,17 +1577,16 @@ ComputeGraph setup_compute_graph( std::vector op_args = input_values; op_args.insert(op_args.end(), output_values.begin(), output_values.end()); - // Invoke the op op_invocations_per_execute times to stack dispatches per - // graph.execute(). The output set_output_value() calls below still happen - // exactly once. - for (int i = 0; i < op_invocations_per_execute; ++i) { - opFn(graph, op_args); - } + // Nodes added before the op are staging uploads; nodes added after are + // staging downloads. Only the op's own nodes are repeated by benchmarks. + const size_t op_begin = graph->execute_nodes().size(); + opFn(*graph, op_args); + const size_t op_end = graph->execute_nodes().size(); for (size_t i = 0; i < output_values.size(); ++i) { - graph.set_output_value(output_values[i]); + graph->set_output_value(output_values[i]); } - return graph; + return {std::move(graph), {op_begin, op_end}}; } // Test execution utilities @@ -1512,16 +1604,20 @@ BenchmarkResult execute_test_case( api::context()->initialize_querypool(); } - // Build the measurement graph with the requested chained_dispatches factor. - // The caller (typically execute_test_cases) decides what it should be — - // this function is a pure "run at the given chained_dispatches" primitive. - ComputeGraph graph = setup_compute_graph( + // Build the operator once. Benchmark repetition is encoded separately so + // persistent graph allocations are not duplicated. + BenchmarkGraph benchmark = setup_compute_graph( test_case, test_case.operator_name(), chained_dispatches); + ComputeGraph& graph = *benchmark.graph; // Prepare the graph graph.prepare(); graph.prepack(); + // Run resize functions once so force_resize exercises resize formulas even + // though the record/replay path below never calls propagate_resize(). + graph.propagate_resize(); + // Copy input data into the graph's staging buffers size_t graph_input_idx = 0; for (size_t i = 0; i < test_case.num_inputs(); ++i) { @@ -1568,9 +1664,12 @@ BenchmarkResult execute_test_case( ++graph_input_idx; } + RepeatedGraphExecutor graph_executor( + graph, chained_dispatches, benchmark.op_nodes); + // Warmup runs for (int run = 0; run < warmup_runs; ++run) { - graph.execute(); + graph_executor.execute(); } // Benchmark runs - collect individual iteration timings @@ -1582,7 +1681,7 @@ BenchmarkResult execute_test_case( for (int run = 0; run < benchmark_runs; ++run) { // Measure CPU time for each execute() call auto cpu_start = std::chrono::high_resolution_clock::now(); - graph.execute(); + graph_executor.execute(); auto cpu_end = std::chrono::high_resolution_clock::now(); auto cpu_duration = std::chrono::duration_cast( @@ -1743,17 +1842,11 @@ TestResult execute_test_cases( // Compute reference once for prototype bool ref_computed = false; - std::vector> ref_data; if (reference_compute_func) { try { reference_compute_func(prototype); ref_computed = true; - - // Cache the reference output for this group - for (const auto& output : prototype.outputs()) { - ref_data.push_back(output.get_ref_float_data()); - } - } catch (const std::invalid_argument& _) { + } catch (const std::invalid_argument&) { // Reference computation skipped for this group } } @@ -1771,15 +1864,21 @@ TestResult execute_test_cases( const auto& src = prototype.inputs()[j]; if (dest.is_tensor() && src.is_tensor() && dest.sizes == src.sizes && dest.dtype == src.dtype) { - dest.copy_data_from(src); + dest.share_data_from(src); } } // Copy reference output data if available if (ref_computed) { - for (size_t j = 0; j < tc.outputs().size() && j < ref_data.size(); + for (size_t j = 0; + j < tc.outputs().size() && j < prototype.outputs().size(); ++j) { - tc.outputs()[j].get_ref_float_data() = ref_data[j]; + const auto& src = prototype.outputs()[j]; + auto& dest = tc.outputs()[j]; + if (dest.is_tensor() && src.is_tensor() && dest.sizes == src.sizes && + dest.dtype == src.dtype) { + dest.share_reference_from(src); + } } } } @@ -1924,6 +2023,8 @@ TestResult execute_test_cases( // Add result to collection results.add_result(std::move(result)); + + test_case.clear(); } } diff --git a/backends/vulkan/test/custom_ops/utils.h b/backends/vulkan/test/custom_ops/utils.h index 2174ceb5618..c7d537e6a2b 100644 --- a/backends/vulkan/test/custom_ops/utils.h +++ b/backends/vulkan/test/custom_ops/utils.h @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -220,21 +221,8 @@ struct ValueSpec { bool is_constant_tensor; bool is_none_flag; bool is_int4_tensor; - bool data_generated_ = false; - - std::vector float_data; - std::vector int32_data; - std::vector half_data; // Using uint16_t as substitute for half - std::vector int8_data; // For kChar (signed 8-bit) - std::vector uint8_data; // For kByte (unsigned 8-bit) std::string string_data; - std::vector ref_float_data; - std::vector ref_int32_data; - std::vector ref_half_data; - std::vector ref_int8_data; - std::vector ref_uint8_data; - ValueSpec( const std::vector& sizes, vkapi::ScalarType dtype, @@ -250,7 +238,8 @@ struct ValueSpec { is_none_flag(false), is_int4_tensor(false), data_generated_(false) { - // Data generation is deferred until ensure_data_generated() is called + // Data generation is deferred until first access (any data getter or + // ensure_data_generated() triggers it). } // Constructor for tensor with custom data generation type @@ -270,7 +259,8 @@ struct ValueSpec { is_none_flag(false), is_int4_tensor(false), data_generated_(false) { - // Data generation is deferred until ensure_data_generated() is called + // Data generation is deferred until first access (any data getter or + // ensure_data_generated() triggers it). } // Constructor for single int @@ -285,7 +275,7 @@ struct ValueSpec { is_none_flag(false), is_int4_tensor(false), data_generated_(true) { - int32_data.push_back(value); + data_->int32_data.push_back(value); } // Constructor for single float @@ -300,7 +290,7 @@ struct ValueSpec { is_none_flag(false), is_int4_tensor(false), data_generated_(true) { - float_data.push_back(value); + data_->float_data.push_back(value); } // Constructor for single bool @@ -315,7 +305,7 @@ struct ValueSpec { is_none_flag(false), is_int4_tensor(false), data_generated_(true) { - int32_data.push_back(value ? 1 : 0); + data_->int32_data.push_back(value ? 1 : 0); } // Constructor for int list @@ -329,8 +319,9 @@ struct ValueSpec { is_constant_tensor(false), is_none_flag(false), is_int4_tensor(false), - data_generated_(true), - int32_data(values) {} + data_generated_(true) { + data_->int32_data = values; + } // Factory method for string (avoids ambiguity with vector constructor) static ValueSpec make_string(const std::string& value) { @@ -385,98 +376,136 @@ struct ValueSpec { } int32_t get_int_value() const { - return int32_data.empty() ? 0 : int32_data[0]; + ensure_data_generated(); + return data_->int32_data.empty() ? 0 : data_->int32_data[0]; } float get_float_value() const { - return float_data.empty() ? 0.0f : float_data[0]; + ensure_data_generated(); + return data_->float_data.empty() ? 0.0f : data_->float_data[0]; } bool get_bool_value() const { - return int32_data.empty() ? false : (int32_data[0] != 0); + ensure_data_generated(); + return data_->int32_data.empty() ? false : (data_->int32_data[0] != 0); } const std::string& get_string_value() const { return string_data; } const std::vector& get_int_list() const { - return int32_data; + ensure_data_generated(); + return data_->int32_data; } const std::vector& get_tensor_sizes() const { return sizes; } + // References and pointers into tensor data must not be held across any other + // access to the same spec: a mutable access may detach the shared payload, + // leaving a previously returned reference bound to the old payload. Consume + // immediately. const std::vector& get_float_data() const { - return float_data; + ensure_data_generated(); + return data_->float_data; } const std::vector& get_int32_data() const { - return int32_data; + ensure_data_generated(); + return data_->int32_data; } const std::vector& get_half_data() const { - return half_data; + ensure_data_generated(); + return data_->half_data; } const std::vector& get_int8_data() const { - return int8_data; + ensure_data_generated(); + return data_->int8_data; } const std::vector& get_uint8_data() const { - return uint8_data; + ensure_data_generated(); + return data_->uint8_data; } std::vector& get_float_data() { - return float_data; + ensure_data_generated(); + ensure_unique_data(); + return data_->float_data; } std::vector& get_int32_data() { - return int32_data; + ensure_data_generated(); + ensure_unique_data(); + return data_->int32_data; } std::vector& get_half_data() { - return half_data; + ensure_data_generated(); + ensure_unique_data(); + return data_->half_data; } std::vector& get_int8_data() { - return int8_data; + ensure_data_generated(); + ensure_unique_data(); + return data_->int8_data; } std::vector& get_uint8_data() { - return uint8_data; + ensure_data_generated(); + ensure_unique_data(); + return data_->uint8_data; } const std::vector& get_ref_float_data() const { - return ref_float_data; + return reference_data_->float_data; } const std::vector& get_ref_int32_data() const { - return ref_int32_data; + return reference_data_->int32_data; } const std::vector& get_ref_half_data() const { - return ref_half_data; + return reference_data_->half_data; } const std::vector& get_ref_int8_data() const { - return ref_int8_data; + return reference_data_->int8_data; } const std::vector& get_ref_uint8_data() const { - return ref_uint8_data; + return reference_data_->uint8_data; } std::vector& get_ref_float_data() { - return ref_float_data; + ensure_unique_reference_data(); + return reference_data_->float_data; } std::vector& get_ref_int32_data() { - return ref_int32_data; + ensure_unique_reference_data(); + return reference_data_->int32_data; } std::vector& get_ref_half_data() { - return ref_half_data; + ensure_unique_reference_data(); + return reference_data_->half_data; } std::vector& get_ref_int8_data() { - return ref_int8_data; + ensure_unique_reference_data(); + return reference_data_->int8_data; } std::vector& get_ref_uint8_data() { - return ref_uint8_data; + ensure_unique_reference_data(); + return reference_data_->uint8_data; } void resize_data(size_t new_size); void* get_mutable_data_ptr(); float get_element(size_t index) const; - // Data generation methods for deferred generation and caching + // Data generation methods for deferred generation and caching. + // + // ValueSpec is not thread-safe: lazy materialization and copy-on-write + // detach mutate shared state from const methods. Test cases are built and + // executed on a single thread. + // + // Implicit materialization (any data getter, resize_data) consumes the + // global seed counter. Callers needing deterministic data must call + // ensure_data_generated(explicit_seed) before any other access; a later + // seeded call is a no-op once data is generated. bool is_data_generated() const { return data_generated_; } - void ensure_data_generated(int seed = -1); - void copy_data_from(const ValueSpec& other); + void ensure_data_generated(int seed = -1) const; + void share_data_from(const ValueSpec& other); + void share_reference_from(const ValueSpec& other); // Set/get constant flag bool is_constant() const { @@ -484,10 +513,6 @@ struct ValueSpec { } void set_constant(bool is_constant) { is_constant_tensor = is_constant; - // Constant tensors need data immediately for test case setup - if (is_constant && is_tensor()) { - ensure_data_generated(); - } } // Set/get none flag @@ -517,7 +542,22 @@ struct ValueSpec { float rel_tolerance = 1e-3f) const; private: - void generate_tensor_data(int seed = -1); + struct TensorData { + std::vector float_data; + std::vector int32_data; + std::vector half_data; + std::vector int8_data; + std::vector uint8_data; + }; + + void ensure_unique_data() const; + void ensure_unique_reference_data() const; + void generate_tensor_data(int seed = -1) const; + + mutable bool data_generated_ = false; + mutable std::shared_ptr data_ = std::make_shared(); + mutable std::shared_ptr reference_data_ = + std::make_shared(); }; // @@ -581,9 +621,9 @@ class TestCase { return shader_filter_; } - // Manual override for the number of times the op is dispatched per - // graph.execute() (a.k.a. chained_dispatches). If > 0, the framework uses - // this directly and skips the probe phase. 0 (the default) means adaptive + // Manual override for the number of chained dispatches per measurement + // iteration (a.k.a. chained_dispatches). If > 0, the framework uses this + // directly and skips the probe phase. 0 (the default) means adaptive // (probe-then-scale). void set_op_invocations_per_execute(int n) { op_invocations_per_execute_ = n; @@ -605,13 +645,15 @@ class TestCase { // When true, the ComputeGraph built for this test case sets // GraphConfig::force_resize, so every DynamicDispatchNode runs its resize - // function on each execute() even when no input shape changed. Because the - // output is already allocated at the swept shape, the resize must recompute - // the same shape from the current input — a wrong resize formula resizes the - // output to a mismatched shape and surfaces as a test failure. Default true - // (opt-out): every custom_ops test exercises its resize formulas across the - // swept shapes. Call set_force_resize(false) for the rare op whose resize fn - // is intentionally not shape-preserving under a fixed output allocation. + // function once during measurement setup (execute_test_case runs + // propagate_resize() after prepack) even when no input shape changed. + // Because the output is already allocated at the swept shape, the resize + // must recompute the same shape from the current input — a wrong resize + // formula resizes the output to a mismatched shape and surfaces as a test + // failure. Default true (opt-out): every custom_ops test exercises its + // resize formulas across the swept shapes. Call set_force_resize(false) for + // the rare op whose resize fn is intentionally not shape-preserving under a + // fixed output allocation. void set_force_resize(bool force_resize) { force_resize_ = force_resize; } @@ -900,9 +942,55 @@ int64_t default_flop_calculator(const TestCase& test_case); using ReferenceComputeFunc = std::function; -// Runs a measurement at the given chained_dispatches factor (how many times -// the op is stacked inside one graph.execute()). This is a primitive; the -// probe-then-scale orchestration lives in execute_test_cases(). +// Half-open index range of the operator's own dispatch nodes within a +// benchmark graph's execute_nodes(). Staging upload nodes precede it, staging +// download nodes follow it. +struct OpNodeRange { + size_t begin = 0; + size_t end = 0; +}; + +// A benchmark graph plus the location of its repeatable operator nodes. The +// graph is heap-held: ComputeGraph owns its Context and must never be moved +// (a moved-from graph's destructor dereferences a null context). +struct BenchmarkGraph { + std::unique_ptr graph; + OpNodeRange op_nodes; +}; + +// Benchmark-only executor that records a graph's execute nodes into a single +// reusable command buffer, then replays it on every execute(). Staging uploads +// are encoded once, operator nodes N times, staging downloads once, so each +// iteration performs the same work as the old stacked-nodes layout (1 upload + +// N ops + 1 download) and the per-invocation divisor is unchanged. Production +// ComputeGraph::execute() behavior is unchanged. +// +// Notes for interpreting benchmark numbers: +// - Submit granularity differs from stacking N distinct nodes: all encodings +// live in one command buffer with one submit per iteration (the old path +// could split across command buffers at the node-count threshold), so +// per-dispatch times may shift systematically against older data. +// - Resize functions run once via propagate_resize() in execute_test_case +// before recording; replay itself never re-triggers resize. +// - Repeated encodings share one node/dispatch id, so per-repetition +// querypool attribution is unavailable (aggregation keys on kernel name). +class RepeatedGraphExecutor final { + public: + RepeatedGraphExecutor( + ComputeGraph& graph, + int repetitions, + OpNodeRange op_nodes); + void execute(); + + private: + ComputeGraph& graph_; + std::unique_ptr command_; +}; + +// Runs a measurement at the given chained_dispatches factor. The operator is +// built once, then its execute nodes are encoded that many times into a +// benchmark-only reusable command buffer. The probe-then-scale orchestration +// lives in execute_test_cases(). // // write_outputs controls whether the graph's staging output buffers are copied // back into test_case.outputs() at the end of the run. The probe path needs @@ -980,11 +1068,11 @@ void compute_weight_sums_4bit_grouped( uint16_t float_to_half(float value); float half_to_float(uint16_t half_val); -// Setup compute graph based on TestCase and operation name. The op function -// is invoked op_invocations_per_execute times so that one graph.execute() -// dispatches the op that many times (Google Benchmark-style stacking). The -// output set_output_value() calls still happen once at the end. -ComputeGraph setup_compute_graph( +// Setup compute graph based on TestCase and operation name. The op function is +// invoked once. op_invocations_per_execute is used only to reserve enough +// descriptor capacity for benchmark-only repeated command encoding. Returns +// the graph plus the range of the operator's own nodes for repeated encoding. +BenchmarkGraph setup_compute_graph( TestCase& test_case, std::string op_name, int op_invocations_per_execute = 1); diff --git a/backends/vulkan/test/custom_ops/utils_test.cpp b/backends/vulkan/test/custom_ops/utils_test.cpp new file mode 100644 index 00000000000..6d2de13d71e --- /dev/null +++ b/backends/vulkan/test/custom_ops/utils_test.cpp @@ -0,0 +1,227 @@ +// 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. + +#include + +#include "utils.h" + +namespace executorch::vulkan::prototyping { +namespace { + +int operator_build_count = 0; +int encode_count = 0; + +class CountingEncodeNode final : public ExecuteNode { + public: + explicit CountingEncodeNode(int& count) : count_(count) {} + void encode(ComputeGraph* graph) override { + (void)graph; + ++count_; + } + + private: + int& count_; +}; + +void counting_operator(ComputeGraph& graph, const std::vector& args) { + (void)args; + ++operator_build_count; + graph.execute_nodes().emplace_back( + std::make_unique(encode_count)); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(test_etvk.counting_operator.default, counting_operator); +} + +TEST(ValueSpecTest, SetConstant_DoesNotMaterializeTensorData) { + ValueSpec value( + {16}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ONES); + + value.set_constant(true); + + EXPECT_FALSE(value.is_data_generated()); +} + +TEST(ValueSpecTest, ShareDataFrom_SharesImmutableTensorData) { + ValueSpec source( + {16}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ONES); + source.ensure_data_generated(); + + ValueSpec copy( + {16}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kTexture3D, + vkcompute::utils::kChannelsPacked, + DataGenType::ONES); + copy.share_data_from(source); + + const ValueSpec& const_source = source; + const ValueSpec& const_copy = copy; + EXPECT_EQ( + const_source.get_float_data().data(), const_copy.get_float_data().data()); +} + +TEST(ValueSpecTest, MutableDataAccess_DetachesSharedTensorData) { + ValueSpec source( + {4}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ONES); + source.ensure_data_generated(); + + ValueSpec copy = source; + copy.get_float_data()[0] = 7.0f; + + const ValueSpec& const_source = source; + const ValueSpec& const_copy = copy; + EXPECT_FLOAT_EQ(const_source.get_float_data()[0], 1.0f); + EXPECT_FLOAT_EQ(const_copy.get_float_data()[0], 7.0f); + EXPECT_NE( + const_source.get_float_data().data(), const_copy.get_float_data().data()); +} + +TEST(ValueSpecTest, CopyConstruction_SharesImmutableReferenceData) { + ValueSpec source( + {4}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ZEROS); + source.get_ref_float_data() = {1.0f, 2.0f, 3.0f, 4.0f}; + + ValueSpec copy = source; + + const ValueSpec& const_source = source; + const ValueSpec& const_copy = copy; + EXPECT_EQ( + const_source.get_ref_float_data().data(), + const_copy.get_ref_float_data().data()); +} + +TEST(ValueSpecTest, MutableReferenceAccess_DetachesSharedReferenceData) { + ValueSpec source( + {2}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ZEROS); + source.get_ref_float_data() = {1.0f, 2.0f}; + + ValueSpec copy = source; + copy.get_ref_float_data()[0] = 9.0f; + + const ValueSpec& const_source = source; + const ValueSpec& const_copy = copy; + EXPECT_FLOAT_EQ(const_source.get_ref_float_data()[0], 1.0f); + EXPECT_FLOAT_EQ(const_copy.get_ref_float_data()[0], 9.0f); + EXPECT_NE( + const_source.get_ref_float_data().data(), + const_copy.get_ref_float_data().data()); +} + +TEST(ValueSpecTest, ConstGetter_MaterializesDeferredData) { + ValueSpec value( + {4}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ONES); + const ValueSpec& const_value = value; + EXPECT_FALSE(const_value.is_data_generated()); + EXPECT_FLOAT_EQ(const_value.get_float_data()[0], 1.0f); + EXPECT_FLOAT_EQ(const_value.get_float_value(), 1.0f); + EXPECT_TRUE(const_value.is_data_generated()); +} + +TEST(ValueSpecTest, ResizeData_PreservesGeneratedPattern) { + ValueSpec value( + {4}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ONES); + value.resize_data(8); + const ValueSpec& const_value = value; + const std::vector expected( + {1.0f, 1.0f, 1.0f, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f}); + EXPECT_EQ(const_value.get_float_data(), expected); +} + +TEST(ValueSpecTest, MutableDataPtr_DetachesSharedTensorData) { + ValueSpec source( + {4}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ONES); + source.ensure_data_generated(); + + ValueSpec copy = source; + auto* mutable_ptr = static_cast(copy.get_mutable_data_ptr()); + ASSERT_NE(mutable_ptr, nullptr); + mutable_ptr[0] = 7.0f; + + const ValueSpec& const_source = source; + const ValueSpec& const_copy = copy; + EXPECT_FLOAT_EQ(const_source.get_float_data()[0], 1.0f); + EXPECT_FLOAT_EQ(const_copy.get_float_data()[0], 7.0f); +} + +TEST(ValueSpecTest, ShareReferenceFrom_IgnoresNonTensorSpecs) { + ValueSpec scalar(3); + ValueSpec tensor( + {4}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ZEROS); + tensor.get_ref_float_data() = {1.0f, 2.0f, 3.0f, 4.0f}; + const void* before = tensor.get_ref_float_data().data(); + tensor.share_reference_from(scalar); + EXPECT_EQ(tensor.get_ref_float_data().data(), before); + const std::vector expected({1.0f, 2.0f, 3.0f, 4.0f}); + EXPECT_EQ(tensor.get_ref_float_data(), expected); +} + +TEST(BenchmarkGraphTest, ChainedDispatches_BuildOperatorOnce) { + if (!vkcompute::api::available()) { + return; + } + + TestCase test_case; + operator_build_count = 0; + encode_count = 0; + + BenchmarkGraph benchmark = setup_compute_graph( + test_case, + "test_etvk.counting_operator.default", + /*op_invocations_per_execute=*/8); + ComputeGraph& graph = *benchmark.graph; + + EXPECT_EQ(operator_build_count, 1); + ASSERT_EQ(graph.execute_nodes().size(), 1u); + EXPECT_EQ(benchmark.op_nodes.begin, 0u); + EXPECT_EQ(benchmark.op_nodes.end, 1u); + + // The replay path must encode the operator node once per repetition. + RepeatedGraphExecutor graph_executor( + graph, /*repetitions=*/8, benchmark.op_nodes); + EXPECT_EQ(encode_count, 8); + graph_executor.execute(); +} + +} // namespace +} // namespace executorch::vulkan::prototyping diff --git a/backends/vulkan/test/op_tests/CMakeLists.txt b/backends/vulkan/test/op_tests/CMakeLists.txt index 0f8456accf5..5facea5a14b 100644 --- a/backends/vulkan/test/op_tests/CMakeLists.txt +++ b/backends/vulkan/test/op_tests/CMakeLists.txt @@ -73,6 +73,8 @@ function(vulkan_op_test test_name test_src) add_executable(${test_name} ${test_src}) target_include_directories(${test_name} PRIVATE ${COMMON_INCLUDES}) + # ATen headers require C++20. + set_target_properties(${test_name} PROPERTIES CXX_STANDARD 20) target_link_libraries( ${test_name} PRIVATE GTest::gtest_main @@ -90,6 +92,8 @@ endfunction() if(TARGET vulkan_backend AND LIB_TORCH) add_library(test_utils ${CMAKE_CURRENT_SOURCE_DIR}/test_utils.cpp) target_include_directories(test_utils PRIVATE ${COMMON_INCLUDES}) + # ATen headers require C++20. + set_target_properties(test_utils PROPERTIES CXX_STANDARD 20) target_link_libraries( test_utils PRIVATE vulkan_backend ${LIB_TORCH} ${LIB_TORCH_CPU} ) diff --git a/backends/vulkan/test/scripts/test_model.sh b/backends/vulkan/test/scripts/test_model.sh index 40ec88bae70..4adbdc407df 100755 --- a/backends/vulkan/test/scripts/test_model.sh +++ b/backends/vulkan/test/scripts/test_model.sh @@ -95,7 +95,7 @@ clean_build_directory() { } recompile() { - cmake --build cmake-out -j64 --target install + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install } build_core_libraries_and_devtools() { @@ -119,7 +119,7 @@ build_core_libraries_and_devtools() { -DEXECUTORCH_BUILD_VULKAN=ON \ -DEXECUTORCH_BUILD_XNNPACK=ON \ -Bcmake-out && \ - cmake --build cmake-out -j64 --target install + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install # Build devtools example runner cmake examples/devtools \ @@ -127,7 +127,7 @@ build_core_libraries_and_devtools() { -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \ -DEXECUTORCH_BUILD_VULKAN=ON \ -Bcmake-out/examples/devtools && \ - cmake --build cmake-out/examples/devtools -j16 --config Release + cmake --build cmake-out/examples/devtools -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release } run_example_runner() { diff --git a/backends/vulkan/test/scripts/test_op.sh b/backends/vulkan/test/scripts/test_op.sh index 797089e54dc..d3967f98abf 100755 --- a/backends/vulkan/test/scripts/test_op.sh +++ b/backends/vulkan/test/scripts/test_op.sh @@ -149,7 +149,7 @@ build_core_libraries() { -DEXECUTORCH_BUILD_XNNPACK=ON \ -DEXECUTORCH_BUILD_TESTS=ON \ -Bcmake-out && \ - cmake --build cmake-out -j64 --target install + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install } build_operator_tests() { @@ -195,13 +195,13 @@ build_operator_tests() { # Build operator tests cmake "${CMAKE_ARGS[@]}" \ -Bcmake-out/backends/vulkan/test/op_tests && \ - cmake --build cmake-out/backends/vulkan/test/op_tests -j16 + cmake --build cmake-out/backends/vulkan/test/op_tests -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) } recompile() { echo "Recompiling..." - cmake --build cmake-out -j64 --target install - cmake --build cmake-out/backends/vulkan/test/op_tests -j16 + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install + cmake --build cmake-out/backends/vulkan/test/op_tests -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) } run_operator_test() { diff --git a/backends/vulkan/test/test_serialization.py b/backends/vulkan/test/test_serialization.py index 71a6980635a..513150ba9fe 100644 --- a/backends/vulkan/test/test_serialization.py +++ b/backends/vulkan/test/test_serialization.py @@ -19,6 +19,8 @@ ) from executorch.backends.vulkan.serialization.vulkan_graph_schema import ( + Double, + DoubleList, IntList, OperatorCall, String, @@ -269,3 +271,77 @@ def test_serialize_deserialize_vkgraph(self) -> None: out_vk_graph = flatbuffer_to_vk_graph(bs) self.assertEqual(in_vk_graph, out_vk_graph) + + def _round_trip(self, values, chain=None) -> VkGraph: + in_vk_graph = VkGraph( + version="1", + chain=chain if chain is not None else [], + values=values, + input_ids=[], + output_ids=[], + constants=[], + shaders=[], + ) + out_vk_graph = flatbuffer_to_vk_graph(convert_to_flatbuffer(in_vk_graph)) + self.assertEqual(in_vk_graph, out_vk_graph) + return out_vk_graph + + def test_serialize_deserialize_non_finite_scalars(self) -> None: + # Python's json module spells the infinities "Infinity" / "-Infinity" + # while flatc spells them "inf" / "-inf" and rejects Python's spelling, + # so both directions need translating. A graph picks up a non-finite + # scalar whenever the model has one -- the -inf fill value of a + # transformer attention mask being the usual source. + self._round_trip( + [ + VkValue(value=Double(double_val=float("-inf"))), + VkValue(value=Double(double_val=float("inf"))), + VkValue(value=Double(double_val=1.5)), + ] + ) + + def test_serialize_deserialize_non_finite_floats_in_list(self) -> None: + # json only emits a float as a chunk of its own inside an object; in a + # list the chunk carries the delimiter with it, so a rewrite that works + # on the scalar above can still miss every element of a DoubleList. + self._round_trip( + [ + VkValue(value=DoubleList(items=[float("-inf")])), + VkValue( + value=DoubleList(items=[1.5, float("inf"), 2.5, float("-inf")]) + ), + VkValue(value=DoubleList(items=[])), + ] + ) + + def test_serialize_nan_float_raises(self) -> None: + # flatc rejects nan, NaN and Nan alike for a value inside a union, and + # every float in the Vulkan schema is a member of the VkValue union, so + # report it here rather than emitting JSON that flatc cannot read. + for value in ( + Double(double_val=float("nan")), + DoubleList(items=[1.0, float("nan")]), + ): + vk_graph = VkGraph( + version="1", + chain=[], + values=[VkValue(value=value)], + input_ids=[], + output_ids=[], + constants=[], + shaders=[], + ) + with self.assertRaisesRegex(ValueError, "NaN"): + convert_to_flatbuffer(vk_graph) + + def test_serialize_deserialize_leaves_strings_alone(self) -> None: + # The token rewrites run over the serialized JSON, so they must not + # reach into string literals in either direction. + self._round_trip( + [ + VkValue(value=String(string_val="value: inf, -inf")), + VkValue(value=String(string_val="Infinity NaN nan")), + VkValue(value=String(string_val='quoted "inf" and \\ inf')), + ], + chain=[OperatorCall(node_id=1, name="inf_shader", args=[])], + ) diff --git a/backends/vulkan/test/test_vulkan_delegate.py b/backends/vulkan/test/test_vulkan_delegate.py index c6915d37684..05c5084718b 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -13,6 +13,7 @@ import executorch.backends.vulkan.test.utils as test_utils import torch +import torch.nn.functional as F from executorch.backends.transforms.convert_dtype_pass import I64toI32 from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.backends.vulkan.vulkan_preprocess import VulkanBackend @@ -538,6 +539,20 @@ def forward(self, x): self.lower_module_and_test_output(ClampModule(), sample_inputs) + def test_vulkan_backend_dynamic_float_clamp(self): + class ClampModule(torch.nn.Module): + def forward(self, x): + return torch.clamp(x, max=x.shape[0]) + + sample_inputs = (torch.arange(32).reshape(8, 4).float(),) + length = Dim("length", min=2, max=16) + self.lower_module_and_test_output( + ClampModule(), + sample_inputs, + dynamic_shapes={"x": {0: length}}, + test_inputs=[(torch.arange(12).reshape(3, 4).float(),)], + ) + def test_vulkan_backend_cos(self): class CosModule(torch.nn.Module): def __init__(self): @@ -1065,6 +1080,20 @@ def forward(self, x): sample_inputs, ) + def test_vulkan_backend_binary_op_zero_dim(self): + # Both operands, and therefore the output, are 0-dimensional. This is + # what a reduction to a scalar followed by arithmetic produces, e.g. the + # log-mel normalisation in Whisper's preprocessor. + class ZeroDimModule(torch.nn.Module): + def forward(self, x): + m = x.max() + return (m - (m - 1.0)).reshape(1) + + self.lower_module_and_test_output( + ZeroDimModule(), + (torch.randn(size=(64,), dtype=torch.float32),), + ) + @disable_test("layer norm compute shader not working with swiftshader") def test_vulkan_backend_native_layer_norm(self): class NativeLayerNormModule(torch.nn.Module): @@ -1464,6 +1493,32 @@ def forward(self, x): sample_inputs, ) + def test_vulkan_backend_constant_pad_nd_symbolic_pad(self): + """A pad amount derived from a dynamic dim, as LSTM padding produces. + + Without the guard this partitions and then aborts at prepack with + "Expected value to have type IntList, got VALUELIST instead", because + the pad list is serialized as a VALUELIST of Int/SymInt. + """ + + class TestModule(torch.nn.Module): + def forward(self, x): + # Pad up to a static length, the shape every unrolled LSTM + # wants its input in. + return torch.nn.functional.pad(x, (0, 0, 0, 16 - x.shape[1])) + + sample_inputs = (torch.randn(size=(1, 12, 8), dtype=torch.float32),) + seq = Dim("seq", min=2, max=16) + self.lower_module_and_test_output( + TestModule(), + sample_inputs, + dynamic_shapes={"x": {1: seq}}, + test_inputs=[ + (torch.randn(size=(1, 4, 8), dtype=torch.float32),), + (torch.randn(size=(1, 16, 8), dtype=torch.float32),), + ], + ) + def test_vulkan_backend_repeat(self): class TestModule(torch.nn.Module): def __init__(self): @@ -1625,6 +1680,190 @@ def forward(self, x): sample_inputs, ) + def test_vulkan_backend_index_select_batch_dynamic_channels(self): + # Selecting along the batch dim walks the z axis in units of channel + # texels, so the step depends on the channel count. Vary the channels + # below the built size: a step frozen at build time reads the wrong + # texel once the count drops. + class IndexSelectModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.index = torch.tensor([1, 3, 0, 2]) + + def forward(self, x): + return torch.index_select(x, 0, self.index) + + sample_inputs = (torch.randn(size=(5, 8, 3, 4), dtype=torch.float32),) + dynamic_shapes = {"x": {1: Dim("channels", min=1, max=8)}} + test_inputs = [ + (torch.randn(5, 1, 3, 4),), + (torch.randn(5, 3, 3, 4),), + (torch.randn(5, 4, 3, 4),), + (torch.randn(5, 5, 3, 4),), + (torch.randn(5, 8, 3, 4),), + ] + + self.lower_module_and_test_output( + IndexSelectModule(), + sample_inputs, + dynamic_shapes=dynamic_shapes, + test_inputs=test_inputs, + ) + + def test_vulkan_backend_index_select_width_dynamic_shapes(self): + class IndexSelectModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.index = torch.tensor([2, 0, 1]) + + def forward(self, x): + return torch.index_select(x, 2, self.index) + + sample_inputs = (torch.randn(size=(2, 3, 4, 6), dtype=torch.float32),) + dynamic_shapes = {"x": {3: Dim("width", min=1, max=6)}} + test_inputs = [ + (torch.randn(2, 3, 4, 1),), + (torch.randn(2, 3, 4, 3),), + (torch.randn(2, 3, 4, 6),), + ] + + self.lower_module_and_test_output( + IndexSelectModule(), + sample_inputs, + dynamic_shapes=dynamic_shapes, + test_inputs=test_inputs, + ) + + def test_vulkan_backend_index_select_channel_dynamic_shapes(self): + # The channel path takes a separate shader and a separate resize + # callback from the one above. + class IndexSelectModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.index = torch.tensor([3, 1, 0, 1]) + + def forward(self, x): + return torch.index_select(x, 1, self.index) + + sample_inputs = (torch.randn(size=(2, 5, 4, 6), dtype=torch.float32),) + dynamic_shapes = {"x": {3: Dim("width", min=1, max=6)}} + test_inputs = [ + (torch.randn(2, 5, 4, 1),), + (torch.randn(2, 5, 4, 4),), + (torch.randn(2, 5, 4, 6),), + ] + + self.lower_module_and_test_output( + IndexSelectModule(), + sample_inputs, + dynamic_shapes=dynamic_shapes, + test_inputs=test_inputs, + ) + + def test_vulkan_backend_embedding_dynamic_shapes(self): + # The output picks up the index tensor's shape, so it has to be resized + # with it rather than left at the size it was built with. + class EmbeddingModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.embedding = torch.nn.Embedding(10, 8) + + def forward(self, x): + return self.embedding(x) + + sample_inputs = (torch.randint(0, 10, (2, 6), dtype=torch.int32),) + dynamic_shapes = {"x": {1: Dim("seq", min=1, max=6)}} + test_inputs = [ + (torch.randint(0, 10, (2, 1), dtype=torch.int32),), + (torch.randint(0, 10, (2, 3), dtype=torch.int32),), + (torch.randint(0, 10, (2, 6), dtype=torch.int32),), + ] + + self.lower_module_and_test_output( + EmbeddingModule(), + sample_inputs, + dynamic_shapes=dynamic_shapes, + test_inputs=test_inputs, + ) + + def test_vulkan_backend_index_tensor_nonzero_axis(self): + class IndexTensorModule(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + self.index = torch.tensor([0, 2]) + + def forward(self, x): + indices = [slice(None)] * x.dim() + indices[self.dim] = self.index + return x[tuple(indices)] + + sample_inputs = (torch.arange(24).reshape(1, 3, 8).float(),) + for dim in (1, 2): + self.lower_module_and_test_output( + IndexTensorModule(dim), + sample_inputs, + ) + + def test_vulkan_backend_dynamic_replicate_pad_time_reduction(self): + class TimeReductionModule(torch.nn.Module): + def forward(self, x): + padded_frames = 8 * ((x.shape[1] + 7) // 8) + x = F.pad( + x, + (0, 0, 0, padded_frames - x.shape[1]), + mode="replicate", + ) + return x.view(x.shape[0], -1, 640) + + sample_inputs = (torch.randn(1, 24, 80),) + frames = Dim("frames", min=1, max=24) + self.lower_module_and_test_output( + TimeReductionModule(), + sample_inputs, + dynamic_shapes={"x": {1: frames}}, + test_inputs=[ + (torch.randn(1, 8, 80),), + (torch.randn(1, 9, 80),), + (torch.randn(1, 17, 80),), + ], + ) + + def test_vulkan_backend_dynamic_arange_float_step(self): + class ArangeModule(torch.nn.Module): + def __init__(self, end_scale, step): + super().__init__() + self.end_scale = end_scale + self.step = step + + def forward(self, x): + return torch.arange(0, self.end_scale * x.shape[0], self.step) + + sample_inputs = (torch.randn(8),) + length = Dim("length", min=2, max=16) + for end_scale, step in ((1, 0.5), (-1, -0.5)): + with self.subTest(end_scale=end_scale, step=step): + self.lower_module_and_test_output( + ArangeModule(end_scale, step), + sample_inputs, + dynamic_shapes={"x": {0: length}}, + test_inputs=[(torch.randn(3),), (torch.randn(7),)], + ) + + def test_vulkan_backend_dynamic_arange_start(self): + class ArangeModule(torch.nn.Module): + def forward(self, x): + return torch.arange(x.shape[0], 32, 2) + + sample_inputs = (torch.randn(8),) + length = Dim("length", min=2, max=16) + self.lower_module_and_test_output( + ArangeModule(), + sample_inputs, + dynamic_shapes={"x": {0: length}}, + test_inputs=[(torch.randn(3),), (torch.randn(15),)], + ) + def test_vulkan_backend_arange_int(self): class ArangeModule(torch.nn.Module): def __init__(self, input): diff --git a/backends/vulkan/test/test_vulkan_graph_builder.py b/backends/vulkan/test/test_vulkan_graph_builder.py new file mode 100644 index 00000000000..65afc3a2542 --- /dev/null +++ b/backends/vulkan/test/test_vulkan_graph_builder.py @@ -0,0 +1,60 @@ +# 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. + +import unittest + +import torch +from executorch.backends.vulkan.serialization.vulkan_graph_builder import VkGraphBuilder +from executorch.backends.vulkan.vulkan_preprocess import apply_passes +from executorch.exir import to_edge +from executorch.exir.backend.utils import DelegateMappingBuilder +from executorch.exir.passes import SpecPropPass + + +class TestVkGraphBuilderInputIds(unittest.TestCase): + """The serialized input list has to match the delegate call's arguments. + + VulkanBackend::execute walks `args` positionally against + ComputeGraph::inputs() and rejects the call when the counts disagree, so + every placeholder that the delegate call passes must appear in input_ids, + including ones this graph happens not to use. Unused placeholders are not + hypothetical: passes that run after partitioning can fold away a + placeholder's only consumers, leaving the argument list and the serialized + graph out of step. + """ + + def _build(self, module: torch.nn.Module, inputs) -> VkGraphBuilder: + edge = to_edge(torch.export.export(module, inputs, strict=True)) + # The builder reads node specs, which the backend's own preprocess + # populates before it gets here. + program = apply_passes(edge.exported_program(), [SpecPropPass()]) + builder = VkGraphBuilder( + program, DelegateMappingBuilder(generated_identifiers=True) + ) + builder.build_graph() + return builder + + def test_unused_placeholder_is_still_declared_as_an_input(self) -> None: + class UsesOnlyTheFirstInput(torch.nn.Module): + def forward(self, used, unused): + return used + used + + builder = self._build( + UsesOnlyTheFirstInput(), (torch.randn(2, 3), torch.randn(2, 3)) + ) + self.assertEqual(len(builder.input_ids), 2) + + def test_used_placeholders_are_declared_in_order(self) -> None: + class UsesBothInputs(torch.nn.Module): + def forward(self, first, second): + return first + second + + builder = self._build(UsesBothInputs(), (torch.randn(2, 3), torch.randn(2, 3))) + self.assertEqual(len(builder.input_ids), 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/vulkan/test/test_vulkan_shader_names.py b/backends/vulkan/test/test_vulkan_shader_names.py new file mode 100644 index 00000000000..1535e6c9653 --- /dev/null +++ b/backends/vulkan/test/test_vulkan_shader_names.py @@ -0,0 +1,148 @@ +# 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. + +"""Checks that generated shader names match the names the dispatcher asks for. + +A kernel name is built at runtime rather than looked up: `add_binary_op_node` +concatenates "binary_", the op, a storage suffix and a dtype suffix, then hands +the result to `VK_KERNEL_FROM_STR`. The yaml, meanwhile, is free to name a +variant anything at all. Nothing connects the two, so a variant whose name does +not follow that shape compiles a shader nothing references and leaves the name +the dispatcher wants missing. + +That failure is invisible until dispatch. The op is still registered as +supported, so the partitioner claims it, the export succeeds, and the model +aborts on device with "Could not find ShaderInfo with name ...". These tests +close the gap at codegen time instead. + +The generator is loaded by file path, and the suffixes are read out of the C++ +that produces them, so this needs neither a built runtime nor a GPU and does not +drift when a dtype is added. +""" + +import importlib.util +import re +import unittest +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_VULKAN_ROOT = _REPO_ROOT / "backends" / "vulkan" +_GLSL_DIR = _VULKAN_ROOT / "runtime" / "graph" / "ops" / "glsl" +_SHADER_NAME_UTILS = ( + _VULKAN_ROOT / "runtime" / "graph" / "ops" / "utils" / "ShaderNameUtils.cpp" +) + +_spec = importlib.util.spec_from_file_location( + "gen_vulkan_spv", _VULKAN_ROOT / "runtime" / "gen_vulkan_spv.py" +) +gen_vulkan_spv = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(gen_vulkan_spv) + +# The yaml templates whose variants `add_binary_op_node` dispatches into. +_BINARY_TEMPLATES = ("binary_op_buffer", "binary_op_texture") + + +def _function_body(text: str, name: str) -> str: + """Source of a top-level C++ function, from its signature to its closing brace.""" + start = next( + ( + i + for i, line in enumerate(text.splitlines()) + if re.match(rf"^\w[\w:<>&* ]*\b{re.escape(name)}\(", line) + ), + None, + ) + if start is None: + raise AssertionError(f"{name} not found in {_SHADER_NAME_UTILS}") + lines = text.splitlines()[start:] + end = next(i for i, line in enumerate(lines) if line == "}") + return "\n".join(lines[: end + 1]) + + +def _suffixes(function_name: str) -> tuple: + """Every literal suffix one ShaderNameUtils.cpp helper can append. + + Read from the C++ rather than restated here: a dtype added to + `add_dtype_suffix` becomes legal in a shader name the moment it is added, + and a test carrying its own copy of the list would reject it. + """ + body = _function_body(_SHADER_NAME_UTILS.read_text(), function_name) + found = re.findall(r'kernel_name \+= "(_[a-z0-9]+)";', body) + if not found: + raise AssertionError(f"no suffixes parsed out of {function_name}") + return tuple(dict.fromkeys(found)) + + +STORAGE_SUFFIXES = _suffixes("add_storage_type_suffix") +DTYPE_SUFFIXES = _suffixes("add_dtype_suffix") + + +def _generated_names() -> dict: + """Variant names the shader codegen produces, keyed by yaml template.""" + env = dict(gen_vulkan_spv.DEFAULT_ENV) + env.update(gen_vulkan_spv.TYPE_MAPPINGS) + env.update(gen_vulkan_spv.UTILITY_FNS) + generator = gen_vulkan_spv.SPVGenerator([str(_GLSL_DIR)], env, glslc_path=None) + return { + template: [variant["NAME"] for variant in variants] + for template, variants in generator.shader_template_params.items() + } + + +class TestShaderNames(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.names = _generated_names() + + def test_suffixes_are_parsed_from_the_cpp(self) -> None: + # Guards the two tests below: a parse that silently returned something + # empty or wrong would make them pass by vacuously accepting any name. + self.assertIn("_buffer", STORAGE_SUFFIXES) + self.assertIn("_texture3d", STORAGE_SUFFIXES) + self.assertIn("_float", DTYPE_SUFFIXES) + self.assertIn("_int32", DTYPE_SUFFIXES) + + def test_int32_eq_shaders_are_named_for_the_dispatcher(self) -> None: + """An int32 `aten.eq.Tensor` must find a shader on both storage types. + + Declared as `binary_eq_int32_{buffer,texture3d}` for a while, which put + the dtype in the middle and so generated a name no dispatch could ever + build. Kokoro's synthesizer aborted on it. + """ + for template, expected in ( + ("binary_op_buffer", "binary_eq_buffer_int32"), + ("binary_op_texture", "binary_eq_texture3d_int32"), + ): + with self.subTest(template=template): + self.assertIn(expected, self.names[template]) + + def test_binary_variants_end_in_a_storage_and_dtype_suffix(self) -> None: + """The general form of the same bug, for every binary op at once. + + `add_binary_op_node` appends the storage suffix and then the dtype + suffix, in that order, to every name it builds. A generated variant that + does not end that way cannot be reached from the dispatcher, whatever + else is true about it. + """ + legal = tuple( + storage + dtype for storage in STORAGE_SUFFIXES for dtype in DTYPE_SUFFIXES + ) + unreachable = [ + name + for template in _BINARY_TEMPLATES + for name in self.names[template] + if not name.endswith(legal) + ] + self.assertEqual( + unreachable, + [], + "these shaders are generated but no dispatch can name them; " + "see add_binary_op_node in BinaryOp.cpp", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/vulkan/test/test_vulkan_tensor_repr.py b/backends/vulkan/test/test_vulkan_tensor_repr.py index 5a0fc664c17..83d68195407 100644 --- a/backends/vulkan/test/test_vulkan_tensor_repr.py +++ b/backends/vulkan/test/test_vulkan_tensor_repr.py @@ -605,6 +605,29 @@ def test_unary_op_construction(self): self.assertEqual(op_repsets.primary_arg_idx, 0) self.assertTrue(op_repsets.sync_primary_io_repr) + def test_single_tensor_output_in_list_construction(self): + """An op declared to return Tensor[] that yields exactly one tensor. + + num_tensors_in_node() counts tensors rather than nesting, so such a + node reports 1 while meta["val"] is a one-element list rather than a + bare FakeTensor. + """ + arg = _make_tensor_arg_node((1, 3, 8, 8)) + node = _make_op_node( + target=torch.ops.aten.split_with_sizes_copy.default, + args=(arg, [3]), + output_val=[_make_fake_tensor((1, 3, 8, 8))], + ) + + op_repsets = OpRepSets( + TensorRepSetList(ANY_STORAGE), + TensorRepSetList(ANY_STORAGE), + node, + DEFAULT_TEXTURE_LIMITS, + ) + + self.assertFalse(op_repsets.any_is_empty()) + def test_binary_op_syncs_args(self): """When a single repset covers all inputs, sync_args_repr is True.""" op_repsets = self._make_binary_op() diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index 95776e42304..a0ca49cff6b 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -32,6 +32,7 @@ #include #include +#include #include @@ -2152,6 +2153,443 @@ TEST(VulkanComputeGraphTest, test_simple_graph_with_symint) { } } +TEST(VulkanComputeGraphTest, was_value_updated_tracks_tensor_changes) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef tensor = graph.add_tensor({2, 4}, vkapi::kFloat); + + EXPECT_FALSE(graph.was_value_updated(kDummyValueRef)); + EXPECT_FALSE(graph.was_value_updated(tensor)); + + graph.virtual_resize(tensor, {2, 4}); + EXPECT_FALSE(graph.was_value_updated(tensor)); + + graph.virtual_resize(tensor, {1, 4}); + EXPECT_TRUE(graph.was_value_updated(tensor)); +} + +TEST(VulkanComputeGraphTest, was_value_updated_tracks_symint_changes) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef symint = graph.add_symint(3); + + EXPECT_FALSE(graph.was_value_updated(symint)); + + graph.set_symint(symint, 3); + EXPECT_FALSE(graph.was_value_updated(symint)); + + graph.set_symint(symint, 5); + EXPECT_TRUE(graph.was_value_updated(symint)); +} + +TEST(VulkanComputeGraphTest, was_value_updated_checks_nested_value_lists) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef unchanged = graph.add_symint(1); + const ValueRef changed = graph.add_symint(2); + const ValueRef inner_list = graph.add_value_list({unchanged, changed}); + const ValueRef outer_list = + graph.add_value_list({kDummyValueRef, inner_list}); + + EXPECT_FALSE(graph.was_value_updated(inner_list)); + EXPECT_FALSE(graph.was_value_updated(outer_list)); + + graph.set_symint(changed, 3); + + EXPECT_FALSE(graph.was_value_updated(unchanged)); + EXPECT_TRUE(graph.was_value_updated(changed)); + EXPECT_TRUE(graph.was_value_updated(inner_list)); + EXPECT_TRUE(graph.was_value_updated(outer_list)); +} + +TEST(VulkanComputeGraphTest, execute_node_resize_tracks_read_arg_updates) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef output = graph.add_symint(1); + const ValueRef input = graph.add_symint(2); + size_t resize_count = 0; + ExecuteNode node( + [&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }, + {}, + {{output, vkapi::kWrite}, {input, vkapi::kRead}}); + + graph.set_symint(input, 3); + + EXPECT_TRUE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 1); +} + +TEST(VulkanComputeGraphTest, execute_node_resize_tracks_write_arg_updates) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef output = graph.add_symint(1); + const ValueRef input = graph.add_symint(2); + size_t resize_count = 0; + ExecuteNode node( + [&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }, + {}, + {{output, vkapi::kWrite}, {input, vkapi::kRead}}); + + graph.set_symint(output, 3); + + EXPECT_TRUE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 1); +} + +TEST(VulkanComputeGraphTest, execute_node_resize_tracks_read_write_updates) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef value = graph.add_symint(1); + size_t resize_count = 0; + ExecuteNode node( + [&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }, + {}, + {{value, vkapi::kReadWrite}}); + + graph.set_symint(value, 2); + + EXPECT_TRUE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 1); +} + +TEST(VulkanComputeGraphTest, execute_node_resize_tracks_nested_resize_args) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef value = graph.add_symint(1); + const ValueRef inner_list = graph.add_value_list({value}); + const ValueRef outer_list = graph.add_value_list({inner_list}); + size_t resize_count = 0; + ExecuteNode node( + [&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }, + {outer_list}); + + graph.set_symint(value, 2); + + EXPECT_TRUE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 1); +} + +TEST(VulkanComputeGraphTest, execute_node_resize_skips_unchanged_args) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef output = graph.add_symint(1); + const ValueRef input = graph.add_symint(2); + size_t resize_count = 0; + ExecuteNode node( + [&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }, + {}, + {{output, vkapi::kWrite}, {input, vkapi::kRead}}); + + EXPECT_FALSE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 0); +} + +TEST(VulkanComputeGraphTest, execute_node_force_resize_ignores_arg_updates) { + GraphConfig config; + config.force_resize = true; + ComputeGraph graph(config); + + size_t resize_count = 0; + ExecuteNode node([&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }); + + EXPECT_TRUE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 1); +} + +TEST( + VulkanComputeGraphTest, + execute_node_data_dependent_resize_is_unconditional) { + GraphConfig config; + ComputeGraph graph(config); + + size_t resize_count = 0; + ExecuteNode node( + [&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }, + {}, + {}, + "data_dependent_node", + true); + + EXPECT_TRUE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 1); +} + +TEST(VulkanComputeGraphTest, resize_input_marks_staging_value_updated) { + GraphConfig config; + ComputeGraph graph(config); + + const IOValueRef input = graph.add_input_tensor({2, 4}, vkapi::kFloat); + + EXPECT_FALSE(graph.was_value_updated(input.value)); + EXPECT_FALSE(graph.was_value_updated(input.staging)); + + graph.resize_input(0, {2, 4}); + + EXPECT_FALSE(graph.was_value_updated(input.value)); + EXPECT_TRUE(graph.was_value_updated(input.staging)); +} + +TEST(VulkanComputeGraphTest, execute_advances_value_update_generation) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef symint = graph.add_symint(1); + const ValueRef values = graph.add_value_list({symint}); + + graph.prepare(); + graph.set_symint(symint, 2); + + EXPECT_TRUE(graph.was_value_updated(symint)); + EXPECT_TRUE(graph.was_value_updated(values)); + + graph.execute(); + + EXPECT_FALSE(graph.was_value_updated(symint)); + EXPECT_FALSE(graph.was_value_updated(values)); + + graph.set_symint(symint, 3); + + EXPECT_TRUE(graph.was_value_updated(symint)); + EXPECT_TRUE(graph.was_value_updated(values)); + + graph.execute(); + + EXPECT_FALSE(graph.was_value_updated(symint)); + EXPECT_FALSE(graph.was_value_updated(values)); +} + +TEST(VulkanComputeGraphTest, choose_qparams_handles_dynamic_row_counts) { + constexpr int64_t kMaxM = 8; + constexpr int64_t kK = 128; + + GraphConfig config; + config.enable_querypool = true; + config.expect_dynamic_shapes = true; + ComputeGraph graph(config); + + const IOValueRef input = + graph.add_input_tensor({kMaxM, kK}, vkapi::kFloat, utils::kBuffer); + const ValueRef quant_min = graph.add_scalar(-128); + const ValueRef quant_max = graph.add_scalar(127); + const ValueRef scales = graph.add_tensor( + {kMaxM}, vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked); + const ValueRef zero_points = graph.add_tensor( + {kMaxM}, vkapi::kChar, utils::kTexture3D, utils::kWidthPacked); + + VK_GET_OP_FN("etvk.choose_qparams_per_row.default") + (graph, {input.value, quant_min, quant_max, scales, zero_points}); + + const ValueRef scales_staging = graph.set_output_tensor(scales); + const ValueRef zero_points_staging = graph.set_output_tensor(zero_points); + + graph.prepare(); + graph.prepack(); + + for (const int64_t M : std::vector{kMaxM, 4, 1, kMaxM}) { + graph.resize_input(0, {M, kK}); + graph.propagate_resize(); + + EXPECT_EQ(graph.sizes_of(scales), std::vector({M})); + EXPECT_EQ(graph.sizes_of(zero_points), std::vector({M})); + + std::vector input_data(M * kK); + for (int64_t m = 0; m < M; ++m) { + std::fill_n(input_data.begin() + m * kK, kK, float(m + 1)); + } + graph.maybe_cast_and_copy_into_staging( + input.staging, input_data.data(), input_data.size(), vkapi::kFloat); + + graph.execute(); + + std::vector scale_data(M); + std::vector zero_point_data(M); + graph.maybe_cast_and_copy_from_staging( + scales_staging, scale_data.data(), scale_data.size(), vkapi::kFloat); + graph.maybe_cast_and_copy_from_staging( + zero_points_staging, + zero_point_data.data(), + zero_point_data.size(), + vkapi::kChar); + + for (int64_t m = 0; m < M; ++m) { + EXPECT_NEAR(scale_data[m], float(m + 1) / 255.0f, 1e-6f); + EXPECT_EQ(zero_point_data[m], -128); + } + + graph.context()->querypool().extract_results(); + const auto shader_results = + graph.context()->querypool().get_shader_timestamp_data(); + const auto choose_result = std::find_if( + shader_results.begin(), shader_results.end(), [](const auto& result) { + return result.kernel_name.find("choose_qparams_per_row") != + std::string::npos; + }); + ASSERT_NE(choose_result, shader_results.end()); + EXPECT_EQ(choose_result->metadata.gwg[0], 1u); + EXPECT_EQ( + choose_result->metadata.gwg[1], + utils::div_up_4(utils::safe_downcast(M))); + EXPECT_EQ(choose_result->metadata.gwg[2], 1u); + EXPECT_EQ(choose_result->metadata.lwg[0], 64u); + EXPECT_EQ(choose_result->metadata.lwg[1], 1u); + EXPECT_EQ(choose_result->metadata.lwg[2], 1u); + } +} + +void test_quantize_and_pack_handles_dynamic_row_counts( + const int64_t group_size_value, + const utils::uvec3& expected_local_wg_size) { + if (!api::context()->adapter_ptr()->supports_int8_dot_product()) { + GTEST_SKIP() << "Quantize and pack requires integer dot product support"; + } + + constexpr int64_t kMaxM = 8; + constexpr int64_t kK = 128; + const int64_t num_groups = kK / group_size_value; + const int64_t max_m4 = utils::div_up(kMaxM, int64_t(4)); + + GraphConfig config; + config.enable_querypool = true; + config.expect_dynamic_shapes = true; + ComputeGraph graph(config); + + const IOValueRef input = + graph.add_input_tensor({kMaxM, kK}, vkapi::kFloat, utils::kBuffer); + const ValueRef quant_min = graph.add_scalar(-128); + const ValueRef quant_max = graph.add_scalar(127); + const ValueRef scales = graph.add_tensor( + {kMaxM}, vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked); + const ValueRef zero_points = graph.add_tensor( + {kMaxM}, vkapi::kChar, utils::kTexture3D, utils::kWidthPacked); + + VK_GET_OP_FN("etvk.choose_qparams_per_row.default") + (graph, {input.value, quant_min, quant_max, scales, zero_points}); + + const ValueRef packed_input = graph.add_tensor( + {kMaxM, kK}, vkapi::kInt8x4, utils::kBuffer, utils::kPackedInt8_4H4W); + const ValueRef input_sums = graph.add_tensor( + {num_groups * max_m4 * 4}, + vkapi::kInt, + utils::kBuffer, + utils::kWidthPacked); + const ValueRef group_size = graph.add_scalar(group_size_value); + const QuantizationConfig input_quant_config( + 8, kPerChannel, {1, kK}, false, true); + + add_quantize_and_pack_4h4w_with_group_sums_node( + graph, + input_quant_config, + input.value, + input_sums, + scales, + zero_points, + packed_input, + group_size); + + const ValueRef packed_input_staging = graph.set_output_tensor(packed_input); + const ValueRef input_sums_staging = graph.set_output_tensor(input_sums); + + graph.prepare(); + graph.prepack(); + + for (const int64_t M : std::vector{kMaxM, 4, 1, kMaxM}) { + graph.resize_input(0, {M, kK}); + graph.propagate_resize(); + + std::vector input_data(M * kK); + for (int64_t m = 0; m < M; ++m) { + std::fill_n(input_data.begin() + m * kK, kK, float(m + 1)); + } + graph.maybe_cast_and_copy_into_staging( + input.staging, input_data.data(), input_data.size(), vkapi::kFloat); + + graph.execute(); + + graph.context()->querypool().extract_results(); + const auto shader_results = + graph.context()->querypool().get_shader_timestamp_data(); + const auto quantize_result = std::find_if( + shader_results.begin(), shader_results.end(), [](const auto& result) { + return result.kernel_name.find( + "quantize_and_pack_4h4w_with_group_sums") != + std::string::npos; + }); + + if (M == 1) { + EXPECT_EQ(quantize_result, shader_results.end()); + continue; + } + + ASSERT_NE(quantize_result, shader_results.end()); + EXPECT_EQ( + quantize_result->metadata.gwg[0], + utils::safe_downcast(num_groups)); + EXPECT_EQ( + quantize_result->metadata.gwg[1], + utils::div_up_4(utils::safe_downcast(M))); + EXPECT_EQ(quantize_result->metadata.gwg[2], 1u); + EXPECT_EQ(quantize_result->metadata.lwg[0], expected_local_wg_size[0]); + EXPECT_EQ(quantize_result->metadata.lwg[1], expected_local_wg_size[1]); + EXPECT_EQ(quantize_result->metadata.lwg[2], expected_local_wg_size[2]); + + const size_t packed_numel = graph.staging_buffer_numel_of(packed_input); + std::vector packed_data(packed_numel); + graph.maybe_cast_and_copy_from_staging( + packed_input_staging, + packed_data.data(), + packed_data.size(), + vkapi::kInt8x4); + for (int64_t i = 0; i < M * kK / 4; ++i) { + EXPECT_EQ(packed_data[i], 0x7f7f7f7f); + } + + std::vector sums_data(num_groups * max_m4 * 4); + graph.maybe_cast_and_copy_from_staging( + input_sums_staging, sums_data.data(), sums_data.size(), vkapi::kInt); + const int64_t current_m4 = utils::div_up(M, int64_t(4)); + for (int64_t group = 0; group < num_groups; ++group) { + for (int64_t m = 0; m < M; ++m) { + EXPECT_EQ( + sums_data[group * current_m4 * 4 + m], 127 * group_size_value); + } + } + } +} + +TEST( + VulkanComputeGraphTest, + quantize_and_pack_handles_dynamic_row_counts_with_small_groups) { + test_quantize_and_pack_handles_dynamic_row_counts(32, {4u, 1u, 16u}); +} + +TEST( + VulkanComputeGraphTest, + quantize_and_pack_handles_dynamic_row_counts_with_large_groups) { + test_quantize_and_pack_handles_dynamic_row_counts(128, {2u, 1u, 32u}); +} + #define CREATE_WEIGHT_TENSOR(name, sizes, dtype, val) \ std::vector data_##name(utils::multiply_integers(sizes)); \ std::fill(data_##name.begin(), data_##name.end(), val); \ diff --git a/backends/vulkan/utils.py b/backends/vulkan/utils.py index 84b901b6b6e..c43f318ae00 100644 --- a/backends/vulkan/utils.py +++ b/backends/vulkan/utils.py @@ -1181,6 +1181,33 @@ def make_tensor_repset(tensor_repr: TensorRepr) -> TensorRepSet: raise RuntimeError(f"Unsupported storage type {tensor_repr.storage_type}") +def upper_bound_size(dim: Union[int, torch.SymInt]) -> Optional[int]: + """Largest value a (possibly symbolic) tensor dimension can take. + + Returns None if no finite bound is known. + + A symbolic dim compares against a limit using its *hint* -- the size of the + example input the model happened to be traced with -- not the maximum the + exported range allows. Sizing decisions must use the bound instead, or a + model traced with a small example will make a choice that is invalid once it + runs at a larger size. + """ + if not isinstance(dim, torch.SymInt): + return int(dim) + if not dim.node.expr.free_symbols: + return int(dim.node.expr) + shape_env = dim.node.shape_env + if shape_env is None: + return None + try: + upper = shape_env.bound_sympy(dim.node.expr).upper + except Exception: + return None + if upper is None or not upper.is_finite: + return None + return int(upper) + + def filter_invalid_reprs( tensor_val: FakeTensor, tensor_repset: TensorRepSet, @@ -1198,11 +1225,17 @@ def filter_invalid_reprs( can be used to produce a valid image texture for the given tensor (i.e. fits within texture limits). """ + # Size the texture by what the dimension CAN be, not by the example the + # model was traced with. An unbounded dim cannot be shown to fit, so it + # falls back to buffer storage. + bounds = [upper_bound_size(d) for d in tensor_val.shape] valid_texture_layouts = set() - for memory_layout in tensor_repset.valid_texture_layouts: - extents = required_image_extents(tensor_val.shape, memory_layout) - if extents_are_valid(extents, texture_limits): - valid_texture_layouts.add(memory_layout) + if all(b is not None for b in bounds): + max_shape = torch.Size(bounds) + for memory_layout in tensor_repset.valid_texture_layouts: + extents = required_image_extents(max_shape, memory_layout) + if extents_are_valid(extents, texture_limits): + valid_texture_layouts.add(memory_layout) # High dimensional tensors require buffer storage if len(tensor_val.shape) > 4: @@ -1442,8 +1475,15 @@ def __init__( # noqa: C901 outs_repset_list = TensorRepSetList([]) common_out_repset = ANY_STORAGE_INCL_PACKED_INT8 if num_tensors_in_node(op_node) == 1: + out_val = op_node.meta["val"] + # num_tensors_in_node counts tensors, not nesting: an op declared + # to return Tensor[] still lands here when it happens to produce + # exactly one, and meta["val"] is then a one-element list rather + # than a bare FakeTensor. + if isinstance(out_val, (list, tuple)): + out_val = out_val[0] common_out_repset = filter_invalid_reprs( - op_node.meta["val"], outputs_repsets[0], texture_limits + out_val, outputs_repsets[0], texture_limits ) outs_repset_list.append(common_out_repset) # Multiple output tensors diff --git a/backends/webgpu/README.md b/backends/webgpu/README.md index da23a74e8bf..5b9877b9601 100644 --- a/backends/webgpu/README.md +++ b/backends/webgpu/README.md @@ -218,4 +218,4 @@ backends/webgpu/ - **Linux:** Vulkan-capable GPU and drivers - **Browser:** A WebGPU-enabled browser; the benchmark harness uses Chrome Canary -- **Build:** CMake 3.19+ and a Python environment with ExecuTorch installed +- **Build:** CMake 3.24+ and a Python environment with ExecuTorch installed diff --git a/backends/xnnpack/_passes/__init__.py b/backends/xnnpack/_passes/__init__.py index 22147fa4215..675336f6274 100644 --- a/backends/xnnpack/_passes/__init__.py +++ b/backends/xnnpack/_passes/__init__.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. @@ -28,6 +29,9 @@ from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass from executorch.backends.xnnpack._passes.insert_pad_qdq import InsertPadQDQPass +from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import ( + LiftConstantScalarOperandsPass, +) from executorch.backends.xnnpack._passes.prelu_reshape_pass import PReLUReshapePass from executorch.backends.xnnpack._passes.propagate_custom_meta_pass import ( PropagateCustomMetaPass, @@ -35,6 +39,7 @@ from executorch.backends.xnnpack._passes.remove_redundant_copy_pass import ( RemoveRedundantCopyPass, ) +from executorch.backends.xnnpack._passes.rewrite_fp16_silu import RewriteFp16SiluPass from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass from executorch.exir.pass_base import ExportPass @@ -69,6 +74,7 @@ def __init__( if not passes: # All the XNNPACK passes self.passes = [ + RewriteFp16SiluPass, XNNPACKRemoveCloneOpsTransform, # TODO - remove this pass once we have a better support for dim_order ops lowering DimOrderOpsRevertPass, @@ -76,6 +82,7 @@ def __init__( ConvertToLinearPass, PropagateCustomMetaPass, ConvertToSDPAPass, + LiftConstantScalarOperandsPass, ConstPropPass, FuseBatchNormPass, DecomposeBatchNorm, diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index c74c35532b0..5238c9689ee 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -7,11 +7,13 @@ # pyre-unsafe from enum import Enum -from typing import Optional, Tuple +from typing import List, Optional, Tuple import torch from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass from executorch.backends.xnnpack.utils.quant_utils import ( + extract_qdq_affine_op_args_for_decomposed_ops, + is_affine_qdq, is_dequant, is_dynamic_qdq, is_tagged_as_implicit_q_dq, @@ -365,6 +367,23 @@ def input_dim_order( else ChannelsLastTaggedReshapePass.is_nhwc_node(input_node) ) + @staticmethod + def redirect_dynamic_chain_to_nhwc( + input_node: torch.fx.Node, + input_node_nhwc: torch.fx.Node, + chain: List[torch.fx.Node], + ) -> None: + """Point the traversed chain's quantize and the choose_qparams feeding it + at the NHWC copy; consumers outside the chain keep the source. + """ + quantize = chain[-1] + quantize_args = quantize.args + if is_affine_qdq(quantize): + quantize_args = extract_qdq_affine_op_args_for_decomposed_ops(quantize) + qparam = quantize_args[1].args[0] + quantize.replace_input_with(input_node, input_node_nhwc) + qparam.replace_input_with(input_node, input_node_nhwc) + def input_to_nhwc( self, graph_module: torch.fx.GraphModule, @@ -409,12 +428,15 @@ def input_to_nhwc( # Check if input uses dynamic quantization is_dynamic_input = is_dynamic_qdq(input_node) + dynamic_chain = [] if is_dynamic_input: - # Trace back to original source node. Stop if args[0] is not - # a Node (e.g., immutable_list from cat). - while getattr(input_node, "args", None) and isinstance( + # Trace back over this consumer's own q/dq chain to the source + # node, so the copy lands ahead of the quantize, and remember the + # chain: only it is redirected below. + while is_dynamic_qdq(input_node) and isinstance( input_node.args[0], torch.fx.Node ): + dynamic_chain.append(input_node) input_node = input_node.args[0] with graph_module.graph.inserting_after(input_node): @@ -427,10 +449,10 @@ def input_to_nhwc( # Use static method for consistency ChannelsLastTaggedReshapePass.mark_as_nhwc_node(input_node_nhwc) - if is_dynamic_input: - # Replace downstream input_nodes with NHWC node - input_node.replace_all_uses_with(input_node_nhwc) - input_node_nhwc.args = (input_node,) + if dynamic_chain: + self.redirect_dynamic_chain_to_nhwc( + input_node, input_node_nhwc, dynamic_chain + ) self.insert_copy_and_assign_partner_nodes_quantization_sensitive( graph_module=graph_module, diff --git a/backends/xnnpack/_passes/convert_to_sdpa.py b/backends/xnnpack/_passes/convert_to_sdpa.py index c7982db750f..b926d62df8f 100644 --- a/backends/xnnpack/_passes/convert_to_sdpa.py +++ b/backends/xnnpack/_passes/convert_to_sdpa.py @@ -1,19 +1,22 @@ # 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. import logging +from copy import deepcopy from typing import Optional import torch from executorch.backends.transforms import get_shape - +from executorch.backends.xnnpack._passes.remove_noop_expand_copy_pass import ( + RemoveNoopExpandCopyPass, +) from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass from executorch.backends.xnnpack.partition.graphs import sdpa from executorch.exir.dialects._ops import ops as exir_ops - from torch.fx.passes.infra.pass_base import PassResult from torch.fx.passes.utils.matcher_utils import InternalMatch, SubgraphMatcher @@ -24,31 +27,31 @@ class ConvertToSDPAPass(XNNPACKPass): def get_scale(self, match: InternalMatch) -> Optional[float]: """ - Returns the scale of the SDPA op. + Return the SDPA scale recovered from the matched pre-QK^T multiplications. - Scale: Optional[float] doesn't change the graph pattern. - The default value can be calulated however we need to extract - it for lowering when it is the user supplied value anyway. + The decomposition applies the square root of the attention scale before + QK^T, so the extracted multiplier is squared to recover the original value. """ for node in match.nodes_map.values(): if ( - node.op == "call_function" - and node.target == exir_ops.edge.aten.mul.Scalar + node.op != "call_function" + or node.target != exir_ops.edge.aten.mul.Scalar ): - scale = node.args[1] + continue - dtype = torch.float - mul_val = node.meta.get("val", None) - if mul_val is not None: - dtype = mul_val.dtype + scale = node.args[1] - if isinstance(scale, float): - # Convert scale value to fp16 (reducing precision) - scale = torch.tensor(scale, dtype=dtype).item() + dtype = torch.float + mul_val = node.meta.get("val", None) + if mul_val is not None: + dtype = mul_val.dtype - # since scale we extracted this before the QK^T. - return scale**2 - break + if isinstance(scale, float): + # Convert scale value to fp16 (reducing precision) + scale = torch.tensor(scale, dtype=dtype).item() + + # since scale we extracted this before the QK^T. + return scale**2 return None def assert_2d_mask(self, match: InternalMatch) -> None: @@ -99,11 +102,16 @@ def call(self, graph_module: torch.fx.GraphModule): logger.debug("ConvertToSDPA Begin: ") logger.debug(graph_module.print_readable(print_output=False)) - for pattern in sdpa.get_graphs(): - sm = SubgraphMatcher(pattern.graph, ignore_literals=True) - matches = list(sm.match(graph_module.graph)) - for partition_to_replace in matches: - self.create_sdpa(graph_module, partition_to_replace) + for scalar_pattern in sdpa.get_graphs(): + normalized_pattern = RemoveNoopExpandCopyPass()( + deepcopy(scalar_pattern) + ).graph_module + + for pattern in (scalar_pattern, normalized_pattern): + sm = SubgraphMatcher(pattern.graph, ignore_literals=True) + matches = list(sm.match(graph_module.graph)) + for partition_to_replace in matches: + self.create_sdpa(graph_module, partition_to_replace) graph_module.recompile() graph_module = super().call(graph_module).graph_module diff --git a/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py new file mode 100644 index 00000000000..01524a3268d --- /dev/null +++ b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py @@ -0,0 +1,119 @@ +# 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 numbers import Number +from typing import Dict, Optional, Union + +import torch +from executorch.backends.transforms.utils import create_constant_placeholder +from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.pass_base import PassResult +from torch._ops import OpOverload +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + +ScalarOp = Union[EdgeOpOverload, OpOverload] + + +class LiftConstantScalarOperandsPass(XNNPACKPass): + """ + Lift scalar operands into tensor constants for selected binary ops. + + XNNPACK already supports the tensor overloads for these binary operations. + This pass converts explicitly listed scalar overloads to their tensor + overloads by replacing constant scalar operands with small tensor constants. + The constants are registered as exported-program constant tensor inputs. + Keep the op map narrow until each new scalar overload is covered by tests. + """ + + default_scalar_to_tensor_ops: Dict[ScalarOp, ScalarOp] = { + exir_ops.edge.aten.mul.Scalar: exir_ops.edge.aten.mul.Tensor, + } + + def __init__( + self, + exported_program: ExportedProgram, + scalar_to_tensor_ops: Optional[Dict[ScalarOp, ScalarOp]] = None, + ) -> None: + super().__init__(exported_program) + self.scalar_to_tensor_ops = ( + scalar_to_tensor_ops + if scalar_to_tensor_ops is not None + else self.default_scalar_to_tensor_ops + ) + + def _create_constant_node( + self, + graph_module: torch.fx.GraphModule, + node: torch.fx.Node, + value: Number, + ) -> torch.fx.Node: + input_node = node.args[0] + if not isinstance(input_node, torch.fx.Node): + raise RuntimeError("Expected scalar op input to be an FX node.") + + input_value = input_node.meta["val"] + tensor = torch.tensor(value, dtype=input_value.dtype, device=input_value.device) + name = self._get_new_constant_name(graph_module) + first_placeholder = next( + graph_node + for graph_node in graph_module.graph.nodes + if graph_node.op == "placeholder" + ) + with graph_module.graph.inserting_before(first_placeholder): + return create_constant_placeholder( + self.exported_program, + graph_module.graph, + name, + InputKind.CONSTANT_TENSOR, + tensor, + ) + + def _get_new_constant_name(self, graph_module: torch.fx.GraphModule) -> str: + prefix = "_tensor_constant_" + existing_names = {node.name for node in graph_module.graph.nodes} + existing_names.update(self.exported_program.constants) + existing_names.update(self.exported_program.state_dict) + index = 0 + while f"{prefix}{index}" in existing_names: + index += 1 + return f"{prefix}{index}" + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + modified = False + + for node in list(graph_module.graph.nodes): + if ( + node.op != "call_function" + or node.target not in self.scalar_to_tensor_ops + or len(node.args) != 2 + or not isinstance(node.args[0], torch.fx.Node) + or not isinstance(node.args[1], Number) + ): + continue + + input_value = node.args[0].meta.get("val") + output_value = node.meta.get("val") + if ( + not isinstance(input_value, torch.Tensor) + or not isinstance(output_value, torch.Tensor) + or input_value.dtype != output_value.dtype + ): + continue + + tensor_arg = self._create_constant_node(graph_module, node, node.args[1]) + node.args = (node.args[0], tensor_arg) + node.target = self.scalar_to_tensor_ops[node.target] + modified = True + + graph_module.graph.eliminate_dead_code() + graph_module.graph.lint() + graph_module.recompile() + + return PassResult(graph_module, modified) diff --git a/backends/xnnpack/_passes/rewrite_fp16_silu.py b/backends/xnnpack/_passes/rewrite_fp16_silu.py new file mode 100644 index 00000000000..775dcf7bc18 --- /dev/null +++ b/backends/xnnpack/_passes/rewrite_fp16_silu.py @@ -0,0 +1,44 @@ +# 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 torch + +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult + + +class RewriteFp16SiluPass(ExportPass): + """Rewrite preserved FP16 SiLU into FP16 sigmoid and multiply.""" + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + graph = graph_module.graph + modified = False + + for node in list(graph.nodes): + if node.target != exir_ops.edge.aten.silu.default: + continue + + input_node = node.args[0] + if not isinstance(input_node, torch.fx.Node): + continue + + with graph.inserting_before(node): + sigmoid = graph.call_function( + exir_ops.edge.aten.sigmoid.default, (input_node,) + ) + mul = graph.call_function( + exir_ops.edge.aten.mul.Tensor, (input_node, sigmoid) + ) + node.replace_all_uses_with(mul) + modified = True + + if not modified: + return PassResult(graph_module, False) + + graph.eliminate_dead_code() + graph.lint() + graph_module.recompile() + graph_module = super().call(graph_module).graph_module + return PassResult(graph_module, True) diff --git a/backends/xnnpack/partition/config/__init__.py b/backends/xnnpack/partition/config/__init__.py index c6c54f083d6..02ddc64874d 100644 --- a/backends/xnnpack/partition/config/__init__.py +++ b/backends/xnnpack/partition/config/__init__.py @@ -43,6 +43,7 @@ MeanDimConfig, MinimumConfig, MulConfig, + MulScalarConfig, NegConfig, PermuteConfig, PowConfig, @@ -50,6 +51,7 @@ ReciprocalSquareRootConfig, ReLUConfig, SigmoidConfig, + SiluConfig, SinConfig, SliceCopyConfig, SoftmaxConfig, @@ -106,6 +108,7 @@ MinimumConfig, MMConfig, MulConfig, + MulScalarConfig, NegConfig, PermuteConfig, PowConfig, @@ -115,6 +118,7 @@ TanhConfig, ToDimOrderCopyConfig, SigmoidConfig, + SiluConfig, SinConfig, CosConfig, SliceCopyConfig, diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index f2b946e412d..907e972cedf 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/backends/xnnpack/partition/config/generic_node_configs.py @@ -8,6 +8,7 @@ # pyre-unsafe import logging +from numbers import Number from typing import cast, List, Optional import numpy as np @@ -409,6 +410,29 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]: return [ConfigPrecisionType.FP32] +class SiluConfig(GenericNodePartitionerConfig): + target_name = "silu.default" + + def supported_precision_types(self) -> List[ConfigPrecisionType]: + return [ConfigPrecisionType.FP32] + + def get_original_aten(self) -> Optional[torch._ops.OpOverload]: + return torch.ops.aten.silu.default + + def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: + if not self.check_common_constraints(node, ep): + return False + + input_value = node.args[0].meta.get("val") + output_value = node.meta.get("val") + return ( + isinstance(input_value, torch.Tensor) + and input_value.dtype == torch.float16 + and isinstance(output_value, torch.Tensor) + and output_value.dtype == torch.float16 + ) + + class MulConfig(GenericNodePartitionerConfig): target_name = "mul.Tensor" @@ -416,6 +440,38 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]: return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT] +class MulScalarConfig(GenericNodePartitionerConfig): + target_name = "mul.Scalar" + + def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: + if node.graph is not ep.graph: + return False + + if not self.check_common_constraints(node, ep): + return False + + if ( + len(node.args) != 2 + or not isinstance(node.args[0], torch.fx.Node) + or not isinstance(node.args[1], Number) + ): + return False + + input_value = node.args[0].meta.get("val") + output_value = node.meta.get("val") + return ( + isinstance(input_value, torch.Tensor) + and isinstance(output_value, torch.Tensor) + and input_value.dtype == output_value.dtype + ) + + def supported_precision_types(self) -> List[ConfigPrecisionType]: + return [ConfigPrecisionType.FP32] + + def get_original_aten(self) -> Optional[torch._ops.OpOverload]: + return torch.ops.aten.mul.Scalar + + class MaximumConfig(GenericNodePartitionerConfig): target_name = "maximum.default" diff --git a/backends/xnnpack/quantizer/BUCK b/backends/xnnpack/quantizer/BUCK index 72dd65c00f8..8257d561197 100644 --- a/backends/xnnpack/quantizer/BUCK +++ b/backends/xnnpack/quantizer/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "xnnpack_quantizer", diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py index fca0f1b14c6..78a989bd75c 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py @@ -1105,6 +1105,9 @@ def _is_share_obs_or_fq_op(op: Callable) -> bool: torch.ops.aten.slice.Tensor, torch.ops.aten.slice_copy.Tensor, torch.ops.aten.flatten.using_ints, + # Identity once not training, which is how a quantized model is deployed. + torch.ops.aten.dropout.default, + torch.ops.aten.dropout_.default, ] @@ -1157,15 +1160,15 @@ def _convert_scalars_to_attrs(model: torch.fx.GraphModule) -> torch.fx.GraphModu prefix = "_tensor_constant_" get_new_attr_name = get_new_attr_name_with_prefix(prefix) tensor_constant_name = get_new_attr_name(model) - float_tensor = torch.tensor(float(args[i])) - model.register_buffer(tensor_constant_name, float_tensor) + scalar_tensor = torch.tensor(args[i], dtype=n.meta["val"].dtype) + model.register_buffer(tensor_constant_name, scalar_tensor) fake_mode = n.meta["val"].fake_mode with model.graph.inserting_before(n): get_attr_node = model.graph.create_node( "get_attr", tensor_constant_name, (), {} ) get_attr_node.meta["val"] = fake_mode.from_tensor( - float_tensor, static_shapes=True + scalar_tensor, static_shapes=True ) new_args.append(get_attr_node) n.args = tuple(new_args) diff --git a/backends/xnnpack/runtime/XNNPACKBackend.cpp b/backends/xnnpack/runtime/XNNPACKBackend.cpp index b9b4e82f6ca..a76a0832def 100644 --- a/backends/xnnpack/runtime/XNNPACKBackend.cpp +++ b/backends/xnnpack/runtime/XNNPACKBackend.cpp @@ -255,6 +255,12 @@ class XnnpackBackend final return first_err; } + public: + /** See xnnpack::get_packed_cache_report(). */ + xnnpack::PackedCacheReport packed_cache_report() const { + return options_.weights_cache_manager().report(); + } + private: mutable xnnpack::XnnpackBackendOptions options_; @@ -272,5 +278,11 @@ Backend backend{xnnpack::xnnpack_backend_key, &backend_instance}; static auto success_with_compiler = register_backend(backend); } // namespace +namespace xnnpack { +PackedCacheReport get_packed_cache_report() { + return backend_instance.packed_cache_report(); +} +} // namespace xnnpack + } // namespace backends } // namespace executorch diff --git a/backends/xnnpack/runtime/XNNPACKBackend.h b/backends/xnnpack/runtime/XNNPACKBackend.h index 1053a206360..df83a500489 100644 --- a/backends/xnnpack/runtime/XNNPACKBackend.h +++ b/backends/xnnpack/runtime/XNNPACKBackend.h @@ -1,5 +1,11 @@ #pragma once +#include +#include +#include +#include +#include + namespace executorch::backends::xnnpack { /// The key for the backend. This is used to register the backend, check /// availability, and get/set options. @@ -61,4 +67,111 @@ enum class WorkspaceSharingMode { // maximum enum value. Count, }; + +/// Outcome of opening the packed-weight cache file. +enum class PackedCacheState : int32_t { + /// No cache path configured — the caller never opted in. + Disabled = 0, + /// The cache file opened. Does NOT imply zero heap; see PackedCacheStats. + FileBacked = 1, + /// A path was configured but the file could not be used. + HeapFallback = 2, +}; + +/** Why an individual allocation was served from heap. */ +enum class PackedCacheHeapReason : int32_t { + None = 0, + /// The instance has no cache path: it never opted into file backing, so + /// heap is the intended behaviour rather than a fallback. Bucketed + /// separately and excluded from heap_bytes — a process that mixes an + /// opted-in model with a non-opted-in one would otherwise report the + /// latter's packed weights as if the former had fallen back. + NotOptedIn = 1, + /// Unnamed constant — can never be reloaded by name. By design. + UnnamedConstant = 2, + /// Incidental re-pack after a successful load. By design *if* the loaded + /// cache is complete; a large volume here means it was not. + RepackAfterLoad = 3, + /// No usable file descriptor at allocation time. + NoFileBacking = 4, + /// ftruncate() to extend the file failed. + GrowFailed = 5, + /// mmap() of the grown region failed. + MmapFailed = 6, + /// Not a reason; bounds the per-reason counters. Matches the + /// WorkspaceSharingMode convention in this header. + Count, +}; + +/** Which step failed when a configured path still ended up on heap. */ +enum class PackedCacheFailure : int32_t { + None = 0, + OpenFailed = 1, + TruncateFailed = 2, + GrowFailed = 3, + MmapFailed = 4, +}; + +/** + * Per-cache counters. `heap_bytes` against `mapped_bytes` is the signal; + * `state` alone calls a partially-loaded cache healthy. + */ +struct PackedCacheStats { + PackedCacheState state{PackedCacheState::Disabled}; + PackedCacheFailure failure{PackedCacheFailure::None}; + int32_t last_errno{0}; + /// Cache file size as of the last successful save. + int64_t file_bytes{0}; + /// Packed bytes served from heap when the file was supposed to serve them. + /// Excludes NotOptedIn, so this is only ever "bytes that should have been + /// file-backed and were not". + int64_t heap_bytes{0}; + /// Packed bytes served from the mmap'd file (clean, file-backed). + int64_t mapped_bytes{0}; + /// Reason accounting for the largest share of heap_bytes. On the aggregate + /// this is the argmax over per-reason bytes summed across caches, not the + /// local reason of whichever cache happened to allocate the most. + PackedCacheHeapReason heap_reason{PackedCacheHeapReason::None}; + /// Heap bytes split by reason, so callers can sum per reason rather than + /// per cache. Index with PackedCacheHeapReason. Excludes nothing — the + /// NotOptedIn slot is populated here but omitted from `heap_bytes`. + std::array(PackedCacheHeapReason::Count)> + heap_bytes_by_reason{}; +}; + +/** One live cache instance and its own counters. */ +struct PackedCacheEntry { + /// Cache file path. Empty for the shared heap-only instance handed to + /// callers that never configured one. + std::string path; + PackedCacheStats stats; +}; + +/** + * Aggregate plus the per-instance breakdown behind it. + * + * Both come from one pass, so the summary and the detail always describe the + * same instant. The breakdown exists because the aggregate alone cannot be + * attributed: a process running several models folds them into one number, so + * a fallback in one model is indistinguishable from a fallback in another. + * The manager already keys caches by path — this stops discarding that. + * + * Takes no per-instance lock: the counters are atomics, so this never waits + * on a model compile. + */ +struct PackedCacheReport { + /// Summed counters. `failure` / `last_errno` are left unset here; read them + /// from the `dominant_fallback` entry so they stay tied to one cache. + PackedCacheStats aggregate; + /// Sorted by path, so repeated calls agree regardless of map iteration + /// order. + std::vector per_cache; + /// Index into `per_cache` of the cache that best explains a fallback: the + /// largest heap contributor, or if none allocated, the first cache in + /// HeapFallback. -1 when nothing fell back. + int32_t dominant_fallback{-1}; +}; + +PackedCacheReport get_packed_cache_report(); + } // namespace executorch::backends::xnnpack diff --git a/backends/xnnpack/runtime/XNNWeightsCache.cpp b/backends/xnnpack/runtime/XNNWeightsCache.cpp index 34479c1c369..3b6b6d310e9 100644 --- a/backends/xnnpack/runtime/XNNWeightsCache.cpp +++ b/backends/xnnpack/runtime/XNNWeightsCache.cpp @@ -80,14 +80,28 @@ static T read_le(const uint8_t* src) { // Open the cache file and take an advisory exclusive lock. Returns the // fd, or -1 if open/flock failed (logs the failure). The caller decides // how to recover (typically: skip the mmap path for this init). -static int open_locked(const std::string& path, int flags) { +// out_errno receives the errno of whichever call failed. Reading errno at the +// call site does not work: the flock path closes the fd first, and close() (or +// ET_LOG) can overwrite it. +static int open_locked(const std::string& path, int flags, int* out_errno) { + if (out_errno != nullptr) { + *out_errno = 0; + } int fd = open(path.c_str(), flags, 0600); if (fd < 0) { - ET_LOG(Error, "open(%s) failed (errno=%d)", path.c_str(), errno); + const int err = errno; + if (out_errno != nullptr) { + *out_errno = err; + } + ET_LOG(Error, "open(%s) failed (errno=%d)", path.c_str(), err); return -1; } if (flock(fd, LOCK_EX | LOCK_NB) != 0) { - ET_LOG(Error, "flock(%s) failed (errno=%d)", path.c_str(), errno); + const int err = errno; + if (out_errno != nullptr) { + *out_errno = err; + } + ET_LOG(Error, "flock(%s) failed (errno=%d)", path.c_str(), err); close(fd); return -1; } @@ -128,6 +142,72 @@ void XNNWeightsCache::reset_for_fresh_write() { } #endif +void XNNWeightsCache::record_cache_failure( + PackedCacheFailure failure, + int err) noexcept { + state_.store( + static_cast(PackedCacheState::HeapFallback), + std::memory_order_relaxed); + failure_.store(static_cast(failure), std::memory_order_relaxed); + last_errno_.store(err, std::memory_order_relaxed); +} + +PackedCacheStats XNNWeightsCache::stats() const noexcept { + PackedCacheStats out; + out.state = + static_cast(state_.load(std::memory_order_relaxed)); + out.failure = + static_cast(failure_.load(std::memory_order_relaxed)); + out.last_errno = last_errno_.load(std::memory_order_relaxed); + out.file_bytes = file_bytes_.load(std::memory_order_relaxed); + out.mapped_bytes = mapped_bytes_.load(std::memory_order_relaxed); + int64_t worst = 0; + for (size_t i = 0; i < heap_bytes_by_reason_.size(); ++i) { + const int64_t bytes = + heap_bytes_by_reason_[i].load(std::memory_order_relaxed); + out.heap_bytes_by_reason[i] = bytes; + if (i == static_cast(PackedCacheHeapReason::NotOptedIn)) { + continue; // intended heap use, not a fallback + } + out.heap_bytes += bytes; + if (bytes > worst) { + worst = bytes; + out.heap_reason = static_cast(i); + } + } + return out; +} + +void XNNWeightsCache::record_heap_alloc( + size_t n, + PackedCacheHeapReason reason) noexcept { + // Re-bucket every reason to NotOptedIn when no path was configured. Such an + // instance is the shared heap-only cache handed to callers that never asked + // for file backing; counting its bytes as a fallback would inflate the + // metric for whichever model in the process *did* opt in. + // packed_cache_path_ is set once before the instance is published and never + // mutated, so this read needs no synchronization. + const PackedCacheHeapReason bucket = + packed_cache_path_.empty() ? PackedCacheHeapReason::NotOptedIn : reason; + heap_bytes_by_reason_[static_cast(bucket)].fetch_add( + static_cast(n), std::memory_order_relaxed); +} + +void XNNWeightsCache::record_mapped_alloc(size_t n) noexcept { + mapped_bytes_.fetch_add(static_cast(n), std::memory_order_relaxed); +} + +void XNNWeightsCache::mark_cache_file_backed() noexcept { + // Only ever upgrades Disabled -> FileBacked. A fallback already recorded + // describes memory the process is carrying, so a later success must not + // mask it. + int32_t expected = static_cast(PackedCacheState::Disabled); + state_.compare_exchange_strong( + expected, + static_cast(PackedCacheState::FileBacked), + std::memory_order_relaxed); +} + Error XNNWeightsCache::initialize_for_runtime( MemoryAllocator* runtime_allocator, const NamedDataMap* named_data_map) { @@ -147,7 +227,13 @@ Error XNNWeightsCache::initialize_for_runtime( // where fresh-write→save→re-init re-enters load_packed_cache and // double-mmaps the same file. if (!name_to_packed_data_metadata_.empty()) { - packed_file_fd_ = open_locked(packed_cache_path_, O_RDWR); + int open_errno = 0; + packed_file_fd_ = open_locked(packed_cache_path_, O_RDWR, &open_errno); + if (packed_file_fd_ < 0) { + record_cache_failure(PackedCacheFailure::OpenFailed, open_errno); + } else { + mark_cache_file_backed(); + } return Error::Ok; } @@ -160,27 +246,47 @@ Error XNNWeightsCache::initialize_for_runtime( "Loaded packed weight cache: %s (%zu entries)", packed_cache_path_.c_str(), name_to_packed_data_metadata_.size()); - packed_file_fd_ = open_locked(packed_cache_path_, O_RDWR); + int open_errno = 0; + packed_file_fd_ = open_locked(packed_cache_path_, O_RDWR, &open_errno); + // Loaded entries are already mmap'd, so reads stay file-backed even if the + // write fd could not be reopened. Record the errno anyway: without it a + // partial cache silently re-packs to heap every launch with no reason. + if (packed_file_fd_ < 0) { + record_cache_failure(PackedCacheFailure::OpenFailed, open_errno); + } + mark_cache_file_backed(); return Error::Ok; } // Fresh write. Skip O_TRUNC in open_locked so a concurrent holder's // mmap stays valid; truncate explicitly only after we hold the lock. - packed_file_fd_ = open_locked(packed_cache_path_, O_RDWR | O_CREAT); + int create_errno = 0; + packed_file_fd_ = + open_locked(packed_cache_path_, O_RDWR | O_CREAT, &create_errno); if (packed_file_fd_ < 0) { + const int err = create_errno; + ET_LOG( + Error, + "open(O_RDWR|O_CREAT) failed for %s (errno=%d); heap fallback this init", + packed_cache_path_.c_str(), + err); + record_cache_failure(PackedCacheFailure::OpenFailed, err); return Error::Ok; } if (ftruncate(packed_file_fd_, 0) != 0) { + const int err = errno; ET_LOG( Error, "ftruncate(0) failed for %s (errno=%d); heap fallback this init", packed_cache_path_.c_str(), - errno); + err); + record_cache_failure(PackedCacheFailure::TruncateFailed, err); close(packed_file_fd_); packed_file_fd_ = -1; return Error::Ok; } reset_for_fresh_write(); + mark_cache_file_backed(); ET_LOG( Info, "Opened packed weight file for writing: %s", @@ -394,6 +500,10 @@ void* XNNWeightsCache::reserve_space(XNNWeightsCache* context, size_t n) { // instead of re-packing into heap (dirty memory) every time. if (context->last_lookup_unnamed_ || (context->loaded_from_disk_ && !seed_mismatch_repack)) { + context->record_heap_alloc( + n, + context->last_lookup_unnamed_ ? PackedCacheHeapReason::UnnamedConstant + : PackedCacheHeapReason::RepackAfterLoad); return context->reserve_space_heap(n); } if (context->packed_file_fd_ >= 0) { @@ -403,13 +513,16 @@ void* XNNWeightsCache::reserve_space(XNNWeightsCache* context, size_t n) { size_t map_size = (n + page_size - 1) & ~(page_size - 1); if (ftruncate(context->packed_file_fd_, file_offset + map_size) != 0) { + const int err = errno; ET_LOG( Error, "reserve_space ftruncate to %zu failed (errno=%d)", file_offset + map_size, - errno); + err); + context->record_cache_failure(PackedCacheFailure::GrowFailed, err); close(context->packed_file_fd_); context->packed_file_fd_ = -1; + context->record_heap_alloc(n, PackedCacheHeapReason::GrowFailed); return context->reserve_space_heap(n); } @@ -421,13 +534,16 @@ void* XNNWeightsCache::reserve_space(XNNWeightsCache* context, size_t n) { context->packed_file_fd_, file_offset); if (ptr == MAP_FAILED) { + const int err = errno; ET_LOG( Error, "reserve_space mmap %zu bytes failed (errno=%d)", map_size, - errno); + err); + context->record_cache_failure(PackedCacheFailure::MmapFailed, err); close(context->packed_file_fd_); context->packed_file_fd_ = -1; + context->record_heap_alloc(n, PackedCacheHeapReason::MmapFailed); return context->reserve_space_heap(n); } @@ -439,12 +555,16 @@ void* XNNWeightsCache::reserve_space(XNNWeightsCache* context, size_t n) { kPackedAllocationAlignment); context->packed_file_used_ = file_offset + map_size; + // n, not map_size: the heap side records the raw request too, and the + // heap:mapped ratio is only meaningful if both measure the same thing. + context->record_mapped_alloc(n); context->file_ptr_to_region_index_[ptr] = context->mmap_regions_.size(); context->mmap_regions_.push_back({ptr, map_size}); context->ptr_to_file_offset_[ptr] = file_offset; return ptr; } #endif + context->record_heap_alloc(n, PackedCacheHeapReason::NoFileBacking); return context->reserve_space_heap(n); } @@ -609,6 +729,8 @@ Error XNNWeightsCache::save_packed_index() { // trailer drops the old entry. Monitoring file_bytes over time tells // us when GC or a size cap is needed. const size_t file_bytes = index_start + buf.size(); + file_bytes_.store( + static_cast(file_bytes), std::memory_order_relaxed); ET_LOG( Info, "Saved packed weight index: %u entries at offset %zu, file_bytes=%zu", @@ -699,6 +821,9 @@ bool XNNWeightsCache::load_packed_cache() { } mmap_regions_.push_back({map, file_size}); + // Bytes actually referenced by the index. Less than file_size whenever an + // earlier run re-packed a name and orphaned its old bytes. + size_t loaded_bytes = 0; const uint8_t* cursor = static_cast(map) + index_start; const uint8_t* end = static_cast(map) + index_region_end; @@ -766,6 +891,7 @@ bool XNNWeightsCache::load_packed_cache() { meta.in_current_runtime = false; meta.from_load = true; meta.seed = seed; + loaded_bytes += static_cast(data_size); name_to_packed_data_metadata_[name] = meta; } @@ -783,6 +909,15 @@ bool XNNWeightsCache::load_packed_cache() { mmap_regions_at_last_save_ = mmap_regions_.size(); mmap_regions_synced_ = mmap_regions_.size(); loaded_from_disk_ = true; + // Success path only: the truncated-entry branch above munmaps and rolls + // back, so counting at the mmap call would over-report. + // + // loaded_bytes, not file_size. The file is append-only, so a same-name + // re-pack leaves the old bytes behind; file_size counts those orphans and + // would inflate mapped_bytes against heap_bytes. Without this a warm launch + // reports heap=0/mapped=0/file=0, identical to the feature being off. + record_mapped_alloc(loaded_bytes); + file_bytes_.store(static_cast(file_size), std::memory_order_relaxed); return true; #else return false; diff --git a/backends/xnnpack/runtime/XNNWeightsCache.h b/backends/xnnpack/runtime/XNNWeightsCache.h index f584199e307..8ac023f63c4 100644 --- a/backends/xnnpack/runtime/XNNWeightsCache.h +++ b/backends/xnnpack/runtime/XNNWeightsCache.h @@ -10,10 +10,13 @@ #include +#include #include #include #include #include +#include +#include #include #include #include @@ -52,6 +55,13 @@ struct PackedDataMeta { uint32_t seed{0}; }; +// Telemetry types live in XNNPACKBackend.h — hosts read them without pulling +// in xnnpack.h through this header. +using xnnpack::PackedCacheFailure; +using xnnpack::PackedCacheHeapReason; +using xnnpack::PackedCacheState; +using xnnpack::PackedCacheStats; + class XNNWeightsCache { public: XNNWeightsCache(); @@ -162,7 +172,50 @@ class XNNWeightsCache { return instance_mutex_; } + /** + * Outcome of the file-backed path for this instance. HeapFallback is + * sticky: once an init has been served from heap the instance keeps + * reporting it, because that is the memory the process is actually + * carrying for the rest of its life. + */ + PackedCacheStats stats() const noexcept; + private: + /** Record a fallback. Overwrites any previous failure for this instance. */ + void record_cache_failure(PackedCacheFailure failure, int err) noexcept; + /** Note a working file-backed path; never downgrades a recorded fallback. */ + void mark_cache_file_backed() noexcept; + /** Attribute `n` packed bytes to heap under `reason`. */ + void record_heap_alloc(size_t n, PackedCacheHeapReason reason) noexcept; + /** Attribute `n` packed bytes to the mmap'd file. */ + void record_mapped_alloc(size_t n) noexcept; + + // Telemetry counters. Written from the XNNPACK callbacks (which run under + // the caller-held instance mutex) and read by hosts through + // XNNWeightsCacheManager::aggregate_stats() with no lock at all — atomics, + // not the mutex, are what make that read safe. The mutex is held across the + // whole of xnn_create_runtime, so a telemetry read that waited on it could + // stall an inference thread for the length of a model compile. + // + // relaxed ordering throughout: these are independent accumulators, and a + // reader that observes one field slightly ahead of another still gets a + // usable picture. There is no invariant spanning them. + // + // Cumulative for the instance's lifetime — delete_packed_data and + // full_unload do not decrement. Decrementing would need a ptr -> reason map + // kept alive purely for telemetry, and hosts sample right after a load or a + // generate, before anything is released, so the two agree in practice. + // Read them as "bytes this cache ever packed", not current residency. + std::atomic state_{static_cast(PackedCacheState::Disabled)}; + std::atomic failure_{static_cast(PackedCacheFailure::None)}; + std::atomic last_errno_{0}; + std::atomic file_bytes_{0}; + std::atomic mapped_bytes_{0}; + std::array< + std::atomic, + static_cast(PackedCacheHeapReason::Count)> + heap_bytes_by_reason_{}; + static constexpr uint32_t kCacheMagic = 0x58505743; // "XPWC" // Bump when the on-disk layout (footer or per-entry record) changes. // v2: per-entry seed added — old v1 files don't carry seeds and would diff --git a/backends/xnnpack/runtime/XNNWeightsCacheManager.cpp b/backends/xnnpack/runtime/XNNWeightsCacheManager.cpp index 0f122aa8ab0..116f7d9f5a0 100644 --- a/backends/xnnpack/runtime/XNNWeightsCacheManager.cpp +++ b/backends/xnnpack/runtime/XNNWeightsCacheManager.cpp @@ -10,6 +10,7 @@ #include +#include #include #include @@ -77,6 +78,98 @@ Error XNNWeightsCacheManager::save_all() { return first_err; } +xnnpack::PackedCacheReport XNNWeightsCacheManager::report() const { + // Snapshot path + instance under the owning mutexes, then read the counters + // without XNNWeightsCache::mutex(). That mutex is held across the whole of + // xnn_create_runtime, so waiting on it here would let a telemetry read stall + // an inference thread for the length of a model compile. + std::vector< + std::pair>> + live; + { + std::scoped_lock lock(meta_mutex_); + live.reserve(caches_.size()); + for (const auto& entry : caches_) { + if (auto cache = entry.second.lock()) { + live.emplace_back(entry.first, std::move(cache)); + } + } + } + { + std::scoped_lock lock(empty_path_mutex_); + if (auto cache = empty_path_cache_.lock()) { + live.emplace_back(std::string{}, std::move(cache)); + } + } + // caches_ is unordered; sort so the report and the index into it are stable + // across calls. + std::sort(live.begin(), live.end(), [](const auto& a, const auto& b) { + return a.first < b.first; + }); + + xnnpack::PackedCacheReport out; + out.per_cache.reserve(live.size()); + for (const auto& [path, cache] : live) { + out.per_cache.push_back(xnnpack::PackedCacheEntry{path, cache->stats()}); + } + + int64_t best_heap = -1; + int32_t first_fallback = -1; + for (size_t i = 0; i < out.per_cache.size(); ++i) { + const auto& s = out.per_cache[i].stats; + auto& agg = out.aggregate; + agg.file_bytes += s.file_bytes; + agg.heap_bytes += s.heap_bytes; + agg.mapped_bytes += s.mapped_bytes; + for (size_t r = 0; r < s.heap_bytes_by_reason.size(); ++r) { + agg.heap_bytes_by_reason[r] += s.heap_bytes_by_reason[r]; + } + // A fallback anywhere is the reportable outcome: if any cache on this + // process went to heap, the process is carrying that memory. + if (s.state == delegate::PackedCacheState::HeapFallback) { + agg.state = s.state; + if (first_fallback < 0) { + first_fallback = static_cast(i); + } + } else if ( + s.state == delegate::PackedCacheState::FileBacked && + agg.state == delegate::PackedCacheState::Disabled) { + agg.state = s.state; + } + if (s.heap_bytes > best_heap) { + best_heap = s.heap_bytes; + if (s.heap_bytes > 0) { + out.dominant_fallback = static_cast(i); + } + } + } + // A cache can fail before it ever allocates, so fall back to the first + // cache in HeapFallback rather than reporting no explanation at all. + if (out.dominant_fallback < 0) { + out.dominant_fallback = first_fallback; + } + + // Argmax over per-reason totals summed across caches — not the local reason + // of whichever cache allocated the most, which can disagree with the global + // picture when one cache mixes reasons. + int64_t worst = 0; + for (size_t r = 0; r < out.aggregate.heap_bytes_by_reason.size(); ++r) { + if (r == static_cast(delegate::PackedCacheHeapReason::NotOptedIn)) { + continue; // intended heap use, not a fallback + } + if (out.aggregate.heap_bytes_by_reason[r] > worst) { + worst = out.aggregate.heap_bytes_by_reason[r]; + out.aggregate.heap_reason = + static_cast(r); + } + } + return out; +} + +delegate::PackedCacheStats XNNWeightsCacheManager::aggregate_stats() const { + return report().aggregate; +} + size_t XNNWeightsCacheManager::live_count() const { std::scoped_lock lock(meta_mutex_); size_t count = 0; diff --git a/backends/xnnpack/runtime/XNNWeightsCacheManager.h b/backends/xnnpack/runtime/XNNWeightsCacheManager.h index c35285b6337..39e5c4b711a 100644 --- a/backends/xnnpack/runtime/XNNWeightsCacheManager.h +++ b/backends/xnnpack/runtime/XNNWeightsCacheManager.h @@ -54,6 +54,22 @@ class XNNWeightsCacheManager { * expired weak_ptrs. */ runtime::Error save_all(); + /** + * Worst outcome across live caches, for host telemetry. A HeapFallback + * anywhere wins over a FileBacked elsewhere: if any cache on this process + * went to heap, the process is carrying that memory. `file_bytes` sums + * across live instances. + */ + delegate::PackedCacheStats aggregate_stats() const; + + /** + * Aggregate plus the per-instance breakdown, from a single pass so the two + * always agree. Callers that need to attribute a fallback to a specific + * model use the breakdown; the aggregate answers "is this process carrying + * heap memory at all". + */ + xnnpack::PackedCacheReport report() const; + /** Test-only: count of live (non-expired) entries. */ size_t live_count() const; diff --git a/backends/xnnpack/test/ops/test_multiply.py b/backends/xnnpack/test/ops/test_multiply.py index 3315200005d..a44b5e6b406 100644 --- a/backends/xnnpack/test/ops/test_multiply.py +++ b/backends/xnnpack/test/ops/test_multiply.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. @@ -29,6 +30,10 @@ def forward(self, x, y): z = torch.mul(x, y) * torch.functional.torch.mul(x, y) return z + class MulScalar(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.mul.Scalar(x, 0.5) + class MulRelu(torch.nn.Module): def forward(self, x, y): z = x * y @@ -58,6 +63,23 @@ def test_fp32_mul(self): inputs = (torch.randn((1, 3)), torch.randn((4, 3))) self._test_mul(inputs) + def test_fp32_mul_scalar(self): + ( + Tester(self.MulScalar(), (torch.randn(2, 3),)) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not( + [ + "executorch_exir_dialects_edge__ops_aten_mul_Tensor", + "executorch_exir_dialects_edge__ops_aten_mul_Scalar", + ] + ) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + def test_qs8_mul(self): inputs = (torch.randn(1, 1, 4, 4), torch.randn(1, 1, 4, 1)) ( diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index adf1c694b22..bfe2ce862c9 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -22,6 +22,7 @@ from executorch.backends.xnnpack.test.tester import Quantize, RunPasses, Tester from executorch.backends.xnnpack.utils.quant_utils import ( is_dequant, + is_dynamic_qdq, is_quant, is_tagged_as_implicit_q_dq, ) @@ -364,6 +365,222 @@ def test_dq_conv2d_channels_last_tagged_reshape_pass(self) -> None: .run_method_and_compare_outputs() ) + class EltwiseConv2dDynamicQuant(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 10, 3) + + def forward(self, x): + return self.conv(torch.sigmoid(x)) + + def test_dq_conv2d_eltwise_source_channels_last_tagged_reshape_pass(self) -> None: + # The conv's input is sigmoid -> q -> dq. Stepping past the q/dq pair leaves + # the sigmoid reading NHWC while its own output stays NCHW, which XNNPACK + # only rejects at runtime. + tester = ( + Tester(self.EltwiseConv2dDynamicQuant().eval(), (torch.randn(1, 3, 8, 8),)) + .quantize( + Quantize( + quantization_config=get_symmetric_quantization_config( + is_dynamic=True + ) + ) + ) + .export() + .to_edge() + .run_passes(self.PassStage) + ) + + artifact = tester.get_artifact(StageType.RUN_PASSES) + graph_module = artifact.exported_program().graph_module + sigmoid_nodes = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.edge.aten.sigmoid.default + ] + self.assertEqual(len(sigmoid_nodes), 1) + sigmoid = sigmoid_nodes[0] + + # The sigmoid keeps its NCHW input and the copy sits on its output instead. + self.assertEqual(sigmoid.args[0].op, "placeholder") + copies = [ + user + for user in sigmoid.users + if user.target == exir_ops.edge.aten._to_copy.default + and user.kwargs.get("memory_format") == torch.channels_last + ] + self.assertEqual(len(copies), 1) + + tester.run_method_and_compare_outputs() + + class SiLUStemSharedConv2dDynamicQuant(torch.nn.Module): + """A SiLU stem ahead of the first convolution, as detection backbones have. + + Two producers here have more than one consumer: the input activation feeds + both the sigmoid and the mul, and the SiLU output feeds both the quantized + convolution and the graph output. Both are reachable by the blanket + ``replace_all_uses_with`` in the dynamic-quant branch of ``input_to_nhwc``. + """ + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3, padding=1) + + def forward(self, x): + act = x * torch.sigmoid(x) + return self.conv(act), act + + def test_dq_conv2d_silu_stem_shared_channels_last_tagged_reshape_pass(self) -> None: + tester = ( + Tester( + self.SiLUStemSharedConv2dDynamicQuant().eval(), + (torch.randn(1, 3, 16, 16),), + ) + .quantize( + Quantize( + quantization_config=get_symmetric_quantization_config( + is_dynamic=True + ) + ) + ) + .export() + .to_edge() + .run_passes(self.PassStage) + ) + + graph_module = ( + tester.get_artifact(StageType.RUN_PASSES).exported_program().graph_module + ) + muls = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.edge.aten.mul.Tensor + ] + self.assertEqual(len(muls), 1) + silu = muls[0] + + # The walk must stop at the SiLU output rather than run on to the input + # activation, which would leave the mul and the sigmoid reading NHWC while + # their own outputs stay NCHW. + for arg in silu.all_input_nodes: + self.assertNotEqual(arg.target, exir_ops.edge.aten._to_copy.default) + + # The SiLU output is converted once, for the convolution only. The graph + # output keeps the unconverted node. + copies = [ + user + for user in silu.users + if user.target == exir_ops.edge.aten._to_copy.default + and user.kwargs.get("memory_format") == torch.channels_last + ] + self.assertEqual(len(copies), 1) + output_node = next( + node for node in graph_module.graph.nodes if node.op == "output" + ) + self.assertIn(silu, output_node.args[0]) + + tester.run_method_and_compare_outputs() + + class SiblingBranchConv2dDynamicQuant(torch.nn.Module): + """A producer feeding both a quantized conv and an ordinary op.""" + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3, padding=1) + + def forward(self, x): + act = torch.sigmoid(x) + return self.conv(act), torch.tanh(act) + + def test_dq_conv2d_sibling_branch_channels_last_tagged_reshape_pass(self) -> None: + tester = ( + Tester( + self.SiblingBranchConv2dDynamicQuant().eval(), + (torch.randn(1, 3, 16, 16),), + ) + .quantize( + Quantize( + quantization_config=get_symmetric_quantization_config( + is_dynamic=True + ) + ) + ) + .export() + .to_edge() + .run_passes(self.PassStage) + ) + + graph_module = ( + tester.get_artifact(StageType.RUN_PASSES).exported_program().graph_module + ) + tanhs = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.edge.aten.tanh.default + ] + self.assertEqual(len(tanhs), 1) + + # Only the quantize wrapper moved to the NHWC copy. + self.assertEqual(tanhs[0].args[0].target, exir_ops.edge.aten.sigmoid.default) + + tester.run_method_and_compare_outputs() + + class SharedLinearConvDynamicQuant(torch.nn.Module): + """Two dynamically quantized siblings sharing one source; only the conv + wants NHWC. + """ + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(8, 8) + self.conv = torch.nn.Conv2d(3, 4, 1) + + def forward(self, x): + act = torch.sigmoid(x) + # Keep linear before conv so it is processed first. + return self.linear(act), self.conv(act) + + def test_dq_shared_linear_conv_channels_last_tagged_reshape_pass(self) -> None: + tester = ( + Tester( + self.SharedLinearConvDynamicQuant().eval(), + (torch.randn(1, 3, 8, 8),), + ) + .quantize( + Quantize( + quantization_config=get_symmetric_quantization_config( + is_dynamic=True + ) + ) + ) + .export() + .to_edge() + .run_passes(self.PassStage) + ) + + graph_module = ( + tester.get_artifact(StageType.RUN_PASSES).exported_program().graph_module + ) + quantizes = [ + node + for node in graph_module.graph.nodes + if is_dynamic_qdq(node) and is_quant(node) + ] + self.assertEqual(len(quantizes), 2) + for quantize in quantizes: + consumer = next(iter(next(iter(quantize.users)).users)) + source = quantize.args[0].target + qparam_source = quantize.args[1].args[0].args[0].target + self.assertEqual(source, qparam_source) + if consumer.target == exir_ops.edge.aten.convolution.default: + # The conv's chain reads the NHWC copy. + self.assertEqual(source, exir_ops.edge.aten._to_copy.default) + else: + # The linear's chain keeps the source. + self.assertEqual(source, exir_ops.edge.aten.sigmoid.default) + + tester.run_method_and_compare_outputs() + class ConvAddConvOutput(torch.nn.Module): def __init__(self): super().__init__() diff --git a/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py new file mode 100644 index 00000000000..dd59aab5204 --- /dev/null +++ b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py @@ -0,0 +1,142 @@ +# 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.xnnpack._passes import XNNPACKPassManager +from executorch.backends.xnnpack._passes.convert_to_sdpa import ConvertToSDPAPass +from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import ( + LiftConstantScalarOperandsPass, +) +from executorch.backends.xnnpack.utils.configs import ( + get_transform_passes, + get_xnnpack_edge_compile_config, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export.graph_signature import InputKind + + +class TestLiftConstantScalarOperandsPass(unittest.TestCase): + def setUp(self): + torch._dynamo.reset() + + class MulScalar(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.mul.Scalar(x, 0.5) + + class AddScalar(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.add.Scalar(x, 0.5) + + class SDPA(torch.nn.Module): + def forward(self, q, k, v, mask): + return torch.nn.functional.scaled_dot_product_attention(q, k, v, mask) + + def _to_edge_program_manager(self, module): + return to_edge( + torch.export.export(module, (torch.randn(2, 3),), strict=True), + compile_config=get_xnnpack_edge_compile_config(skip_dim_order=True), + ) + + def _lift(self, exported_program): + return XNNPACKPassManager( + exported_program, passes=[LiftConstantScalarOperandsPass] + ).transform() + + def test_lifts_mul_scalar_operand(self): + exported_program = self._lift( + self._to_edge_program_manager(self.MulScalar()).exported_program() + ) + graph = exported_program.graph_module.graph + + self.assertFalse( + any(node.target == exir_ops.edge.aten.mul.Scalar for node in graph.nodes) + ) + self.assertTrue( + any(node.target == exir_ops.edge.aten.mul.Tensor for node in graph.nodes) + ) + self.assertFalse(any(node.op == "get_attr" for node in graph.nodes)) + + constant_specs = [ + spec + for spec in exported_program.graph_signature.input_specs + if spec.kind == InputKind.CONSTANT_TENSOR + ] + self.assertEqual(len(constant_specs), 1) + constant_spec = constant_specs[0] + self.assertIn(constant_spec.target, exported_program.constants) + + placeholders = [node for node in graph.nodes if node.op == "placeholder"] + self.assertEqual(placeholders[0].name, constant_spec.arg.name) + mul_node = next( + node for node in graph.nodes if node.target == exir_ops.edge.aten.mul.Tensor + ) + self.assertIs(mul_node.args[1], placeholders[0]) + + def test_is_idempotent(self): + exported_program = self._lift( + self._to_edge_program_manager(self.MulScalar()).exported_program() + ) + exported_program = self._lift(exported_program) + + constant_specs = [ + spec + for spec in exported_program.graph_signature.input_specs + if spec.kind == InputKind.CONSTANT_TENSOR + ] + self.assertEqual(len(constant_specs), 1) + self.assertEqual(len(exported_program.constants), 1) + + def test_keeps_unmapped_scalar_op(self): + exported_program = self._lift( + self._to_edge_program_manager(self.AddScalar()).exported_program() + ) + graph = exported_program.graph_module.graph + + self.assertTrue( + any(node.target == exir_ops.edge.aten.add.Scalar for node in graph.nodes) + ) + self.assertFalse(exported_program.constants) + + def test_converts_sdpa_after_default_transform_passes(self): + q = torch.randn(2, 4, 8, 16) + k = torch.randn(2, 4, 8, 16) + v = torch.randn(2, 4, 8, 16) + mask = torch.randn(8, 8) + for use_default_transforms in (False, True): + with self.subTest(use_default_transforms=use_default_transforms): + edge = to_edge( + torch.export.export(self.SDPA(), (q, k, v, mask), strict=True), + compile_config=get_xnnpack_edge_compile_config(), + ) + if use_default_transforms: + edge = edge.transform(get_transform_passes()) + exported_program = XNNPACKPassManager( + edge.exported_program(), + passes=[ConvertToSDPAPass, LiftConstantScalarOperandsPass], + ).transform() + + graph = exported_program.graph_module.graph + self.assertTrue( + any( + node.target + == exir_ops.edge.aten.scaled_dot_product_attention.default + for node in graph.nodes + ) + ) + self.assertFalse( + any( + node.target == exir_ops.edge.aten.bmm.default + for node in graph.nodes + ) + ) + self.assertFalse( + any( + node.target == exir_ops.edge.aten.mul.Scalar + for node in graph.nodes + ) + ) diff --git a/backends/xnnpack/test/passes/test_rewrite_fp16_silu.py b/backends/xnnpack/test/passes/test_rewrite_fp16_silu.py new file mode 100644 index 00000000000..131be95f846 --- /dev/null +++ b/backends/xnnpack/test/passes/test_rewrite_fp16_silu.py @@ -0,0 +1,104 @@ +# 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.xnnpack._passes import XNNPACKPassManager +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner +from executorch.backends.xnnpack.test.tester import Tester +from executorch.backends.xnnpack.utils.configs import get_xnnpack_edge_compile_config +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops + + +class TestRewriteFp16SiluPass(unittest.TestCase): + edge_mul = "executorch_exir_dialects_edge__ops_aten_mul_Tensor" + edge_sigmoid = "executorch_exir_dialects_edge__ops_aten_sigmoid_default" + edge_silu = "executorch_exir_dialects_edge__ops_aten_silu_default" + + class Silu(torch.nn.Module): + def forward(self, x): + return torch.nn.functional.silu(x) + + def setUp(self): + torch._dynamo.reset() + + def _export_silu(self, dtype): + return torch.export.export( + self.Silu(), + (torch.randn(2, 3, dtype=dtype),), + strict=True, + ) + + def _get_preserved_edge_fp16_silu(self): + edge_program = to_edge( + self._export_silu(torch.float16), + compile_config=get_xnnpack_edge_compile_config(), + ).exported_program() + graph = edge_program.graph_module.graph + input_node = next(node for node in graph.nodes if node.op == "placeholder") + output_node = next(node for node in graph.nodes if node.op == "output") + silu = output_node.args[0][0] + silu.target = exir_ops.edge.aten.silu.default + silu.args = (input_node,) + silu.kwargs = {} + graph.eliminate_dead_code() + graph.lint() + edge_program.graph_module.recompile() + return edge_program + + def test_ops_to_not_decompose_filters_for_fp16(self): + partitioner = XnnpackPartitioner() + + fp16_program = self._export_silu(torch.float16) + preserved_ops, filter_fn = partitioner.ops_to_not_decompose(fp16_program) + fp16_silu = next( + node + for node in fp16_program.graph.nodes + if node.target == torch.ops.aten.silu.default + ) + self.assertIn(torch.ops.aten.silu.default, preserved_ops) + self.assertIsNotNone(filter_fn) + self.assertTrue(filter_fn(fp16_silu)) + + fp32_program = self._export_silu(torch.float32) + _, filter_fn = partitioner.ops_to_not_decompose(fp32_program) + fp32_silu = next( + node + for node in fp32_program.graph.nodes + if node.target == torch.ops.aten.silu.default + ) + self.assertIsNotNone(filter_fn) + self.assertFalse(filter_fn(fp32_silu)) + + def test_preprocess_rewrites_preserved_fp16_silu(self): + result = XNNPACKPassManager(self._get_preserved_edge_fp16_silu()).transform() + targets = [node.target for node in result.graph.nodes] + + self.assertNotIn(exir_ops.edge.aten.silu.default, targets) + self.assertEqual(targets.count(exir_ops.edge.aten.sigmoid.default), 1) + self.assertEqual(targets.count(exir_ops.edge.aten.mul.Tensor), 1) + + def test_fp32_silu_uses_default_decomposition(self): + ( + Tester(self.Silu(), (torch.randn(2, 3),)) + .export() + .to_edge() + .check_count({self.edge_silu: 0, self.edge_sigmoid: 1, self.edge_mul: 1}) + ) + + def test_to_edge_transform_and_lower_delegates_fp16_silu(self): + ( + Tester(self.Silu(), (torch.randn(2, 3, dtype=torch.float16),)) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not([self.edge_silu]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) diff --git a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py index 1e1a473dd59..27d6a8f65fb 100644 --- a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py +++ b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py @@ -1122,6 +1122,30 @@ def forward(self, x): node_list, ) + def test_int64_scalar_add_used_as_index(self): + """Scalars lifted to attrs must keep the op's output dtype; an int64 + add chain used as an index must not be promoted to float32.""" + + class M(torch.nn.Module): + def forward(self, x): + return x[:, torch.arange(4) + 0] + + quantizer = XNNPACKQuantizer() + quantization_config = get_symmetric_quantization_config(is_per_channel=True) + quantizer.set_global(quantization_config) + example_inputs = (torch.randn(1, 4, 5),) + m = export(M(), example_inputs, strict=True).module() + m = quantizer.transform_for_annotation(m) + lifted_constants = [ + m.get_buffer(n.target) + for n in m.graph.nodes + if n.op == "get_attr" and n.target.startswith("_tensor_constant_") + ] + self.assertEqual(len(lifted_constants), 1) + self.assertEqual(lifted_constants[0].dtype, torch.int64) + m = prepare_pt2e(m, quantizer) + m(*example_inputs) + def test_cat_same_node(self): """Ensure that concatenating the same node does not cause any unexpected behavior""" diff --git a/backends/xnnpack/test/runtime/test_weight_cache.cpp b/backends/xnnpack/test/runtime/test_weight_cache.cpp index d2c079c057a..384edbaec8c 100644 --- a/backends/xnnpack/test/runtime/test_weight_cache.cpp +++ b/backends/xnnpack/test/runtime/test_weight_cache.cpp @@ -11,13 +11,19 @@ #include #include #include + #include #include #include #include +#include using namespace ::testing; +using executorch::backends::xnnpack::get_packed_cache_report; +using executorch::backends::xnnpack::packed_cache_path_option_key; +using executorch::backends::xnnpack::PackedCacheHeapReason; +using executorch::backends::xnnpack::save_weight_cache_on_disk_option_key; using executorch::backends::xnnpack::weight_cache_option_key; using executorch::backends::xnnpack::workspace_sharing_mode_option_key; using executorch::backends::xnnpack::WorkspaceSharingMode; @@ -148,3 +154,42 @@ TEST(RuntimeSpec, OverridesGlobalWeightCache) { get_option(xnnpack_backend_key, read_option); ASSERT_EQ(std::get(read_option.value), true); } + +TEST(PackedCacheStats, GlobalAccessorReachesTheBackendSingleton) { + executorch::runtime::runtime_init(); + + // The wiring under test is get_packed_cache_report() -> the registered + // backend instance -> XnnpackBackendOptions -> XNNWeightsCacheManager. + // Hosts call only this entry point, and nothing else in the suite exercises + // it. Absolute values depend on what else has run in the process, so this + // asserts reachability and invariants rather than specific counts. + const auto report = get_packed_cache_report(); + const auto& stats = report.aggregate; + + EXPECT_GE(stats.heap_bytes, 0); + EXPECT_GE(stats.mapped_bytes, 0); + EXPECT_GE(stats.file_bytes, 0); + EXPECT_LT( + static_cast(stats.heap_reason), + static_cast(PackedCacheHeapReason::Count)); + EXPECT_NE(stats.heap_reason, PackedCacheHeapReason::NotOptedIn) + << "NotOptedIn is excluded from heap_bytes and must never be reported"; + + // The breakdown must be able to attribute the aggregate: several models + // share a process, and a single folded number cannot say which one fell + // back. + for (const auto& entry : report.per_cache) { + EXPECT_GE(entry.stats.heap_bytes, 0); + EXPECT_GE(entry.stats.mapped_bytes, 0); + } +} + +// NOTE: the warm path — load_packed_cache() succeeding and its mapped bytes +// being counted — is deliberately NOT covered here. Producing a loadable +// cache file needs a model with packed weights and a populated index; the +// models wired into this target (ModuleAddLarge / ModuleSubLarge) are +// elementwise and pack nothing, so a cold run writes a zero-entry trailer +// that load_packed_cache correctly rejects. Covering it needs ModuleLinear +// plus its external .ptd and a NamedDataMap, as test_xnn_data_separation +// does. Until then the warm path is verified on device by the +// PackedWeights log line reporting non-zero mapped/cache_file. diff --git a/backends/xnnpack/test/runtime/test_xnn_weights_cache_manager.cpp b/backends/xnnpack/test/runtime/test_xnn_weights_cache_manager.cpp index 06bc74211ad..c2764776f6f 100644 --- a/backends/xnnpack/test/runtime/test_xnn_weights_cache_manager.cpp +++ b/backends/xnnpack/test/runtime/test_xnn_weights_cache_manager.cpp @@ -14,8 +14,11 @@ #include #include +#include #include +#include #include +#include #include #include @@ -31,14 +34,34 @@ class XNNWeightsCacheManagerTest : public ::testing::Test { manager_ = std::make_unique(); } + void TearDown() override { + for (const auto& path : temp_paths_) { + std::remove(path.c_str()); + } + } + + // Unique per test and per process. A leftover file from an earlier run + // flips initialize_for_runtime between the load and fresh-create branches, + // and two concurrent runs would race on the same path. + std::string TempPath(const char* tag) { + const auto* info = ::testing::UnitTest::GetInstance()->current_test_info(); + auto path = std::string(::testing::TempDir()) + "xnnwc_" + info->name() + + "_" + tag + "_" + std::to_string(static_cast(::getpid())) + + ".bin"; + std::remove(path.c_str()); + temp_paths_.push_back(path); + return path; + } + std::unique_ptr manager_; + std::vector temp_paths_; }; // --- Core dedup semantics --- TEST_F(XNNWeightsCacheManagerTest, SamePathReturnsSameInstance) { - auto a = manager_->get_or_create("/tmp/test_cache_same.bin"); - auto b = manager_->get_or_create("/tmp/test_cache_same.bin"); + auto a = manager_->get_or_create(TempPath("same")); + auto b = manager_->get_or_create(TempPath("same")); ASSERT_TRUE(a.ok()); ASSERT_TRUE(b.ok()); EXPECT_EQ(a.get().get(), b.get().get()) @@ -46,8 +69,8 @@ TEST_F(XNNWeightsCacheManagerTest, SamePathReturnsSameInstance) { } TEST_F(XNNWeightsCacheManagerTest, DifferentPathsReturnDifferentInstances) { - auto a = manager_->get_or_create("/tmp/test_cache_a.bin"); - auto b = manager_->get_or_create("/tmp/test_cache_b.bin"); + auto a = manager_->get_or_create(TempPath("a")); + auto b = manager_->get_or_create(TempPath("b")); ASSERT_TRUE(a.ok()); ASSERT_TRUE(b.ok()); EXPECT_NE(a.get().get(), b.get().get()) @@ -85,7 +108,7 @@ TEST_F(XNNWeightsCacheManagerTest, EmptyPathRecreatedAfterAllRefsDrop) { TEST_F(XNNWeightsCacheManagerTest, EmptyPathDoesNotShareWithMmapPath) { auto empty = manager_->get_or_create(""); - auto mmap = manager_->get_or_create("/tmp/test_cache_isolation.bin"); + auto mmap = manager_->get_or_create(TempPath("isolation")); ASSERT_TRUE(empty.ok()); ASSERT_TRUE(mmap.ok()); // Empty-path cache stays separate from any mmap-path cache — @@ -100,7 +123,7 @@ TEST_F(XNNWeightsCacheManagerTest, EmptyPathDoesNotShareWithMmapPath) { TEST_F(XNNWeightsCacheManagerTest, ExpiredEntryDoesNotLeak) { { - auto a = manager_->get_or_create("/tmp/test_cache_expire.bin"); + auto a = manager_->get_or_create(TempPath("expire")); ASSERT_TRUE(a.ok()); EXPECT_EQ(manager_->live_count(), 1u); } @@ -112,13 +135,13 @@ TEST_F(XNNWeightsCacheManagerTest, ExpiredEntryDoesNotLeak) { TEST_F(XNNWeightsCacheManagerTest, ExpiredEntryRecreatedOnNextCall) { void* first_addr = nullptr; { - auto a = manager_->get_or_create("/tmp/test_cache_recreate.bin"); + auto a = manager_->get_or_create(TempPath("recreate")); ASSERT_TRUE(a.ok()); first_addr = a.get().get(); } // Address re-use is allowed but not required; the only guarantee is // that we get a usable instance, not a dangling shared_ptr. - auto b = manager_->get_or_create("/tmp/test_cache_recreate.bin"); + auto b = manager_->get_or_create(TempPath("recreate")); ASSERT_TRUE(b.ok()); ASSERT_NE(b.get(), nullptr); // Live count should be 1 again — the stale entry was erased and @@ -136,15 +159,18 @@ TEST_F(XNNWeightsCacheManagerTest, ConcurrentSamePathSameInstance) { std::vector threads; threads.reserve(kThreads); std::atomic ready{0}; + // Resolve the path up front: TempPath() appends to temp_paths_, which is + // not safe to call from the racing threads. + const std::string race_path = TempPath("race"); for (int i = 0; i < kThreads; ++i) { - threads.emplace_back([this, &results, &ready, i] { + threads.emplace_back([this, &results, &ready, &race_path, i] { // Spin to maximize the chance of true concurrent entry into // get_or_create. ready.fetch_add(1, std::memory_order_acq_rel); while (ready.load(std::memory_order_acquire) < kThreads) { std::this_thread::yield(); } - auto r = manager_->get_or_create("/tmp/test_cache_race.bin"); + auto r = manager_->get_or_create(race_path); ASSERT_TRUE(r.ok()); results[i] = r.get(); }); @@ -169,10 +195,14 @@ TEST_F(XNNWeightsCacheManagerTest, ConcurrentDifferentPathsIndependent) { std::vector> results(kThreads); std::vector threads; threads.reserve(kThreads); + std::vector paths; + paths.reserve(kThreads); for (int i = 0; i < kThreads; ++i) { - threads.emplace_back([this, &results, i] { - std::string path = "/tmp/test_cache_diff_" + std::to_string(i) + ".bin"; - auto r = manager_->get_or_create(path); + paths.push_back(TempPath(("diff_" + std::to_string(i)).c_str())); + } + for (int i = 0; i < kThreads; ++i) { + threads.emplace_back([this, &results, &paths, i] { + auto r = manager_->get_or_create(paths[i]); ASSERT_TRUE(r.ok()); results[i] = r.get(); }); @@ -195,8 +225,8 @@ TEST_F(XNNWeightsCacheManagerTest, SaveAllNoLiveInstancesIsOk) { } TEST_F(XNNWeightsCacheManagerTest, SaveAllWalksLiveCaches) { - auto a = manager_->get_or_create("/tmp/test_cache_save_a.bin"); - auto b = manager_->get_or_create("/tmp/test_cache_save_b.bin"); + auto a = manager_->get_or_create(TempPath("save_a")); + auto b = manager_->get_or_create(TempPath("save_b")); ASSERT_TRUE(a.ok()); ASSERT_TRUE(b.ok()); EXPECT_EQ(manager_->live_count(), 2u); @@ -208,7 +238,7 @@ TEST_F(XNNWeightsCacheManagerTest, SaveAllWalksLiveCaches) { TEST_F(XNNWeightsCacheManagerTest, SaveAllSkipsExpiredEntries) { { - auto a = manager_->get_or_create("/tmp/test_cache_save_expired.bin"); + auto a = manager_->get_or_create(TempPath("save_expired")); ASSERT_TRUE(a.ok()); } // The entry's weak_ptr is now expired. save_all must not crash on @@ -220,7 +250,268 @@ TEST_F(XNNWeightsCacheManagerTest, SaveAllSkipsExpiredEntries) { // --- Path is set on the instance before publishing --- TEST_F(XNNWeightsCacheManagerTest, NonEmptyPathRegistersInMap) { - auto a = manager_->get_or_create("/tmp/test_cache_register.bin"); + auto a = manager_->get_or_create(TempPath("register")); ASSERT_TRUE(a.ok()); EXPECT_EQ(manager_->live_count(), 1u); } + +// --- Packed-cache telemetry (host-visible fallback reporting) --- + +TEST_F(XNNWeightsCacheManagerTest, StatsDisabledWhenNoCacheEverUsed) { + const auto stats = manager_->aggregate_stats(); + EXPECT_EQ( + stats.state, + executorch::backends::xnnpack::delegate::PackedCacheState::Disabled); + EXPECT_EQ(stats.last_errno, 0); + EXPECT_EQ(stats.file_bytes, 0); + EXPECT_EQ(stats.heap_bytes, 0); + EXPECT_EQ(stats.mapped_bytes, 0); +} + +TEST_F(XNNWeightsCacheManagerTest, StatsReportOpenFailureWithErrno) { + // A path whose parent directory does not exist: open(O_RDWR|O_CREAT) fails + // with ENOENT, which is the same branch a full disk takes with ENOSPC. + auto cache = manager_->get_or_create("/nonexistent_dir_xnnwc/cache.bin"); + ASSERT_TRUE(cache.ok()); + { + std::lock_guard lock(cache.get()->mutex()); + ASSERT_EQ(cache.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok) + << "a fallback must stay non-fatal"; + } + + const auto report = manager_->report(); + EXPECT_EQ( + report.aggregate.state, + executorch::backends::xnnpack::delegate::PackedCacheState::HeapFallback) + << "an unusable path must be reported as a heap fallback, not silently"; + + // failure/errno are deliberately absent from the aggregate: they belong to + // one cache. dominant_fallback names which one, even though this cache + // never allocated (open failed before any pack). + ASSERT_GE(report.dominant_fallback, 0); + const auto& dominant = + report.per_cache[static_cast(report.dominant_fallback)].stats; + EXPECT_EQ( + dominant.failure, + executorch::backends::xnnpack::delegate::PackedCacheFailure::OpenFailed); + EXPECT_NE(dominant.last_errno, 0) + << "errno is what distinguishes ENOSPC from a path problem"; + EXPECT_EQ( + report.aggregate.failure, + executorch::backends::xnnpack::delegate::PackedCacheFailure::None) + << "the aggregate must not adopt one cache's failure"; +} + +TEST_F(XNNWeightsCacheManagerTest, HeapReasonIsGlobalArgmaxNotPerCache) { + // Two caches whose local dominant reasons disagree with the global one. + // Cache A: a failed grow plus a smaller unnamed pack. Cache B: unnamed only, + // larger in total than A's grow contribution. Summing per cache would report + // A's reason; summing per reason reports UnnamedConstant, which is correct. + auto a = manager_->get_or_create(TempPath("argmax_a")); + auto b = manager_->get_or_create(TempPath("argmax_b")); + ASSERT_TRUE(a.ok()); + ASSERT_TRUE(b.ok()); + + const auto unnamed_pack = [](XNNWeightsCache* c, size_t n) { + auto* provider = c->get(); + int dummy = 0; + xnn_weights_cache_look_up_key key{}; + key.kernel = &dummy; + key.bias = nullptr; + provider->look_up(provider->context, &key); + ASSERT_NE(provider->reserve_space(provider->context, n), nullptr); + }; + + { + std::lock_guard lock(a.get()->mutex()); + ASSERT_EQ(a.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + unnamed_pack(a.get().get(), 8192); + } + { + std::lock_guard lock(b.get()->mutex()); + ASSERT_EQ(b.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + unnamed_pack(b.get().get(), 16384); + unnamed_pack(b.get().get(), 16384); + } + + const auto report = manager_->report(); + EXPECT_EQ( + report.aggregate.heap_reason, + executorch::backends::xnnpack::delegate::PackedCacheHeapReason:: + UnnamedConstant); + const auto unnamed_idx = + static_cast(executorch::backends::xnnpack::delegate:: + PackedCacheHeapReason::UnnamedConstant); + EXPECT_EQ( + report.aggregate.heap_bytes_by_reason[unnamed_idx], + report.aggregate.heap_bytes) + << "per-reason totals must sum to the same heap_bytes"; +} + +TEST_F(XNNWeightsCacheManagerTest, PerCacheIsSortedByPathForStableIndices) { + // dominant_fallback is an index into per_cache, so the order must not + // depend on unordered_map iteration. + auto z = manager_->get_or_create(TempPath("zzz")); + auto a = manager_->get_or_create(TempPath("aaa")); + ASSERT_TRUE(z.ok()); + ASSERT_TRUE(a.ok()); + + const auto report = manager_->report(); + ASSERT_GE(report.per_cache.size(), 2u); + for (size_t i = 1; i < report.per_cache.size(); ++i) { + EXPECT_LE(report.per_cache[i - 1].path, report.per_cache[i].path); + } +} + +TEST_F(XNNWeightsCacheManagerTest, HeapFallbackWinsOverFileBackedInAggregate) { + auto bad = manager_->get_or_create("/nonexistent_dir_xnnwc/cache.bin"); + auto good = manager_->get_or_create(TempPath("stats_ok")); + ASSERT_TRUE(bad.ok()); + ASSERT_TRUE(good.ok()); + { + std::lock_guard lock(good.get()->mutex()); + ASSERT_EQ(good.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + } + { + std::lock_guard lock(bad.get()->mutex()); + ASSERT_EQ(bad.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + } + + EXPECT_EQ( + manager_->aggregate_stats().state, + executorch::backends::xnnpack::delegate::PackedCacheState::HeapFallback) + << "if any live cache fell back, the process is carrying that memory"; +} + +// A binary "did the file open" flag is not enough: a cache can load +// successfully and still serve most of its packed bytes from heap. These +// cover the byte accounting that distinguishes the two. + +TEST_F(XNNWeightsCacheManagerTest, MappedBytesCountedWhenFileBacked) { + auto cache = manager_->get_or_create(TempPath("bytes_mapped")); + ASSERT_TRUE(cache.ok()); + { + std::lock_guard lock(cache.get()->mutex()); + ASSERT_EQ(cache.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + auto* provider = cache.get()->get(); + ASSERT_NE(provider->reserve_space(provider->context, 4096), nullptr); + } + + const auto stats = manager_->aggregate_stats(); + EXPECT_GT(stats.mapped_bytes, 0); + EXPECT_EQ(stats.heap_bytes, 0) << "the healthy case must report zero heap"; +} + +TEST_F(XNNWeightsCacheManagerTest, HeapBytesAttributedToUnnamedConstant) { + auto cache = manager_->get_or_create(TempPath("bytes_unnamed")); + ASSERT_TRUE(cache.ok()); + { + std::lock_guard lock(cache.get()->mutex()); + ASSERT_EQ(cache.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + auto* provider = cache.get()->get(); + // A look_up whose kernel pointer was never named marks the next + // reserve_space as an unnamed constant, which routes to heap. + int dummy = 0; + xnn_weights_cache_look_up_key key{}; + key.kernel = &dummy; + key.bias = nullptr; + provider->look_up(provider->context, &key); + ASSERT_NE(provider->reserve_space(provider->context, 4096), nullptr); + } + + const auto stats = manager_->aggregate_stats(); + EXPECT_GT(stats.heap_bytes, 0) << "heap bytes must be counted, not hidden"; + EXPECT_EQ( + stats.heap_reason, + executorch::backends::xnnpack::delegate::PackedCacheHeapReason:: + UnnamedConstant); +} + +TEST_F(XNNWeightsCacheManagerTest, FileBackedStateDoesNotImplyZeroHeap) { + // The case a state flag alone reports as healthy: the file opened fine, so + // state is FileBacked, yet packed bytes still went to heap. Only the byte + // split makes that visible. + auto cache = manager_->get_or_create(TempPath("bytes_split")); + ASSERT_TRUE(cache.ok()); + { + std::lock_guard lock(cache.get()->mutex()); + ASSERT_EQ(cache.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + auto* provider = cache.get()->get(); + int dummy = 0; + xnn_weights_cache_look_up_key key{}; + key.kernel = &dummy; + key.bias = nullptr; + provider->look_up(provider->context, &key); + ASSERT_NE(provider->reserve_space(provider->context, 4096), nullptr); + } + + const auto stats = manager_->aggregate_stats(); + EXPECT_EQ( + stats.state, + executorch::backends::xnnpack::delegate::PackedCacheState::FileBacked) + << "the file opened, so state alone looks healthy"; + EXPECT_GT(stats.heap_bytes, 0) + << "but heap bytes are non-zero — this is what state alone hides"; +} + +TEST_F(XNNWeightsCacheManagerTest, AggregateStatsTakesNoInstanceLock) { + // aggregate_stats() must not wait on XNNWeightsCache::mutex(): that mutex is + // held across all of xnn_create_runtime, so a telemetry read that blocked on + // it would stall inference for the length of a model compile. + auto cache = manager_->get_or_create(TempPath("nolock")); + ASSERT_TRUE(cache.ok()); + std::lock_guard held(cache.get()->mutex()); + const auto stats = manager_->aggregate_stats(); // must not deadlock + EXPECT_EQ(stats.heap_bytes, 0); +} + +TEST_F(XNNWeightsCacheManagerTest, EmptyPathHeapIsNotCountedAsFallback) { + // The shared heap-only instance handed to callers that never configured a + // path. Its heap use is intended, so it must not inflate heap_bytes for a + // model in the same process that did opt in. + auto opted_out = manager_->get_or_create(""); + ASSERT_TRUE(opted_out.ok()); + { + std::lock_guard lock(opted_out.get()->mutex()); + ASSERT_EQ( + opted_out.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + auto* provider = opted_out.get()->get(); + int dummy = 0; + xnn_weights_cache_look_up_key key{}; + key.kernel = &dummy; + key.bias = nullptr; + provider->look_up(provider->context, &key); + ASSERT_NE(provider->reserve_space(provider->context, 4096), nullptr); + } + + const auto stats = manager_->aggregate_stats(); + EXPECT_EQ(stats.heap_bytes, 0) + << "a cache that never opted into file backing is not a fallback"; +} + +TEST_F(XNNWeightsCacheManagerTest, OptedInHeapStillCountedAlongsideOptedOut) { + // Both kinds live at once: only the opted-in instance's heap bytes count. + auto opted_out = manager_->get_or_create(""); + auto opted_in = manager_->get_or_create(TempPath("mixed")); + ASSERT_TRUE(opted_out.ok()); + ASSERT_TRUE(opted_in.ok()); + for (auto* cache : {opted_out.get().get(), opted_in.get().get()}) { + std::lock_guard lock(cache->mutex()); + ASSERT_EQ(cache->initialize_for_runtime(nullptr, nullptr), Error::Ok); + auto* provider = cache->get(); + int dummy = 0; + xnn_weights_cache_look_up_key key{}; + key.kernel = &dummy; + key.bias = nullptr; + provider->look_up(provider->context, &key); + ASSERT_NE(provider->reserve_space(provider->context, 8192), nullptr); + } + + const auto stats = manager_->aggregate_stats(); + EXPECT_GT(stats.heap_bytes, 0) << "the opted-in instance's heap must count"; + EXPECT_EQ( + stats.heap_reason, + executorch::backends::xnnpack::delegate::PackedCacheHeapReason:: + UnnamedConstant) + << "and NotOptedIn must never win the argmax"; +} diff --git a/backends/xnnpack/test/test_xnnpack_partitioner.py b/backends/xnnpack/test/test_xnnpack_partitioner.py index 894fab4098f..ff25c69655e 100644 --- a/backends/xnnpack/test/test_xnnpack_partitioner.py +++ b/backends/xnnpack/test/test_xnnpack_partitioner.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. @@ -30,6 +31,71 @@ def __init__(self): def forward(self, x): return self.linear(x) + class MulScalar(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.mul.Scalar(x, 0.5) + + class MulScalarCond(torch.nn.Module): + def forward(self, pred, x): + def true_fn(value): + return torch.ops.aten.mul.Scalar(value, 0.5) + + def false_fn(value): + return torch.ops.aten.mul.Scalar(value, 2.0) + + return torch.cond(pred, true_fn, false_fn, (x,)) + + def test_mul_scalar_ops_to_not_decompose(self): + partitioner = XnnpackPartitioner() + exported_program = export(self.MulScalar(), (torch.randn(2, 3),)) + ops, filter_fn = partitioner.ops_to_not_decompose(exported_program) + + self.assertIn(torch.ops.aten.mul.Scalar, ops) + self.assertIsNotNone(filter_fn) + mul_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.mul.Scalar + ) + self.assertTrue(filter_fn(mul_node)) + + def test_mul_scalar_ops_to_not_decompose_rejects_unsupported_dtype(self): + partitioner = XnnpackPartitioner() + exported_program = export( + self.MulScalar(), (torch.ones(2, 3, dtype=torch.int32),) + ) + _, filter_fn = partitioner.ops_to_not_decompose(exported_program) + + self.assertIsNotNone(filter_fn) + mul_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.mul.Scalar + ) + self.assertFalse(filter_fn(mul_node)) + + def test_mul_scalar_ops_to_not_decompose_rejects_cond_branches(self): + partitioner = XnnpackPartitioner() + exported_program = export( + self.MulScalarCond(), (torch.tensor(True), torch.randn(2, 3)) + ) + _, filter_fn = partitioner.ops_to_not_decompose(exported_program) + + self.assertIsNotNone(filter_fn) + cond_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.higher_order.cond + ) + for branch_node in cond_node.args[1:3]: + branch = exported_program.graph_module.get_submodule(branch_node.target) + mul_node = next( + node + for node in branch.graph.nodes + if node.target == torch.ops.aten.mul.Scalar + ) + self.assertFalse(filter_fn(mul_node)) + 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/docs/source/backends-cadence.md b/docs/source/backends-cadence.md index c5a5fc8497a..ddd258cdca5 100644 --- a/docs/source/backends-cadence.md +++ b/docs/source/backends-cadence.md @@ -311,7 +311,7 @@ cmake -DCMAKE_BUILD_TYPE=Debug \ -Bcmake-out/examples/cadence \ examples/cadence -cmake --build cmake-out/examples/cadence -j8 -t cadence_executorch_example +cmake --build cmake-out/examples/cadence -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) -t cadence_executorch_example ``` After having succesfully run the above step you should see two binary files in their CMake output directory. diff --git a/docs/source/backends-qualcomm.md b/docs/source/backends-qualcomm.md index caf2426ed16..cf9369cb214 100644 --- a/docs/source/backends-qualcomm.md +++ b/docs/source/backends-qualcomm.md @@ -36,15 +36,24 @@ Currently, this ExecuTorch Backend can delegate AI computations to Hexagon proce ### Host OS -The QNN Backend is currently verified on the following Linux host operating systems: +The QNN Backend is verified on the following host operating systems: - **Ubuntu 22.04 LTS (x64)** - **CentOS Stream 9** +- **Windows 10 / 11 (x64)** +- **Windows 10 / 11 (ARM64)** with Qualcomm NPU - **Windows Subsystem for Linux (WSL)** with Ubuntu 22.04 In general, we verify the backend on the same OS versions that the QNN SDK is officially validated against. The exact supported versions are documented in the QNN SDK. +#### Windows (x64 / ARM64) Setup + +To build on native Windows platforms, the MSVC toolchain must be installed. +The required MSVC Build Tools can be installed through **Visual Studio Installer**. + +For installation instructions, refer to the official [Microsoft Visual Studio Downloads page](https://visualstudio.microsoft.com/downloads/). + #### Windows (WSL) Setup To install Ubuntu 22.04 on WSL, run the following command in PowerShell or Windows Terminal: @@ -55,24 +64,23 @@ wsl --install -d ubuntu 22.04 This command will install WSL and set up Ubuntu 22.04 as the default Linux distribution. -For more details and troubleshooting, refer to the official Microsoft WSL installation guide: -👉 [Install WSL | Microsoft Learn](https://learn.microsoft.com/en-us/windows/wsl/install) +For more details and troubleshooting, refer to the official Microsoft WSL installation guide: [Install WSL | Microsoft Learn](https://learn.microsoft.com/en-us/windows/wsl/install). ### Hardware: -You will need an Android / Linux device with adb-connected running on one of Qualcomm SoCs listed in `QcomChipset`. Please navigate to [qc_schema.py](https://github.com/pytorch/executorch/blob/main/backends/qualcomm/serialization/qc_schema.py). -This example is verified with SM8550 and SM8450. +The QNN backend runs on Qualcomm SoCs (Systems on Chips) across two device families: + +- **Android / Linux devices** — connected over `adb`. This example is verified with SM8550 and SM8450. +- **Windows on ARM64 (WoA) devices** — This example is verified with SC8380XP (Qualcomm Snapdragon X Elite). + +The target SoC must be one of those listed in the `QcomChipset` enum; see [qc_schema.py](https://github.com/pytorch/executorch/blob/main/backends/qualcomm/serialization/qc_schema.py). ### Software: - - Follow ExecuTorch recommended Python version. - - A compiler to compile AOT parts, e.g., the GCC compiler comes with Ubuntu LTS. g++ version need to be 13 or higher. - - [Android NDK](https://developer.android.com/ndk). This example is verified with NDK 26c. - - (Optional) Target toolchain for linux embedded platform. - - [Qualcomm AI Engine Direct SDK](https://developer.qualcomm.com/software/qualcomm-ai-engine-direct-sdk) - - Click the "Get Software" button to download the latest version of the QNN SDK. - - Although newer versions are available, we have verified and recommend using QNN 2.37.0 for stability. - - You can download it directly from the following link: [QNN 2.37.0](https://softwarecenter.qualcomm.com/api/download/software/sdks/Qualcomm_AI_Runtime_Community/All/2.37.0.250724/v2.37.0.250724.zip) +[Qualcomm AI Engine Direct SDK](https://developer.qualcomm.com/software/qualcomm-ai-engine-direct-sdk) + - Click the "Get Software" button to download the latest version of the QNN SDK. + - Although newer versions are available, we have verified and recommend using QNN 2.37.0 for stability. + - You can download it directly from the following link: [QNN 2.37.0](https://softwarecenter.qualcomm.com/api/download/software/sdks/Qualcomm_AI_Runtime_Community/All/2.37.0.250724/v2.37.0.250724.zip) The directory with installed Qualcomm AI Engine Direct SDK looks like: ``` @@ -94,6 +102,18 @@ The directory with installed Qualcomm AI Engine Direct SDK looks like: └── share ``` +On Android / Linux devices: + + - Follow ExecuTorch recommended Python version. + - A compiler to compile AOT parts, e.g., the GCC compiler comes with Ubuntu LTS. g++ version need to be 13 or higher. + - [Android NDK](https://developer.android.com/ndk). This example is verified with NDK 26c. + - (Optional) Target toolchain for linux embedded platform. + +On Windows on ARM64 (WoA) devices: + + - Install the **AMD64 version of Python** to run AOT compilation under x64 emulation. This is required because certain Python modules used in the AOT workflow do not currently provide ARM64 prebuilt wheels. + - MSVC Build Tools. + ## Setting up your developer environment @@ -106,9 +126,9 @@ i.e., the directory containing `QNN_README.txt`. `$EXECUTORCH_ROOT` refers to the root of executorch git repository. -### Setup environment variables +### Setup QNN SDK paths and environment variables -Source the QNN SDK environment setup script to configure paths and environment variables: +For Linux platform: ```bash source $QNN_SDK_ROOT/bin/envsetup.sh @@ -116,35 +136,66 @@ source $QNN_SDK_ROOT/bin/envsetup.sh This sets up `LD_LIBRARY_PATH` and other required variables for the QNN SDK tools and libraries. -Additionally, set `PYTHONPATH` for ExecuTorch Python APIs: +For Windows platform: + +```powershell +& "$env:QNN_SDK_ROOT\bin\envsetup.ps1" +``` + +### Setup `PYTHONPATH` for ExecuTorch Python APIs + +For Linux platform: ```bash export PYTHONPATH=$EXECUTORCH_ROOT/..:$PYTHONPATH ``` +For Windows platform: + +```powershell +$env:PYTHONPATH="$env:EXECUTORCH_ROOT\..;$env:PYTHONPATH" +``` + ## Build -An example script for the below building instructions is [here](https://github.com/pytorch/executorch/blob/main/backends/qualcomm/scripts/build.sh). +**On Linux platform**, an example script for the below building instructions is [`build.sh`](https://github.com/pytorch/executorch/blob/main/backends/qualcomm/scripts/build.sh). We recommend to use the script because the ExecuTorch build-command can change from time to time. The above script is actively used. It is updated more frequently than this tutorial. An example usage is ```bash cd $EXECUTORCH_ROOT -# android target +# Android target ./backends/qualcomm/scripts/build.sh -# (optional) linux embedded target +# (Optional) Linux embedded target ./backends/qualcomm/scripts/build.sh --enable_linux_embedded -# for release build +# Android target for release build ./backends/qualcomm/scripts/build.sh --release ``` +**On Windows platform**, use the PowerShell script [`build.ps1`](https://github.com/pytorch/executorch/blob/main/backends/qualcomm/scripts/build.ps1) for the building instructions. Both Windows x64 and ARM64 architectures are supported. +Here's the example usage +```powershell +cd $env:EXECUTORCH_ROOT +# Generate both Windows x64 and ARM64 target libraries +.\backends\qualcomm\scripts\build.ps1 -Release +# Generate only Windows x64 target libraries +.\backends\qualcomm\scripts\build.ps1 -SkipArm64Windows -Release +# Generate only Windows ARM64 target libraries +.\backends\qualcomm\scripts\build.ps1 -SkipX86Windows -Release +``` + +> **Notes** +> +> The script supports building both x64 and cross-compiling ARM64 target artifacts on Windows x64 host. After the build completes, the ARM64 libraries and executables can be copied to a Windows on Snapdragon (WoS) device using `scp`. +> This allows a `.pte` generated on Windows x64 host to be executed on WoS device. ## Deploying and running on device ### AOT compile a model Refer to [this script](https://github.com/pytorch/executorch/blob/main/examples/qualcomm/scripts/deeplab_v3.py) for the exact flow. -We use deeplab-v3-resnet101 as an example in this tutorial. Run below commands to compile: +We use deeplab-v3-resnet101 as an example in this tutorial. +Run below commands to compile on Linux platform: ```bash cd $EXECUTORCH_ROOT @@ -152,14 +203,26 @@ cd $EXECUTORCH_ROOT python -m examples.qualcomm.scripts.deeplab_v3 --build_folder build-android --soc_model SM8550 --compile_only --download ``` -You might see something like below: +For Windows x64 and ARM64 platforms, run the following commands for the AOT compilation: +```powershell +cd $env:EXECUTORCH_ROOT +python -m examples.qualcomm.scripts.deeplab_v3 --build_folder build-x86_64-windows --soc_model SC8380XP --compile_only --download ``` -[INFO][Qnn ExecuTorch] Destroy Qnn context -[INFO][Qnn ExecuTorch] Destroy Qnn device -[INFO][Qnn ExecuTorch] Destroy Qnn backend -Finish compile_only and save to ./deeplab_v3/dlv3_qnn.pte +> **Notes** +> +> AOT compilation on Windows on ARM64 (WoA) device currently relies on an AMD64 Python environment running under x64 emulation, since some AOT dependencies are not yet distributed as ARM64 prebuilt wheels. + +You might see something like below: + +``` +Completed stage: Finalizing Graph Sequence (8966 us) +Starting stage: Completion +Completed stage: Completion (1388 us) +[INFO] [Qnn ExecuTorch]: Destroy Qnn context +[INFO] [Qnn ExecuTorch]: Destroy Qnn device +[INFO] [Qnn ExecuTorch]: Destroy Qnn backend ``` The compiled model is `./deeplab_v3/dlv3_qnn.pte`. @@ -167,26 +230,9 @@ The compiled model is `./deeplab_v3/dlv3_qnn.pte`. Note that the model is compiled for specific backend (e.g., HTP), so you can specify the target backend via `--backend gpu` or `--backend lpai`. If not specified, it will be default to HTP. -### Test model inference on QNN HTP emulator / QNN LPAI emulator - -We can test model inferences before deploying it to a device by HTP emulator. - -Let's build `qnn_executor_runner` for a x64 host: -```bash -# assuming the AOT component is built. -cd $EXECUTORCH_ROOT/build-x86 -cmake ../examples/qualcomm \ - -DCMAKE_PREFIX_PATH="$PWD/lib/cmake/ExecuTorch;$PWD/third-party/gflags;" \ - -DCMAKE_FIND_ROOT_PATH_MODE_PACKAGE=BOTH \ - -DPYTHON_EXECUTABLE=python3 \ - -Bexamples/qualcomm - -cmake --build examples/qualcomm -j$(nproc) +### Test model inference on Linux x64 host with QNN HTP emulator / QNN LPAI emulator -# qnn_executor_runner can be found under examples/qualcomm/executor_runner -# The full path is $EXECUTORCH_ROOT/build-x86/examples/qualcomm/executor_runner/qnn_executor_runner -ls examples/qualcomm/executor_runner -``` +Before deploying a model to a physical device, inference execution can be tested and validated on a Linux x64 host using the HTP / LPAI emulator. To run the HTP emulator / LPAI emulator, the dynamic linker needs to access QNN libraries and `libqnn_executorch_backend.so`. We set the below two paths to `LD_LIBRARY_PATH` environment variable: @@ -199,27 +245,55 @@ The second path is for `libqnn_executorch_backend.so`. So, we can run `./deeplab_v3/dlv3_qnn.pte` by: ```bash -cd $EXECUTORCH_ROOT/build-x86 +cd $EXECUTORCH_ROOT export LD_LIBRARY_PATH=$EXECUTORCH_ROOT/build-x86/lib/:$LD_LIBRARY_PATH -examples/qualcomm/executor_runner/qnn_executor_runner --model_path ../deeplab_v3/dlv3_qnn.pte +build-x86/examples/qualcomm/executor_runner/qnn_executor_runner --model_path ./deeplab_v3/dlv3_qnn.pte ``` We should see some outputs like the below. Note that the emulator can take some time to finish. ```bash -I 00:00:00.354662 executorch:qnn_executor_runner.cpp:213] Method loaded. -I 00:00:00.356460 executorch:qnn_executor_runner.cpp:261] ignoring error from set_output_data_ptr(): 0x2 -I 00:00:00.357991 executorch:qnn_executor_runner.cpp:261] ignoring error from set_output_data_ptr(): 0x2 -I 00:00:00.357996 executorch:qnn_executor_runner.cpp:265] Inputs prepared. - -I 00:01:09.328144 executorch:qnn_executor_runner.cpp:414] Model executed successfully. -I 00:01:09.328159 executorch:qnn_executor_runner.cpp:421] Write etdump to etdump.etdp, Size = 424 -[INFO] [Qnn ExecuTorch]: Destroy Qnn backend parameters +I 00:00:00.174364 executorch:qnn_executor_runner.cpp:416] Method loaded. +E 00:00:00.179250 executorch:method.cpp:1373] Output 0 is memory planned, or is a constant. Cannot override the existing data pointer. +I 00:00:00.179264 executorch:qnn_executor_runner.cpp:473] ignoring error from set_output_data_ptr(): 0x2 +E 00:00:00.183296 executorch:method.cpp:1373] Output 1 is memory planned, or is a constant. Cannot override the existing data pointer. +I 00:00:00.183305 executorch:qnn_executor_runner.cpp:473] ignoring error from set_output_data_ptr(): 0x2 +I 00:00:00.183310 executorch:qnn_executor_runner.cpp:479] Inputs prepared. +I 00:00:00.184008 executorch:qnn_executor_runner.cpp:684] Input list not provided. Inputs prepared with default values set. +I 00:01:19.663283 executorch:qnn_executor_runner.cpp:695] Model executed successfully. +I 00:01:19.663299 executorch:qnn_executor_runner.cpp:698] Perform 0 inferences for warming up +I 00:01:53.881349 executorch:qnn_executor_runner.cpp:715] 1 inferences took 34218.046000 ms, avg 34218.046000 ms +I 00:01:53.881426 executorch:qnn_executor_runner.cpp:727] Write etdump to etdump.etdp, Size = 576 [INFO] [Qnn ExecuTorch]: Destroy Qnn context [INFO] [Qnn ExecuTorch]: Destroy Qnn device [INFO] [Qnn ExecuTorch]: Destroy Qnn backend ``` -### Run model inference on an Android smartphone with Qualcomm SoCs +### Test model inference on Windows x64 host with QNN HTP emulator / QNN LPAI emulator + +Unlike Linux, which set `LD_LIBRARY_PATH` to access shared libraries, Windows uses the `$env:PATH` environment variable. To enable runtime loading of `qnn_executorch_backend.dll`, ensure that it is discoverable by the Windows DLL loader. + +This can be achieved by either: +- Placing `qnn_executorch_backend.dll` in the same directory as `qnn_executor_runner.exe`; or +- Adding the directory containing `qnn_executorch_backend.dll` to `$env:PATH` environment variable. + +The generated artifacts can be found at: +- `$env:EXECUTORCH_ROOT\build-x86_64-windows\examples\qualcomm\executor_runner\Release\qnn_executor_runner.exe` +- `$env:EXECUTORCH_ROOT\build-x86_64-windows\backends\qualcomm\Release\qnn_executorch_backend.dll` + +To add the directory containing `qnn_executorch_backend.dll` to the `$env:PATH` environment variable: +```powershell +$env:PATH="$env:EXECUTORCH_ROOT\build-x86_64-windows\backends\qualcomm\Release;$env:PATH" +``` + +Once configured, `qnn_executorch_backend.dll` will be accessed by `qnn_executor_runner.exe` at runtime. + +To test the model inference on Windows x64 host with QNN HTP emulator / QNN LPAI emulator: +```powershell +cd $env:EXECUTORCH_ROOT\build-x86_64-windows\examples\qualcomm\executor_runner\Release +.\qnn_executor_runner.exe --model_path $env:EXECUTORCH_ROOT\deeplab_v3\dlv3_qnn.pte +``` + +### Run model inference on Android smartphone with Qualcomm SoCs ***Step 1***. We need to push required QNN libraries to the device. @@ -251,7 +325,7 @@ adb push ${QNN_SDK_ROOT}/lib/aarch64-android/libQnnGpu.so ${DEVICE_DIR} adb push ${QNN_SDK_ROOT}/lib/aarch64-android/libQnnLpai.so ${DEVICE_DIR} adb push ${QNN_SDK_ROOT}/lib/aarch64-android/libQnnLpaiStub.so ${DEVICE_DIR} adb push ${QNN_SDK_ROOT}/lib/aarch64-android/libQnnSystem.so ${DEVICE_DIR} -# make sure the skel lib is signed for LPAI backend. +# Make sure the skel lib is signed for LPAI backend. adb push ${QNN_SDK_ROOT}/lib/lpai-v6/signed/libQnnLpaiSkel.so ${DEVICE_DIR} ``` @@ -298,6 +372,54 @@ After the above command, pre-processed inputs and outputs are put in `$EXECUTORC The command-line arguments are written in [utils.py](https://github.com/pytorch/executorch/blob/main/examples/qualcomm/utils.py#L139). The model, inputs, and output location are passed to `qnn_executorch_runner` by `--model_path`, `--input_list_path`, and `--output_folder_path`. +### Run model inference on Windows on Snapdragon (WoS) with Qualcomm SoCs + +Before running inference on Windows on Snapdragon (WoS) with Qualcomm SoCs, ensure that `qnn_executorch_backend.dll` and all required QNN libraries are discoverable by the Windows loader. This can be achieved by either: +- Copying `qnn_executorch_backend.dll` and the required QNN libraries into the same directory as `qnn_executor_runner.exe`; or +- Adding the directories containing these libraries to the `$env:PATH` environment variable. + +The generated artifacts can be found at: +- `$env:EXECUTORCH_ROOT\build-arm64-windows\examples\qualcomm\executor_runner\Release\qnn_executor_runner.exe` +- `$env:EXECUTORCH_ROOT\build-arm64-windows\backends\qualcomm\Release\qnn_executorch_backend.dll` + +Depending on the selected QNN backend, the corresponding QNN libraries can be found under: + +```powershell +# For HTP +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnHtp.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnSystem.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnHtpV69Stub.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnHtpV73Stub.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnHtpV75Stub.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnHtpV79Stub.dll +$env:QNN_SDK_ROOT\lib\hexagon-v69\unsigned\libQnnHtpV69Skel.so +$env:QNN_SDK_ROOT\lib\hexagon-v73\unsigned\libQnnHtpV73Skel.so +$env:QNN_SDK_ROOT\lib\hexagon-v75\unsigned\libQnnHtpV75Skel.so +$env:QNN_SDK_ROOT\lib\hexagon-v79\unsigned\libQnnHtpV79Skel.so +``` + +```powershell +# For GPU +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnGpu.dll +``` + +```powershell +# For LPAI +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnLpai.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnLpaiStub.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnSystem.dll +# Make sure the skel lib is signed for LPAI backend. +$env:QNN_SDK_ROOT\lib\lpai-v6\signed\libQnnLpaiSkel.so +``` + +Once configured, `qnn_executorch_backend.dll` and the required QNN libraries can be accessed by `qnn_executor_runner.exe` at runtime. + +To test the model inference on Windows on Snapdragon (WoS) with Qualcomm SoCs: +```powershell +cd $env:EXECUTORCH_ROOT +.\qnn_executor_runner.exe --model_path .\deeplab_v3\dlv3_qnn.pte +``` + ### Run [Android LlamaDemo](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/android/LlamaDemo) with QNN backend `$DEMO_APP` refers to the root of the executorch android demo, i.e., the directory containing `build.gradle.kts`. diff --git a/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md b/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md index 0a6a250f968..47eb2fceeb7 100644 --- a/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md +++ b/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md @@ -6,6 +6,8 @@ This backend is in **beta**. It has been validated with a set of small models (e The Arm® Cortex®-M backend accelerates quantized model execution on Arm Cortex-M CPUs using [CMSIS-NN](https://arm-software.github.io/CMSIS-NN/latest/) optimized kernels. Unlike delegate-based backends, it operates as an operator library: quantized subgraphs are replaced with CMSIS-NN accelerated kernels during the pass-lowering stage, while unsupported operators fall back to portable fp32 kernels. +The default AOT flow uses channels-last inputs and the existing dim-order representation. The experimental explicit-layout flow uses ordinary contiguous inputs, represents NCHW/NHWC conversions as graph operators, and selects the experimental `cortex_m::*_nhwc` kernels. Enable it with `--cortex-m-explicit-layout`. Layout modes are selected independently of the Cortex-M CPU target and never mix operator families. + ## Target Support The backend targets Arm Cortex-M CPUs via CMSIS-NN, which provides optimized kernel implementations for three instruction set variants: @@ -98,6 +100,9 @@ quantized = convert_pt2e(prepared) quantized_exported_program = torch.export.export(quantized, (example_input,)) ``` +Calibration observes logical tensor values, so calibration inputs do not need +to use the same memory format as the export example. + ### 2. Lower to edge and apply Cortex-M passes Lower to the edge dialect with the backend's `EdgeCompileConfig`, then run the `CortexMPassManager` to replace quantized subgraphs with CMSIS-NN operator implementations: diff --git a/docs/source/backends/arm-ethos-u/arm-ethos-u-partitioner.md b/docs/source/backends/arm-ethos-u/arm-ethos-u-partitioner.md index 26df600045b..4b674b681dd 100644 --- a/docs/source/backends/arm-ethos-u/arm-ethos-u-partitioner.md +++ b/docs/source/backends/arm-ethos-u/arm-ethos-u-partitioner.md @@ -50,3 +50,19 @@ Returns: def EthosUPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None: ``` Register a custom op to be considered supported. + +```python +def EthosUPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram: +``` +Apply required Arm passes before default ATen decompositions. + +EXIR invokes this backend extension hook automatically through +``to_edge_transform_and_lower``. Model export users should not call it +directly. + +Args: +- **exported_program (ExportedProgram)**: The ATen-dialect program to + transform. + +Returns: +- **ExportedProgram**: The transformed ATen-dialect program. diff --git a/docs/source/backends/arm-vgf/VGF_op_support.md b/docs/source/backends/arm-vgf/VGF_op_support.md index fd27dcde42e..98a3623244f 100644 --- a/docs/source/backends/arm-vgf/VGF_op_support.md +++ b/docs/source/backends/arm-vgf/VGF_op_support.md @@ -6,7 +6,7 @@ This page lists VGF-supported PyTorch APIs and the dtype and quantization modes `8x8` means 8-bit activations and 8-bit weights. `16x8` means 16-bit activations and 8-bit weights. `8x4` means 8-bit activations and 4-bit weights. -Total supported PyTorch APIs: **154**. +Total supported PyTorch APIs: **157**. | PyTorch API | Support profile | DType | Quantization mode | | --- | --- | --- | --- | @@ -43,7 +43,7 @@ Total supported PyTorch APIs: **154**. | `torch.conv1d` | FP, INT | `FP32`, `INT8`, `INT4` | 8x8, 8x4 | | `torch.conv2d` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `INT16`, `INT4` | 8x8, 8x4, 16x8 | | `torch.conv3d` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `INT16`, `INT4` | 8x8, 8x4, 16x8 | -| `torch.conv_transpose2d` | FP, INT | `FP16`, `BF16`, `INT8`, `INT16`, `INT4` | 8x8, 8x4, 16x8 | +| `torch.conv_transpose2d` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `INT16`, `INT4` | 8x8, 8x4, 16x8 | | `torch.cos` | FP, INT | `FP16`, `BF16`, `INT8` | 8x8 | | `torch.cosh` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.cumsum` | FP, INT | `FP32`, `INT8` | 8x8 | @@ -63,12 +63,14 @@ Total supported PyTorch APIs: **154**. | `torch.full_like` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.gather` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `BOOL` | 8x8 | | `torch.ge` / `>=` | FP, INT | `FP32`, `INT8`, `INT16` | 8x8, 16x8 | -| `torch.grid_sampler` | FP | `FP32` | - | -| `torch.grid_sampler_2d` | FP | `FP32` | - | +| `torch.grid_sampler` | FP, INT | `FP32`, `INT8` | 8x8 | +| `torch.grid_sampler_2d` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.group_norm` | FP | `FP32` | - | | `torch.gt` / `>` | FP, INT | `FP32`, `INT8`, `INT16` | 8x8, 16x8 | | `torch.index_put_` | INT | `INT8` | 8x8 | | `torch.index_select` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `BOOL` | 8x8 | +| `torch.isinf` | FP | `FP32` | - | +| `torch.isnan` | FP | `FP32` | - | | `torch.layer_norm` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.le` / `<=` | FP, INT | `FP32`, `INT8`, `INT16` | 8x8, 16x8 | | `torch.linspace` | FP, INT | `FP32`, `INT8` | 8x8 | @@ -125,6 +127,7 @@ Total supported PyTorch APIs: **154**. | `torch.relu` / `torch.nn.ReLU` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.remainder` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.repeat_interleave` | FP, INT | `FP16`, `BF16`, `INT8`, `INT16`, `BOOL` | 8x8, 16x8 | +| `torch.roll` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8` | 8x8 | | `torch.round` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8` | 8x8 | | `torch.rsqrt` | FP, INT | `FP32`, `INT8`, `INT16` | 8x8, 16x8 | | `torch.rsub` | FP | `FP32` | - | @@ -147,7 +150,7 @@ Total supported PyTorch APIs: **154**. | `torch.t` / `torch.Tensor.t` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.tan` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.tanh` / `torch.nn.Tanh` | FP, INT | `FP16`, `BF16`, `INT8` | 8x8 | -| `torch.Tensor.__getitem__` / `tensor indexing` | FP, INT | `FP16`, `BF16`, `INT8` | 8x8 | +| `torch.Tensor.__getitem__` / `tensor indexing` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8` | 8x8 | | `torch.Tensor.__getitem__` / `tensor slicing` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8` | 8x8 | | `torch.Tensor.__setitem__` / `tensor indexing assignment` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.Tensor.copy_` | FP, INT | `FP32`, `INT8` | 8x8 | diff --git a/docs/source/backends/arm-vgf/arm-vgf-partitioner.md b/docs/source/backends/arm-vgf/arm-vgf-partitioner.md index 66701acebaa..620088956bd 100644 --- a/docs/source/backends/arm-vgf/arm-vgf-partitioner.md +++ b/docs/source/backends/arm-vgf/arm-vgf-partitioner.md @@ -50,3 +50,19 @@ Returns: def VgfPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None: ``` Register a custom op to be considered supported. + +```python +def VgfPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram: +``` +Apply required Arm passes before default ATen decompositions. + +EXIR invokes this backend extension hook automatically through +``to_edge_transform_and_lower``. Model export users should not call it +directly. + +Args: +- **exported_program (ExportedProgram)**: The ATen-dialect program to + transform. + +Returns: +- **ExportedProgram**: The transformed ATen-dialect program. diff --git a/docs/source/backends/nxp/nxp-kernel-selection.md b/docs/source/backends/nxp/nxp-kernel-selection.md index 307f06d1d02..4cd3bfc33f7 100644 --- a/docs/source/backends/nxp/nxp-kernel-selection.md +++ b/docs/source/backends/nxp/nxp-kernel-selection.md @@ -1,12 +1,12 @@ # NXP eIQ Neutron Kernel Selective Kernel Registration The NXP ExecuTorch backend supports selective Neutron kernel registration for `Neutron-C` targets, which reduces the -size of the Neutron Firmware. During the backend's conversion to the Neutron representation by the Neutron Converter, +size of the Neutron Firmware. During the backend's conversion to the Neutron representation by the Neutron Compiler, microcode for the Neutron accelerator is generated. The microcode consists of kernel calls executed by the Neutron Driver. The code for kernel call functions is distributed in the Neutron Firmware. -The `eiq_neutron_sdk.neutron_converter` optionally generates a `*_kernel_selection.c` file, registering +The `eiq_neutron_sdk.neutron_compiler` optionally generates a `*_kernel_selection.c` file, registering only kernels that are required for a particular model or, in the case of ExecuTorch, a delegated subgraph. This `*_kernel_selection.c`, when used during application linking, takes precedence over the default list of registered kernels in the Neutron Firmware, and allows the linker to include only the necessary Neutron kernels. @@ -21,7 +21,7 @@ final application with unused code. In memory-constrained environments, you can deployed models. This way you can reduce the size of the final application by linking only selected kernels, used in one or more models. -The feature works as follows: The Neutron Converter with the appropriate flag exports a kernel selection file for each +The feature works as follows: The Neutron Compiler with the appropriate flag exports a kernel selection file for each converted subgraph, the kernel selection files are then merged and ready to be included in the MCUXpresso SDK to use for a selection-only build. @@ -52,7 +52,7 @@ python -m eiq_neutron_sdk.neutron_library_utils.merge_kernel_selection_code \ -output-file merged_kernel_selection.c ``` -Each particular model must be converted by the same Neutron converter version, so the `*_kernel_selection.c` files +Each particular model must be compiled by the same Neutron Compiler version, so the `*_kernel_selection.c` files share the same version. ## MCUXpresso SDK build with kernel selection diff --git a/docs/source/backends/nxp/nxp-mcuxpresso-example.md b/docs/source/backends/nxp/nxp-mcuxpresso-example.md new file mode 100644 index 00000000000..8f8bb93fda3 --- /dev/null +++ b/docs/source/backends/nxp/nxp-mcuxpresso-example.md @@ -0,0 +1,139 @@ +# Using the MCUXpresso Example + +This example demonstrates how to build and run the ExecuTorch CifarNet application for the NXP RT700 platform using the MCUXpresso SDK and the GNU Arm Embedded Toolchain. Before building the project, make sure that all required dependencies are installed and that the necessary environment variables are configured correctly. + +> **Tip:** The `test_build_from_scratch.sh` script automates all the steps described in this guide, including downloading the ARM GNU toolchain, preparing the model, and downloading the MCUXpresso SDK using the `west` tool. If you prefer a fully automated setup, you can run it directly instead of following the manual steps below. + +All scripts described in this guide are located in the following directory of the ExecuTorch repository: + +```text +examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/ +``` + +## 1. Install the Arm GNU Toolchain + +First, download the Arm GCC cross-compilation toolchain that is supported by the RT700 platform: + +```text +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 +``` + +After extracting the archive, create an environment variable called `ARMGCC_DIR` that points to the root directory of the toolchain installation. The build scripts use this variable to locate the compiler, linker, and other required tools. + +Example on Linux: + +```bash +export ARMGCC_DIR=/path/to/arm-gnu-toolchain-15.2.rel1-x86_64-arm-none-eabi +``` + +To verify the installation, you can run: + +```bash +$ARMGCC_DIR/bin/arm-none-eabi-gcc --version +``` + +The command should print the installed compiler version. + +## 2. Download the MCUXpresso SDK + +Next, download MCUXpresso SDK for the RT700 device family using the west tool: + +```bash +pip install west +west init -m https://github.com/nxp-mcuxpresso/mcuxsdk-manifests.git mcuxpresso-sdk +pushd mcuxpresso-sdk +west update_board --set board mimxrt700evk +popd +``` + +Afterwards, configure the `SdkRootDirPath` environment variable to point to the mcuxsdk directory in the downloaded dir. + +Example on Linux: + +```bash +export SdkRootDirPath=/path/to/mcuxpresso-sdk/mcuxsdk +``` + +The build system relies on this variable to locate board support packages, middleware components, startup code, linker scripts, and device-specific libraries. + +## 3. Prepare the Model Header File + +Before building the application, a compiled model must be provided as a C header file named `model_pte.h` and placed in the current directory. Run the provided helper script to generate it: + +```bash +./prepare_model.sh +``` + +The script performs the following steps: + +1. Installs ExecuTorch and its Python dependencies. +2. Installs the `eiq-neutron-sdk` Python package in the version that has been tested with the current ExecuTorch release. +3. Compiles the CifarNet model using the NXP ExecuTorch ahead-of-time (AoT) pipeline and produces a `.pte` model file. +4. Converts the `.pte` file into the `model_pte.h` C header, with the correct memory-section attributes for the RT700 target. + +> **Important:** The MCUXpresso SDK package includes a pre-built CifarNet model and a set of Neutron libraries, but this build flow deliberately does **not** use either of them. Instead, `prepare_model.sh` installs the `eiq-neutron-sdk` version that was tested with the current ExecuTorch release, compiles the model from scratch, and the linker later picks up the matching Neutron libraries from that same installation. This keeps the ExecuTorch AoT compiler, the model bytecode, the Neutron driver, the Neutron firmware, and the ExecuTorch runtime all in sync. + +Once the script finishes, verify that `model_pte.h` was created in the project directory before proceeding to the build step. + +## 4. Build the Application + +Once the environment variables have been configured and `model_pte.h` is present in the project directory, set the `NEUTRON_LIB_DIR` variable to the directory that contains the Neutron static libraries shipped with the eiq-neutron-sdk: + +```bash +export NEUTRON_LIB_DIR=/path/to/eiq_neutron_sdk/libs +``` + +The build script expects the following libraries to exist in that directory: + +- `libNeutronDriver.a` +- `libNeutronFirmware.a` + +Then build the project by executing the provided script: + +```bash +./build_example.sh +``` + +The script validates all required inputs, configures CMake, compiles the source code, links the application, and generates the executable image: + +```text +flash_release/executorch_cifarnet.elf +``` + +If the build completes successfully, the ELF file will be available and ready for programming onto the target board. + +## 5. Flash the Application + +The generated application can be programmed onto the RT700 device using SEGGER J-Link tools. + +### Linux + +```bash +echo "loadfile flash_release/executorch_cifarnet.elf" | \ +/opt/SEGGER/JLink_V796k/JLinkExe \ + -IF SWD \ + -speed auto \ + -Device MIMXRT798S_M33_0 +``` + +Before flashing, ensure that: + +- The board is powered on. +- The JLink debugger probe is flashed on device, if not see [documentation](https://mcuxpresso.nxp.com/mcuxsdk/latest/html/boards/RT/mimxrt700evk/gettingStartedXplorer/topics/program_lpc-link2_with_segger_j-link.html) how to flash it. +- The J-Link debugger is connected to the target. +- The SWD interface is available and correctly wired. +- No other debugging application is currently using the J-Link connection. + +The programming process typically takes only a few seconds. Once the image has been loaded successfully, the application can be started directly from flash memory. + +## 6. Running the Example + +After the firmware is programmed, reset the board and open a serial terminal connected to the device's debug UART interface. The application will initialize the hardware, load the embedded CifarNet model, and begin performing image inference. + +During execution, inference results and diagnostic messages are printed to the terminal. The included demonstration image contains a cat, and the model is expected to classify the image accordingly. + +A successful run produces output similar to the following: + +![example](terminal.png "Example") + +This example serves as a basic validation that the ExecuTorch runtime, model integration, SDK configuration, and hardware platform are all functioning correctly. It can also be used as a starting point for evaluating custom neural network models and experimenting with on-device machine learning workloads on the RT700 platform. diff --git a/docs/source/backends/nxp/nxp-overview.md b/docs/source/backends/nxp/nxp-overview.md index 581c375d038..87249faa05c 100644 --- a/docs/source/backends/nxp/nxp-overview.md +++ b/docs/source/backends/nxp/nxp-overview.md @@ -46,12 +46,17 @@ For a quick overview how to convert a custom PyTorch model, take a look at our [ An example runtime application using the eIQ NSYS (eIQ Neutron Simulator) is available [examples/nxp/executor_runner](https://github.com/pytorch/executorch/blob/main/examples/nxp/executor_runner/), described in the tutorial [Getting started with eIQ Neutron NPU ExecuTorch backend](tutorials/nxp-basic-tutorial.md) -To learn how to run the converted model on the NXP hardware, use one of our example projects on using ExecuTorch runtime from MCUXpresso IDE example projects list. +To learn how to run the delegated model on the NXP hardware, use one of our example projects on using ExecuTorch runtime from MCUXpresso IDE example projects list. For more finegrained tutorial, visit [this manual page](https://mcuxpresso.nxp.com/mcuxsdk/latest/html/middleware/eiq/executorch/docs/nxp/topics/example_applications.html). For guideline how to update the eIQ Neutron Runtime on MCUXpresso SDK, follow the instructions from the eIQ Neutron SDK package `docs/NeutronSDKUserGuide.md` available here https://www.nxp.com/design/design-center/software/eiq-ai-development-environment/eiq-toolkit-for-end-to-end-model-development-and-deployment:EIQ-TOOLKIT. +## Using the MCUXpresso Example + +[This page](nxp-mcuxpresso-example.md) demonstrates how to build and run the ExecuTorch CIFARNet example from MCUXpresso SDK. + + ## Reference **→{doc}`nxp-partitioner` — Partitioner options.** diff --git a/docs/source/backends/nxp/nxp-partitioner.rst b/docs/source/backends/nxp/nxp-partitioner.rst index 4ddc38fb2db..b24f4dde78e 100644 --- a/docs/source/backends/nxp/nxp-partitioner.rst +++ b/docs/source/backends/nxp/nxp-partitioner.rst @@ -27,9 +27,9 @@ Following fields can be set: * `extra_flags` - Extra flags for the Neutron compiler. * `operators_not_to_delegate` - List of operators that will not be delegated. * `use_neutron_for_format_conversion` - If True, let the eIQ Neutron NPU to handle conversion between channel-first (NCHW) and channel-last (NHWC) data formats. That is the Neutron backend will insert `Transpose` ops to ensure that the IO matches the executorch partition, which will be delegated to Neutron. -* `fetch_constants_to_sram` - If True, the Neutron Converter will insert microinstructions to prefetch weights from FLASH to SRAM. This should be used when the whole model does not fit into SRAM on Neutron-C devices, like i.MX RT700. -* `dump_kernel_selection_code` - Whether Neutron converter dumps kernel selection code, which is used by the selective kernel registration, see :doc:`Neutron Firmware Kernel Selection support `. -* `use_profiling` - If true Neutron Converter will enable profiling for neutron delegated model. +* `fetch_constants_to_sram` - If True, the Neutron Compiler will insert microinstructions to prefetch weights from FLASH to SRAM. This should be used when the whole model does not fit into SRAM on Neutron-C devices, like i.MX RT700. +* `dump_kernel_selection_code` - Whether Neutron Compiler dumps kernel selection code, which is used by the selective kernel registration, see :doc:`Neutron Firmware Kernel Selection support `. +* `use_profiling` - If true Neutron Compiler will enable profiling for neutron delegated model. ------------------------- Custom Delegation Options diff --git a/docs/source/backends/nxp/nxp-profiling.md b/docs/source/backends/nxp/nxp-profiling.md index 17e352e479d..1b8f3b49035 100644 --- a/docs/source/backends/nxp/nxp-profiling.md +++ b/docs/source/backends/nxp/nxp-profiling.md @@ -7,16 +7,16 @@ to provide visibility into delegated operator execution time. There are three steps required to obtain profiling results for an NXP‑delegated model: -* Convert the model with profiling support enabled. +* Delegate the model with profiling support enabled. * Generate the artifacts consumed by the Developer Tools (`ETRecord`, `ETDump`). * Create and run the Inspector class to consume these artifacts and print the results. --- -## Convert a model with the profiling support +## Delegate a model with the profiling support Profiling data is generated only for a **profilable** model. -To convert a model with profiling enabled, the `--use-profiling` flag must be set. +To delegate a model with profiling enabled, the `--use-profiling` flag must be set. See the `aot_neutron_compile.py` example and its [README](https://github.com/pytorch/executorch/blob/main/examples/nxp/README.md) @@ -89,7 +89,7 @@ A full implementation is available in [aot_neutron_compile.py](https://github.com/pytorch/executorch/blob/main/examples/nxp/aot_neutron_compile.py). The `--use_profiling` flag is used to create a **profilable** model and the corresponding `ETRecord` file -(see [Convert a model with profiling support](#convert-a-model-with-profiling-support) for the full command). +(see [Delegate a model with the profiling support](#delegate-a-model-with-the-profiling-support) for the full command). --- @@ -102,7 +102,7 @@ The next step is to generate an `ETDump`. An `ETDump` contains runtime data coll To generate an `ETDump`, ensure that the ExecuTorch runtime library is integrated with the Developer Tools and built with the `ET_EVENT_TRACER_ENABLED` flag enabled. -Only models converted with profiling support will produce an `ETDump` containing execution times for all Neutron +Only models delegated with profiling support will produce an `ETDump` containing execution times for all Neutron operators. Otherwise, the dump will include only the final delegate execution time. Neutron software provides a profiling mechanism that logs individual operator execution times to a dedicated runtime @@ -176,7 +176,7 @@ The [Inspector](https://docs.pytorch.org/executorch/1.0/model-inspector.html) AP contents of `ETRecord` and `ETDump`, enabling developers to gain insights into model architecture and performance statistics. -`ETRecord` is an optional argument used to obtain a mapping between the original model and the converted Neutron model. +`ETRecord` is an optional argument used to obtain a mapping between the original model and the delegated Neutron model. An `ETDump` generated on the board contains metadata for each Neutron operator, including its unique identifier. To visualize this metadata in the Inspector results table, set the `include_delegate_debug_data = True` argument. diff --git a/docs/source/backends/nxp/nxp-quantization.md b/docs/source/backends/nxp/nxp-quantization.md index 61cd00632df..3ba39fda0bb 100644 --- a/docs/source/backends/nxp/nxp-quantization.md +++ b/docs/source/backends/nxp/nxp-quantization.md @@ -10,6 +10,7 @@ The Neutron delegate supports the following quantization schemes: - Static quantization with 8-bit symmetric weights and 8-bit asymmetric activations (via the PT2E quantization flow), per-tensor granularity. - Following operators are supported at this moment: - `aten.abs.default` + - `aten.adaptive_avg_pool1d.default` - `aten.adaptive_avg_pool2d.default` - `aten.add.Tensor` - `aten.addmm.default` @@ -248,7 +249,7 @@ Moving from PTQ to QAT check-list: #### Known limitations of QAT In the current ExecuTorch/TorchAO implementation, there is an issue when quantizing biasless convolutions during QAT. -The pipeline produces a non‑quantized empty bias, which causes the Neutron Converter to fail. +The pipeline produces a non‑quantized empty bias, which causes the Neutron Compiler to fail. To mitigate this issue, use the `QuantizeFusedConvBnBiasAtenPass` post‑quantization: ```python diff --git a/docs/source/backends/nxp/op-support.csv b/docs/source/backends/nxp/op-support.csv index b6aa870e31a..368b981face 100644 --- a/docs/source/backends/nxp/op-support.csv +++ b/docs/source/backends/nxp/op-support.csv @@ -1,6 +1,7 @@ Operator,Compute DType,Quantization,Constraints aten.abs.default,int8,static int8, -aten._adaptive_avg_pool2d.default,int8,static int8,"ceil_mode=False, count_include_pad=False, divisor_override=False" +aten._adaptive_avg_pool2d.default,int8,static int8,"Must be representable by avg_pool2D" +aten.adaptive_avg_pool1d.default,int8,static int8,"Must be representable by avg_pool2D" aten.addmm.default,int8,static int8,2D tensor only aten.add.Tensor,int8,static int8,"alpha = 1" aten.amax.default,int8,static int8, diff --git a/docs/source/backends/nxp/terminal.png b/docs/source/backends/nxp/terminal.png new file mode 100644 index 00000000000..9a4f1d70ee8 Binary files /dev/null and b/docs/source/backends/nxp/terminal.png differ diff --git a/docs/source/backends/nxp/tutorials/nxp-basic-tutorial.md b/docs/source/backends/nxp/tutorials/nxp-basic-tutorial.md index b2e07bb7c1d..264a3fc2dfc 100644 --- a/docs/source/backends/nxp/tutorials/nxp-basic-tutorial.md +++ b/docs/source/backends/nxp/tutorials/nxp-basic-tutorial.md @@ -13,14 +13,14 @@ You need to install the ExecuTorch. Please follow the tutorial to install the Ex In addition to this, you will need to install the eIQ Neutron Simulator, called NSYS, -and the Neutron Converter for generating the byte-code for the eIQ Neutron NPU, +and the Neutron Compiler for generating the byte-code for the eIQ Neutron NPU, during the model conversion in ExecuTorch AoT flow. To install the eIQ Neutron dependencies, run: ```bash examples/nxp/setup.sh ``` This will install: -* Neutron Converter, for converting the Neutron IR to Neutron byte-code +* Neutron Compiler, for compiling the Neutron IR to Neutron byte-code * eIQ Neutron SDK, containing the eIQ Neutron runtimes (driver and firmware) for various NXP SoC and simulator * eIQ NSYS, the Neutron behavioral simulator diff --git a/docs/source/backends/vulkan/tutorials/etvk-llama-tutorial.md b/docs/source/backends/vulkan/tutorials/etvk-llama-tutorial.md index cb14c72331e..42a7a846d75 100644 --- a/docs/source/backends/vulkan/tutorials/etvk-llama-tutorial.md +++ b/docs/source/backends/vulkan/tutorials/etvk-llama-tutorial.md @@ -93,7 +93,7 @@ cmake . \ -DEXECUTORCH_BUILD_VULKAN=ON \ -DEXECUTORCH_BUILD_TESTS=OFF \ -Bcmake-out-android-so && \ -cmake --build cmake-out-android-so -j16 --target install --config Release +cmake --build cmake-out-android-so -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install --config Release ``` ## Build and push the llama runner binary to Android @@ -111,7 +111,7 @@ cmake examples/models/llama \ -DCMAKE_BUILD_TYPE=Release \ -DPYTHON_EXECUTABLE=python \ -Bcmake-out-android-so/examples/models/llama && \ -cmake --build cmake-out-android-so/examples/models/llama -j16 --config Release +cmake --build cmake-out-android-so/examples/models/llama -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release ``` Once the binary is built, it can be pushed to your Android device. diff --git a/docs/source/backends/vulkan/tutorials/etvk-profiling-tutorial.md b/docs/source/backends/vulkan/tutorials/etvk-profiling-tutorial.md index 07982d81c1c..9330ba370e8 100644 --- a/docs/source/backends/vulkan/tutorials/etvk-profiling-tutorial.md +++ b/docs/source/backends/vulkan/tutorials/etvk-profiling-tutorial.md @@ -71,7 +71,7 @@ cmake . \ -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=ON \ -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \ -Bcmake-out-android-so && \ -cmake --build cmake-out-android-so -j16 --target install --config Release +cmake --build cmake-out-android-so -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install --config Release ``` Once the build completes, we can push the runner binary to device. diff --git a/docs/source/backends/webgpu/webgpu-overview.md b/docs/source/backends/webgpu/webgpu-overview.md index 91a388ac576..ebc46031a53 100644 --- a/docs/source/backends/webgpu/webgpu-overview.md +++ b/docs/source/backends/webgpu/webgpu-overview.md @@ -68,7 +68,7 @@ depend on the selected WebGPU adapter. ## Development Requirements -- CMake 3.19 or later. +- CMake 3.24 or later. - A Python environment with ExecuTorch installed for model export. - Dawn's CMake package for native builds. - Emscripten for browser builds. diff --git a/docs/source/llm/run-with-c-plus-plus.md b/docs/source/llm/run-with-c-plus-plus.md index b6c6082c3a6..69d98325e4f 100644 --- a/docs/source/llm/run-with-c-plus-plus.md +++ b/docs/source/llm/run-with-c-plus-plus.md @@ -12,7 +12,7 @@ Before you begin, make sure you have: - For HuggingFace tokenizers, this is a JSON file `tokenizer.json` - For SentencePiece tokenizers, this is a `tokenizer.model` file and normally lives alongside the weights file 3. CMake and a C++ compiler installed - - CMake version 3.29 or higher + - CMake version 3.26 or higher - g++ or clang compiler ## Model Metadata diff --git a/docs/source/raspberry_pi_llama_tutorial.md b/docs/source/raspberry_pi_llama_tutorial.md index 6075e455c9b..9eb99711ef5 100644 --- a/docs/source/raspberry_pi_llama_tutorial.md +++ b/docs/source/raspberry_pi_llama_tutorial.md @@ -21,7 +21,7 @@ This tutorial demonstrates how to deploy **Llama models on Raspberry Pi 4/5 devi - **Python 3.10-3.14** (ExecuTorch requirement) - **conda** or **venv** for environment management -- **CMake 3.29.6+** +- **CMake 3.26+** - **Git** for repository cloning ### Target Device Requirements @@ -47,7 +47,7 @@ python3 --version # Should be 3.10-3.14 # Check required tools hash cmake git md5sum 2>/dev/null || echo "Missing required tools" -cmake --version # Should be 3.29.6+ at minimum +cmake --version # Should be 3.26+ at minimum ## Development Environment Setup diff --git a/docs/source/tutorial-xnnpack-delegate-lowering.md b/docs/source/tutorial-xnnpack-delegate-lowering.md index 5c88246b0ba..81b26db2f12 100644 --- a/docs/source/tutorial-xnnpack-delegate-lowering.md +++ b/docs/source/tutorial-xnnpack-delegate-lowering.md @@ -166,7 +166,7 @@ cmake \ Then you can build the runtime componenets with ```bash -cmake --build cmake-out -j9 --target install --config Release +cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install --config Release ``` Now you should be able to find the executable built at `./cmake-out/executor_runner` you can run the executable with the model you generated as such diff --git a/docs/source/using-executorch-building-from-source.md b/docs/source/using-executorch-building-from-source.md index e8c7f0e5a3b..03d73420bb9 100644 --- a/docs/source/using-executorch-building-from-source.md +++ b/docs/source/using-executorch-building-from-source.md @@ -79,8 +79,29 @@ portability details. * `--clean`: Removes build artifacts. * `--editable`: Install the ExecuTorch python package in editable mode (see [Editable Install](#editable-install)). * `--minimal`: Install only the minimal set of dependencies required to run ExecuTorch. Do not install dependencies for examples. + * `--optional-dependency `: Install an optional Python dependency set. + Repeat the flag to select more than one. Supported names are `ethos_u`, + `vgf`, and `openvino`. * `--use-pt-pinned-commit`: Install the pinned PyTorch commit or release version. When not specified, the latest PyTorch nightly build is installed. + For example, install the current checkout with the dependencies needed for + Ethos-U ahead-of-time (AOT) export: + + ```bash + ./install_executorch.sh --optional-dependency ethos_u + ``` + + After the base dependencies have already been installed, the equivalent + editable package command is: + + ```bash + pip install -e '.[ethos_u]' --no-build-isolation + ``` + + The `ethos_u` optional dependencies are host-side Python tools used during + AOT export. Embedded toolchains, simulators, and target runtimes are + configured separately by the backend setup and build instructions. + For Intel-based macOS systems, use `--use-pt-pinned-commit --minimal`. As PyTorch does not provide pre-built binaries for Intel Mac, installation requires building PyTorch from source. Instructions can be found in [PyTorch Installation](https://github.com/pytorch/pytorch#installation). Note that only the XNNPACK and CoreML backends are built by default. Additional backends can be enabled or disabled by setting the corresponding CMake flags: diff --git a/docs/source/using-executorch-ios.md b/docs/source/using-executorch-ios.md index 7053c28fd76..c4dc29e198e 100644 --- a/docs/source/using-executorch-ios.md +++ b/docs/source/using-executorch-ios.md @@ -171,7 +171,7 @@ OTHER_LDFLAGS = $(inherited) \ **Note:** In the example above, we link against the Debug version of the ExecuTorch runtime (`libexecutorch_debug`) to preserve the logs. Normally, that does not impact the performance too much. Nevertheless, remember to link against the release version of the runtime (`libexecutorch`) for the best performance and no logs. -**Note:** The MLX backend loads its Metal kernels at runtime from a per-slice metallib inside a resource bundle named `executorch_backend_mlx_resources`, not from the frameworks in `cmake-out`. The build stages the correctly named files (`mlx-ios.metallib`, `mlx-ios-simulator.metallib`, `mlx-macos.metallib`) under `.Package.swift/backend_mlx_resources/`. If you integrate MLX from a source build, ship those files in a bundle of that name for the slices you use, or MLX links and registers but has no kernels to run. +**Note:** The MLX backend loads its Metal kernels at runtime from a per-slice metallib inside a resource bundle named `executorch_backend_mlx_resources`, not from the frameworks in `cmake-out`. The Apple framework presets enable this resource lookup when MLX is available, and the framework build stages the correctly named files (`mlx-ios.metallib`, `mlx-ios-simulator.metallib`, `mlx-macos.metallib`) under `.Package.swift/backend_mlx_resources/`. Generic Apple presets retain MLX's native colocated-metallib lookup. If you use a custom CMake configuration for SwiftPM packaging, enable `EXECUTORCH_MLX_SWIFTPM_RESOURCES` and ship the matching slice in a bundle of that name. You can assign such a config file to your target in Xcode: diff --git a/examples/apple/coreml/llama/BUCK b/examples/apple/coreml/llama/BUCK index 8604a50d9f7..ac9cd6eaf3a 100644 --- a/examples/apple/coreml/llama/BUCK +++ b/examples/apple/coreml/llama/BUCK @@ -4,6 +4,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "no load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target(_kind = runtime.python_library, name = "llama_transformer", srcs = [ diff --git a/examples/apple/coreml/scripts/BUCK b/examples/apple/coreml/scripts/BUCK index 42a97ea893f..2685fe53140 100644 --- a/examples/apple/coreml/scripts/BUCK +++ b/examples/apple/coreml/scripts/BUCK @@ -3,6 +3,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "no # targets.bzl. This file can contain fbcode-only targets. load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary") +oncall("executorch") + fbcode_target(_kind = python_binary, name = "extract_coreml_models", srcs = [ diff --git a/examples/arm/README.md b/examples/arm/README.md index 1a0923f1ab3..8efffd5c377 100644 --- a/examples/arm/README.md +++ b/examples/arm/README.md @@ -11,13 +11,29 @@ This directory contains documentation and scripts to help you setup and run a PyTorch model on the Arm backend via ExecuTorch. +## Python package setup + +For Ethos-U examples, install the current checkout and the dependencies needed +for ahead-of-time (AOT) export in a clean Python environment: + +```bash +python3.12 -m venv .venv +source .venv/bin/activate +./install_executorch.sh --optional-dependency ethos_u +``` + +After the base dependencies are installed, the equivalent editable package +command is `pip install -e '.[ethos_u]' --no-build-isolation`. The `ethos_u` +extra provides host-side export dependencies; it does not install the Arm +toolchain, FVPs, or target runtime. Run `setup.sh` below to install the cross +compiler, FVPs, and backend tools used by these examples. + ## setup.sh `setup.sh` downloads the Arm cross-compilation toolchain and Corstone FVP -simulators, installs the Python dependencies for TOSA, Ethos-U Vela, and -Cortex-M/CMSIS-NN, and generates `setup_path.sh` scripts for adding those tools -to your environment. Optional flags also install VGF/MLSDK and Vulkan -dependencies. +simulators, installs the backend dependencies, and generates `setup_path.sh` +scripts for adding those tools to your environment. Optional flags also install +VGF/MLSDK and Vulkan dependencies. Example to install the default Arm backend dependencies and add them to your current shell: @@ -96,6 +112,8 @@ For Cortex-M testing, use a Cortex-M target and bundled I/O: BundleIO, ETDump, profiling, semihosted files, and backend regression tests. - [ethos-u-porting-guide.md](ethos-u-porting-guide.md) - Notes for adapting the example Ethos-U runtime integration to another target. +- [model-explorer.md](model-explorer.md) - Visualize PTE and TOSA graphs and + overlay per-operator Ethos-U cycle data collected from an FVP PMU trace. - [export_standalone_tosa_graph.py](export_standalone_tosa_graph.py) - Example of exporting a standalone TOSA graph with multiple outputs. - [visualize.py](visualize.py) - Helper used by `run.sh --model_explorer` to diff --git a/examples/arm/ethos_u_minimal_example.ipynb b/examples/arm/ethos_u_minimal_example.ipynb index 11f24019d23..d9298237b95 100644 --- a/examples/arm/ethos_u_minimal_example.ipynb +++ b/examples/arm/ethos_u_minimal_example.ipynb @@ -21,11 +21,17 @@ "This guide demonstrates the full flow for running a module on Arm Ethos-U55 using ExecuTorch.\n", "Tested on Linux x86_64 and macOS aarch64. If something is not working for you, please raise a GitHub issue and tag Arm.\n", "\n", - "Before you begin:\n", - "1. (In a clean virtual environment with a compatible Python version) Install executorch using `./install_executorch.sh`\n", - "2. Install Arm cross-compilation toolchain and simulators using `./examples/arm/setup.sh --i-agree-to-the-contained-eula`\n", + "Before you begin, run these commands from the base `executorch` folder:\n", "\n", - "With all commands executed from the base `executorch` folder.\n", + "```bash\n", + "python3.12 -m venv .venv\n", + "source .venv/bin/activate\n", + "./install_executorch.sh --optional-dependency ethos_u\n", + "./examples/arm/setup.sh --i-agree-to-the-contained-eula\n", + "source examples/arm/arm-scratch/setup_path.sh\n", + "```\n", + "\n", + "`--optional-dependency ethos_u` installs the Python tools needed to export models for Ethos-U, including Vela. The Arm setup script separately installs the cross compiler and FVPs used later in the notebook.\n", "\n", "\n", "\n", diff --git a/examples/arm/executor_runner/CMakeLists.txt b/examples/arm/executor_runner/CMakeLists.txt index ba80df19ddf..6313ceedbbc 100644 --- a/examples/arm/executor_runner/CMakeLists.txt +++ b/examples/arm/executor_runner/CMakeLists.txt @@ -54,6 +54,16 @@ set(ET_NUM_INFERENCES option(ET_LOG_DUMP_INPUT "Dump input in log" OFF) option(ET_LOG_DUMP_OUTPUT "Dump output in log" ON) +option(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING + "Report Ethos-U PMU statistics per delegate" OFF +) +option(ET_ARM_ETHOSU_PROFILE_IO_COPIES + "Count Ethos-U backend I/O memcpy calls and bytes" OFF +) +set(ET_ARM_ETHOSU_MAX_PROFILED_DELEGATES + "16" + CACHE STRING "Maximum delegates tracked by per-delegate profiling" +) option(ET_BUNDLE_IO "Set to compile in BundleIO support" OFF) set(BUNDLED_PROGRAM_LIBRARY_DIR @@ -338,6 +348,33 @@ if(ET_NUM_INFERENCES) ) endif() +if(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + if(NOT ET_ARM_ETHOSU_MAX_PROFILED_DELEGATES MATCHES "^[1-9][0-9]*$") + message( + FATAL_ERROR + "ET_ARM_ETHOSU_MAX_PROFILED_DELEGATES must be a positive integer" + ) + endif() + target_compile_definitions( + arm_executor_runner + PRIVATE + ET_ARM_ETHOSU_PER_DELEGATE_PROFILING + ET_ARM_ETHOSU_MAX_PROFILED_DELEGATES=${ET_ARM_ETHOSU_MAX_PROFILED_DELEGATES} + ) + target_compile_definitions( + executorch_delegate_ethos_u PRIVATE ET_ARM_ETHOSU_PER_DELEGATE_PROFILING + ) +endif() + +if(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + target_compile_definitions( + arm_executor_runner PRIVATE ET_ARM_ETHOSU_PROFILE_IO_COPIES + ) + target_compile_definitions( + executorch_delegate_ethos_u PRIVATE ET_ARM_ETHOSU_PROFILE_IO_COPIES + ) +endif() + if(ET_LOG_DUMP_INPUT) target_compile_definitions(arm_executor_runner PUBLIC ET_LOG_DUMP_INPUT) endif() diff --git a/examples/arm/executor_runner/arm_perf_monitor.cpp b/examples/arm/executor_runner/arm_perf_monitor.cpp index 59daede6920..4f75d93d0c7 100644 --- a/examples/arm/executor_runner/arm_perf_monitor.cpp +++ b/examples/arm/executor_runner/arm_perf_monitor.cpp @@ -40,6 +40,48 @@ uint64_t ethosu_ArmBackendExecuteCycleCount = 0; uint64_t ethosu_ArmWhenNPURunCycleCountStart = 0; uint64_t ethosu_ArmWhenNPURunCycleCount = 0; uint64_t ethosu_pmuCycleCount = 0; +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) +struct IOCopyStats { + uint64_t calls = 0; + uint64_t bytes = 0; +}; +IOCopyStats ethosu_inputCopyStats; +IOCopyStats ethosu_outputCopyStats; +#endif +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) +struct DelegateStats { + const void* handle = nullptr; + uint64_t backend_invocations = 0; + uint64_t npu_invocations = 0; + uint64_t pmu_cycles = 0; + std::array pmu_events{}; +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + IOCopyStats input_copies; + IOCopyStats output_copies; +#endif +}; + +std::array + ethosu_delegateStats; +size_t ethosu_delegateCount = 0; +DelegateStats* ethosu_activeDelegate = nullptr; +bool ethosu_delegateCapacityExceeded = false; + +DelegateStats* get_delegate_stats(const void* handle) { + for (size_t i = 0; i < ethosu_delegateCount; ++i) { + if (ethosu_delegateStats[i].handle == handle) { + return ðosu_delegateStats[i]; + } + } + if (ethosu_delegateCount == ethosu_delegateStats.size()) { + ethosu_delegateCapacityExceeded = true; + return nullptr; + } + DelegateStats& stats = ethosu_delegateStats[ethosu_delegateCount++]; + stats.handle = handle; + return &stats; +} +#endif std::array ethosu_pmuEventCounts = {0}; // ethosu_pmuCountersUsed should match numbers of counters setup in @@ -50,6 +92,43 @@ static_assert(ETHOSU_PMU_NCOUNTERS >= ethosu_pmuCountersUsed); extern "C" { +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) +void EthosUBackend_input_memcpy(size_t size) { + ethosu_inputCopyStats.calls++; + ethosu_inputCopyStats.bytes += size; +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + if (ethosu_activeDelegate != nullptr) { + ethosu_activeDelegate->input_copies.calls++; + ethosu_activeDelegate->input_copies.bytes += size; + } +#endif +} + +void EthosUBackend_output_memcpy(size_t size) { + ethosu_outputCopyStats.calls++; + ethosu_outputCopyStats.bytes += size; +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + if (ethosu_activeDelegate != nullptr) { + ethosu_activeDelegate->output_copies.calls++; + ethosu_activeDelegate->output_copies.bytes += size; + } +#endif +} +#endif + +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) +void EthosUBackend_delegate_begin(const void* handle) { + ethosu_activeDelegate = get_delegate_stats(handle); + if (ethosu_activeDelegate != nullptr) { + ethosu_activeDelegate->backend_invocations++; + } +} + +void EthosUBackend_delegate_end() { + ethosu_activeDelegate = nullptr; +} +#endif + // Callback invoked at start of NPU execution void ethosu_inference_begin(struct ethosu_driver* drv, void*) { // Enable PMU @@ -100,10 +179,27 @@ void ethosu_inference_begin(struct ethosu_driver* drv, void*) { // Callback invoked at end of NPU execution void ethosu_inference_end(struct ethosu_driver* drv, void*) { ethosu_delegation_count++; +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + const uint64_t pmu_cycles = ETHOSU_PMU_Get_CCNTR(drv); + ethosu_pmuCycleCount += pmu_cycles; + if (ethosu_activeDelegate != nullptr) { + ethosu_activeDelegate->npu_invocations++; + ethosu_activeDelegate->pmu_cycles += pmu_cycles; + } +#else ethosu_pmuCycleCount += ETHOSU_PMU_Get_CCNTR(drv); +#endif for (size_t i = 0; i < ethosu_pmuCountersUsed; i++) { +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + const uint64_t event_count = ETHOSU_PMU_Get_EVCNTR(drv, i); + ethosu_pmuEventCounts[i] += event_count; + if (ethosu_activeDelegate != nullptr) { + ethosu_activeDelegate->pmu_events[i] += event_count; + } +#else ethosu_pmuEventCounts[i] += ETHOSU_PMU_Get_EVCNTR(drv, i); +#endif } ETHOSU_PMU_Disable(drv); // Add Cortex-M cycle clock used during this NPU execution @@ -127,10 +223,26 @@ void EthosUBackend_execute_end() { } void StartMeasurements() { +#if defined(__ARM_ARCH_8_1M_MAIN__) + // StopMeasurements() disables the cycle counter after each measurement. + // Server mode starts a new measurement for every inference. + ARM_PMU_Enable(); + ARM_PMU_CNTR_Enable(PMU_CNTENSET_CCNTR_ENABLE_Msk); +#endif ethosu_delegation_count = 0; ethosu_ArmBackendExecuteCycleCount = 0; ethosu_ArmWhenNPURunCycleCount = 0; ethosu_pmuCycleCount = 0; +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + ethosu_inputCopyStats = {}; + ethosu_outputCopyStats = {}; +#endif +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + ethosu_delegateStats = {}; + ethosu_delegateCount = 0; + ethosu_activeDelegate = nullptr; + ethosu_delegateCapacityExceeded = false; +#endif for (size_t i = 0; i < ethosu_pmuCountersUsed; i++) { ethosu_pmuEventCounts[i] = 0; @@ -162,6 +274,32 @@ void StopMeasurements(int num_inferences) { "ethos-u : cycle_cnt : %" PRIu64 " cycles (%.2f per inference)", ethosu_ArmBackendExecuteCycleCount, (double)ethosu_ArmBackendExecuteCycleCount / num_inferences); +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + const uint64_t io_copy_calls = + ethosu_inputCopyStats.calls + ethosu_outputCopyStats.calls; + const uint64_t io_copy_bytes = + ethosu_inputCopyStats.bytes + ethosu_outputCopyStats.bytes; + ET_LOG( + Info, + "Ethos-U IO copy calls: %" PRIu64 " (%.2f per inference)", + io_copy_calls, + (double)io_copy_calls / num_inferences); + ET_LOG( + Info, + "Ethos-U IO copy bytes: %" PRIu64 " bytes (%.2f per inference)", + io_copy_bytes, + (double)io_copy_bytes / num_inferences); + ET_LOG( + Info, + "Ethos-U input copy: %" PRIu64 " calls, %" PRIu64 " bytes", + ethosu_inputCopyStats.calls, + ethosu_inputCopyStats.bytes); + ET_LOG( + Info, + "Ethos-U output copy: %" PRIu64 " calls, %" PRIu64 " bytes", + ethosu_outputCopyStats.calls, + ethosu_outputCopyStats.bytes); +#endif // We could print a list of the cycles used by the other delegates here in the // future but now we only print ethos-u: this means that "Operator(s) total: // ..." will be the same number as ethos-u : cycle_cnt and not the sum of all @@ -221,6 +359,54 @@ void StopMeasurements(int num_inferences) { ethosu_pmuEventCounts[i], (double)ethosu_pmuEventCounts[i] / num_inferences); } +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + ET_LOG(Info, "Ethos-U per-delegate PMU report:"); + for (size_t delegate_id = 0; delegate_id < ethosu_delegateCount; + ++delegate_id) { + const DelegateStats& stats = ethosu_delegateStats[delegate_id]; + ET_LOG( + Info, + "Ethos-U delegate %zu: %" PRIu64 " backend invocations, %" PRIu64 + " NPU invocations", + delegate_id, + stats.backend_invocations, + stats.npu_invocations); + ET_LOG( + Info, + "Ethos-U delegate %zu PMU cycles: %" PRIu64, + delegate_id, + stats.pmu_cycles); + for (size_t event = 0; event < ethosu_pmuCountersUsed; ++event) { + ET_LOG( + Info, + "Ethos-U delegate %zu PMU counter %zu: %" PRIu64, + delegate_id, + event, + stats.pmu_events[event]); + } +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + ET_LOG( + Info, + "Ethos-U delegate %zu input copy: %" PRIu64 " calls, %" PRIu64 " bytes", + delegate_id, + stats.input_copies.calls, + stats.input_copies.bytes); + ET_LOG( + Info, + "Ethos-U delegate %zu output copy: %" PRIu64 " calls, %" PRIu64 + " bytes", + delegate_id, + stats.output_copies.calls, + stats.output_copies.bytes); +#endif + } + if (ethosu_delegateCapacityExceeded) { + ET_LOG( + Error, + "Ethos-U per-delegate profiling exceeded its capacity of %zu delegates", + ethosu_delegateStats.size()); + } +#endif #if defined(ETHOSU55) || defined(ETHOSU65) ET_LOG( Info, diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/README.md b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/README.md index 58eae2b0201..ee2359c171d 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/README.md +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/README.md @@ -1,43 +1,79 @@ -# MobileSAM Prompt Segmentation Example Application - -This end-to-end example shows how to use the Arm Ethos-U backend in -ExecuTorch for transformer-based prompt segmentation. MobileSAM predicts a -binary mask for fixed positive point prompts rather than semantic class IDs. -The host debug flow validates quantization by comparing FP32 and quantized -masks, with an optional binary reference mask when one is available. - -It covers: - -- Loading the MobileSAM `vit_t` checkpoint. -- Freezing one or more positive point prompts into the exported graph. -- Applying post-training quantization with the Ethos-U quantizer. -- Lowering the quantized model to an Ethos-U85-256 ExecuTorch program. -- Producing validation and debugging artifacts such as masks, overlays, - mismatch heatmaps, metrics, and delegation summaries. -- Building a bare-metal Corstone-320 runtime app and running it on FVP. - -The default export uses a reduced `448x448` image input and returns one -low-resolution `[1, 1, 112, 112]` mask-logit tensor. The example prepares the -official MobileSAM GitHub source at a pinned revision in an external checkout -and applies a small configurable-input patch there. Neither the MobileSAM -source nor checkpoint is redistributed in ExecuTorch. - -The export uses int8 activations and int8 weights globally, and A16W8 -quantization for TinyViT attention modules. This keeps the transformer -attention numerically stable while still producing one Ethos-U delegate. - -The exported graph intentionally uses `multimask_output=False` and leaves -mask thresholding outside the model. SAM-style candidate-mask selection can be -numerically sensitive after export and quantization, so this example keeps the -target graph focused on the fixed-prompt image encoder and mask decoder. - -## Layout - -- `model_export/prepare_mobilesam.py` - Prepares the pinned external MobileSAM - checkout and applies the configurable-input patch. -- `model_export/README.md` - Model loading, quantization, lowering, - validation, and debug artifact generation. -- `runtime/README.md` - Bare-metal runtime build, image header generation, and - Corstone-320 FVP execution. -- `runtime/visualize_fvp_output.py` - Decodes the target mask dump, creates an - overlay, and compares FVP output with the host quantized mask. +# MobileSAM Prompt Segmentation on Ethos-U + +This example turns a point on an image into an object mask. It shows the full +ExecuTorch flow: export MobileSAM, quantize it, delegate it to Ethos-U85, run it +on the Corstone-320 FVP, and compare the target result with the host result. + +There is one tested configuration: MobileSAM `vit_t`, a `448x448` input, and +Ethos-U85-256. The image can change at runtime, but the point prompt is embedded +in the exported model. Changing the prompt requires re-exporting the model. + +## Run It + +From the ExecuTorch repository root: + +```bash +python3.12 -m venv .venv +source .venv/bin/activate +./install_executorch.sh --optional-dependency ethos_u +./examples/arm/setup.sh --i-agree-to-the-contained-eula +./examples/arm/mobilesam_prompt_segmentation_example_ethos_u/run.sh +``` + +The final command performs the complete flow and prints +`MobileSAM example: PASS`. Its main result is: + +`arm_test/mobilesam/result/fvp_comparison.png` + +## What It Does + +1. Fetches the pinned official MobileSAM source and checkpoint outside the + repository. +2. Runs `torch.export`, PT2E quantization, and `EthosUPartitioner` to create a + `.pte` containing one Ethos-U delegate. +3. Builds the standard Arm ExecuTorch runner and runs one inference on FVP. +4. Compares the FVP mask with the host quantized mask and requires `0.9` IoU. + +Successful completion creates: + +- Program: `arm_test/mobilesam/export/mobilesam.pte` +- Host masks: `arm_test/mobilesam/export/fp32_mask.png` and + `arm_test/mobilesam/export/quantized_mask.png` +- FVP log: `arm_test/mobilesam/fvp.log` +- Comparison: `arm_test/mobilesam/result/fvp_comparison.png` +- FVP validation: `arm_test/mobilesam/result/metrics.json` +- TOSA and Vela artifacts: `arm_test/mobilesam/export/artifacts` + +The Python installer uses this source checkout and installs the dependencies +needed for ahead-of-time Ethos-U export. The Arm setup script installs the +cross compiler and FVP. Do not install a separate PyPI `executorch` wheel for +this source example. + +On macOS, Docker must be running and the +[FVPs-on-Mac](https://github.com/Arm-Examples/FVPs-on-Mac) wrapper must be on +`PATH`. + +## Code Map + +- [`prepare_mobilesam.py`](model_export/prepare_mobilesam.py) fetches and + verifies the external model. +- [`export_mobilesam.py`](model_export/export_mobilesam.py) contains the model, + quantization, validation, and lowering flow. +- [`run.sh`](run.sh) uses ExecuTorch's standard Arm runner for target execution. +- [`visualize_fvp_output.py`](runtime/visualize_fvp_output.py) checks and plots + the raw output tensor. + +There is no MobileSAM-specific C++ runtime or CMake project. + +## Limitations + +- The exported model accepts one image tensor and uses one fixed positive point. +- It returns a low-resolution mask. Upsampling and thresholding are host-side + post-processing. +- The demo image is also the calibration image. Product use requires a + representative calibration set. +- The default fast FVP mode validates correctness. Its counters are not a + performance benchmark or a measurement of real-device latency. + +See [model export](model_export/README.md) and +[runtime](runtime/README.md) for details of each stage. diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/README.md b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/README.md index 8069c0d8cc4..15c6043f3d0 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/README.md +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/README.md @@ -1,192 +1,29 @@ -# MobileSAM Export, Quantization, and Debugging +# MobileSAM Export -This directory exports MobileSAM `vit_t` for the ExecuTorch Ethos-U backend. -The exporter freezes one or more positive point prompts into the graph, so the -runtime app has one tensor input: the preprocessed image. The model returns one -mask-logit tensor for those prompts. The default export uses a `448x448` input, -which produces a `112x112` mask-logit tensor. +The exporter keeps one tested model configuration so the ExecuTorch steps are +visible without a layer of command-line configuration. -Production SAM applications usually pass prompts dynamically from a UI or -tracking pipeline. This example freezes the prompt embeddings so the FVP -runtime stays small and demonstrates the Ethos-U flow with the same image -encoder and mask decoder used by MobileSAM. - -## Requirements - -- Python 3.10+ with `executorch`. -- Dependencies from - `examples/arm/mobilesam_prompt_segmentation_example_ethos_u/requirements.txt`. -- Git and internet access to prepare the pinned external MobileSAM checkout and - download the official checkpoint, unless both are already cached. -- Ethos-U dependencies from `examples/arm/setup.sh`. - -MobileSAM's A16W8 attention requires `ethos-u-vela>=5.1.0`. Vela 5.0 produces -incorrect Ethos-U85 INT16 reductions, so the exporter rejects that version -before generating a `.pte`. - -## Export - -Run from the ExecuTorch repo root: - -```bash -MOBILE_SAM_SOURCE="$HOME/.cache/executorch/mobilesam/f706ad9c4eb7f219c00d9050e46328518ffb65d2/source" -python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py \ - --source-dir "$MOBILE_SAM_SOURCE" - -python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py \ - --output-path ./mobilesam_point_ethos_u85_448.pte \ - --mobile-sam-source "$MOBILE_SAM_SOURCE" \ - --calibration-image examples/models/dinov2/dog.jpg \ - --eval-image examples/models/dinov2/dog.jpg \ - --point 219 193 \ - --artifact-dir ./mobilesam_point_artifacts \ - --debug-output-dir ./mobilesam_point_debug -``` - -The default configuration is: - -- Model source: `https://github.com/ChaoningZhang/MobileSAM` -- MobileSAM source revision: `f706ad9c4eb7f219c00d9050e46328518ffb65d2` -- External source patch: `0001-Make-TinyViT-image-size-configurable.patch` -- Checkpoint URL: - `https://github.com/ChaoningZhang/MobileSAM/raw/f706ad9c4eb7f219c00d9050e46328518ffb65d2/weights/mobile_sam.pt` -- Checkpoint SHA256: - `6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f` -- MobileSAM source-code license: Apache-2.0 -- Static input shape: `[1, 3, 448, 448]` -- Output mask logits shape: `[1, 1, 112, 112]` -- Positive point prompt: `(224, 224)` in the padded `448x448` model input - frame -- Target: `ethos-u85-256` -- System config: `Ethos_U85_SYS_DRAM_Mid` -- Memory mode: `Dedicated_Sram_384KB` -- Calibration samples: `4` -- Validation samples: `4` - -The MobileSAM source and checkpoint are not redistributed by this example. The -preparation script clones the pinned official source into a managed external -cache and applies the configurable-input patch. The exporter downloads the -pinned official checkpoint and verifies its SHA256 unless `--checkpoint-path` -is provided. The source repository and checkpoint may have separate terms; the -export metadata records only the source-code license and does not assign a -license to the checkpoint. - -For a quick offline smoke test after caching the official checkpoint and source -checkout: - -```bash -python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py \ - --source-dir "$MOBILE_SAM_SOURCE" \ - --local-files-only - -python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py \ - --output-path /tmp/mobilesam_smoke.pte \ - --local-files-only \ - --mobile-sam-source "$MOBILE_SAM_SOURCE" \ - --calibration-image examples/models/dinov2/dog.jpg \ - --eval-image examples/models/dinov2/dog.jpg \ - --point 219 193 \ - --num-calibration-samples 1 \ - --num-eval-samples 1 \ - --minimum-fp32-quantized-iou 0.9 -``` - -Repeat `--point X Y` to freeze a multi-point prompt into the graph. Validate -every prompt set because adding positive points can substantially change the -predicted object after quantization: +From the repository root, run: ```bash -python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py \ - --output-path /tmp/mobilesam_multipoint_smoke.pte \ - --local-files-only \ - --mobile-sam-source "$MOBILE_SAM_SOURCE" \ - --calibration-image examples/models/dinov2/dog.jpg \ - --eval-image examples/models/dinov2/dog.jpg \ - --point 166 158 \ - --point 219 193 \ - --point 289 184 \ - --num-calibration-samples 1 \ - --num-eval-samples 1 \ - --debug-output-dir /tmp/mobilesam_multipoint_debug +python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py +python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py ``` -Pass `--input-size 1024` to reproduce the original MobileSAM resolution. Smaller -inputs export and run faster, but very small inputs such as `224` or `256` -usually produce lower-quality masks because the checkpoint was trained for -`1024x1024` images. - -To validate against a known binary mask, pass one `--eval-mask` per -`--eval-image`. Non-zero mask pixels are treated as foreground. When no -reference mask is provided, validation reports FP32/quantized mask agreement. -Use `--minimum-fp32-quantized-iou` in automated runs to reject an inaccurate -quantized model before lowering. - -The export flow: - -1. Loads the MobileSAM `vit_t` checkpoint through the patched external API. -2. Builds a fixed-prompt wrapper containing the image encoder and mask decoder. -3. Calibrates PT2E quantization with `EthosUQuantizer`. -4. Uses stable softmax decomposition for transformer attention blocks. -5. Lowers the quantized graph with `EthosUPartitioner`. -6. Writes an ExecuTorch `.pte` program. - -The quantization recipe uses int8 activations and int8 weights globally, with -A16W8 quantization for the TinyViT attention modules. Static int8 activation -quantization collapses MobileSAM attention features, while the selective A16W8 -attention path preserves mask quality and still lowers as one Ethos-U delegate. - -## Outputs - -For `--output-path ./mobilesam_point_ethos_u85_448.pte`, the script writes: - -- `mobilesam_point_ethos_u85_448.pte` - Ethos-U-ready ExecuTorch program. -- `mobilesam_point_ethos_u85_448.json` - Export metadata. -- `mobilesam_point_ethos_u85_448_metrics.json` - FP32/quantized mask - agreement and optional reference-mask IoU. -- `mobilesam_point_ethos_u85_448_delegation.txt` - Operator delegation - summary. -- `mobilesam_point_artifacts/` - Optional TOSA/Vela intermediate artifacts. -- `mobilesam_point_debug/` - Optional per-sample masks, overlays, mismatch - heatmaps, and mask summaries. - -## Interpreting the Debug Artifacts - -Each debug sample contains: - -- `input.png` - The resized RGB input used by the exported model. -- `reference_mask.png` - Optional binary reference mask resized to the model - output-mask size. -- `fp32_mask.png` - Host-side FP32 model prediction. -- `fp32_overlay.png` - Colored FP32 prediction blended over the input image. -- `quantized_mask.png` - Host-side PT2E quantized prediction before lowering. -- `quantized_overlay.png` - Colored quantized prediction blended over the input - image. -- `mismatch_heatmap.png` - Green for FP32/quantized agreement and red for - mismatch. -- `mask_summary.json` - Foreground/background pixel counts and FP32/quantized - IoU. +The first script prepares the pinned official MobileSAM source and verified +checkpoint in `~/.cache/executorch/mobilesam`. It applies the included patch +there to support the smaller input size. Neither source nor checkpoint is +copied into the ExecuTorch repository. -The runtime app thresholds mask logits on target, logs a mask hash and -foreground/background counts, and can dump the thresholded mask as RLE chunks. -Host-side debug masks are intentionally generated before lowering so users can -inspect quantization quality without needing target-side image output. +The second script performs the model flow directly: -## Limitations +1. Wrap MobileSAM with the fixed point prompt. +2. Export with `torch.export`. +3. Calibrate and convert with PT2E, using A8W8 generally and A16W8 attention + activations to preserve mask quality. +4. Check the FP32 and quantized masks have at least `0.9` IoU. +5. Lower with `EthosUPartitioner` and require one delegated subgraph. +6. Save `arm_test/mobilesam/export/mobilesam.pte`. -- The example freezes positive point prompts into the exported graph to keep - the target app to a single image input. Changing the prompt requires - re-exporting the `.pte`. -- The export uses `multimask_output=False` and does not include candidate-mask - argmax selection, upsampling, or thresholding in the graph. Keep those steps - in host-side or target-side post-processing when comparing mask quality. -- The runtime app logs a mask hash and foreground/background counts; it does - not render a color image on target. -- The first supported runtime target is Corstone-320/Ethos-U85-256. -- Reduced input sizes require the patch applied by `prepare_mobilesam.py`; the - MobileSAM source remains outside the ExecuTorch checkout. -- The default `448x448` input is the smallest size in the local sweep that - retained at least `0.95` host quantized/FP32 mask IoU on the demo image. - `512x512` retained about `0.975` IoU, while smaller sizes were inconsistent. -- MobileSAM PTQ is sensitive to the calibration set and the mask-logit - threshold. Inspect `*_metrics.json` and the debug overlays before treating - the quantized mask as an accuracy result. +The export directory also contains the input tensor, host masks, validation +metrics, delegation summary, and TOSA/Vela intermediate artifacts. diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py index e7c22518f00..c2fc992bfe7 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py @@ -3,25 +3,16 @@ # 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 hashlib -import importlib -import inspect import json -import os import sys -import urllib.error -import urllib.request -from dataclasses import dataclass -from importlib.metadata import PackageNotFoundError, version from pathlib import Path from typing import Any, cast +import executorch.kernels.quantized # noqa: F401 + import numpy as np import torch -import tqdm # type: ignore[import] +import torch.nn.functional as F from executorch.backends.arm.common.pipeline_config import ( ArmPassPipelineConfig, SoftmaxDecompositionConfig, @@ -39,856 +30,165 @@ to_edge_transform_and_lower, ) from executorch.extension.export_util.utils import save_pte_program -from packaging.version import Version from PIL import Image -from torchao.quantization.pt2e.quantize_pt2e import ( # type: ignore[import] - convert_pt2e, - prepare_pt2e, +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +ROOT = Path(__file__).resolve().parents[4] +WORK_DIR = ROOT / "arm_test" / "mobilesam" +SOURCE_DIR = ( + Path.home() + / ".cache" + / "executorch" + / "mobilesam" + / "f706ad9c4eb7f219c00d9050e46328518ffb65d2" + / "source" ) - -MOBILE_SAM_SOURCE_URL = "https://github.com/ChaoningZhang/MobileSAM" -MOBILE_SAM_SOURCE_REVISION = "f706ad9c4eb7f219c00d9050e46328518ffb65d2" -MOBILE_SAM_PATCH = "0001-Make-TinyViT-image-size-configurable.patch" -DEFAULT_CHECKPOINT_FILENAME = "mobile_sam.pt" -DEFAULT_CHECKPOINT_URL = ( - f"{MOBILE_SAM_SOURCE_URL}/raw/{MOBILE_SAM_SOURCE_REVISION}/weights/" - f"{DEFAULT_CHECKPOINT_FILENAME}" -) -DEFAULT_CHECKPOINT_SHA256 = ( - "6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f" -) -MOBILE_SAM_SOURCE_LICENSE = "Apache-2.0" -MINIMUM_VELA_VERSION = Version("5.1.0") -DEFAULT_INPUT_SIZE = 448 -MOBILE_SAM_INPUT_ALIGNMENT = 16 - - -@dataclass -class PreparedSample: - name: str - image: Image.Image - pixel_values: torch.Tensor - labels: torch.Tensor | None - - -def load_mobile_sam( - checkpoint_path: str, - mobile_sam_source: str | None, - input_size: int, -) -> torch.nn.Module: - mobile_sam_module = import_mobile_sam_module(mobile_sam_source) - builder = mobile_sam_module.sam_model_registry["vit_t"] - if "image_size" not in inspect.signature(builder).parameters: - raise RuntimeError( - "The MobileSAM checkout does not provide configurable image sizes. " - "Run model_export/prepare_mobilesam.py and pass the prepared checkout " - "with --mobile-sam-source." - ) - return builder(checkpoint=checkpoint_path, image_size=input_size).eval() +CHECKPOINT = SOURCE_DIR.parent / "mobile_sam.pt" +IMAGE = ROOT / "examples" / "models" / "dinov2" / "dog.jpg" +POINT = (219.0, 193.0) +INPUT_SIZE = 448 +MINIMUM_IOU = 0.9 class MobileSAMFixedPrompt(torch.nn.Module): - image_encoder: Any - mask_decoder: Any - - def __init__( - self, - sam: torch.nn.Module, - point_prompts: list[tuple[float, float]], - ) -> None: + def __init__(self, sam: torch.nn.Module) -> None: super().__init__() sam = cast(Any, sam) - self.image_encoder = sam.image_encoder self.mask_decoder = sam.mask_decoder with torch.no_grad(): points = ( - torch.tensor([point_prompts], dtype=torch.float32), - torch.ones((1, len(point_prompts)), dtype=torch.int64), - ) - sparse_embeddings, dense_embeddings = sam.prompt_encoder( - points=points, - boxes=None, - masks=None, + torch.tensor([[POINT]], dtype=torch.float32), + torch.ones((1, 1), dtype=torch.int64), ) + sparse, dense = sam.prompt_encoder(points=points, boxes=None, masks=None) image_pe = sam.prompt_encoder.get_dense_pe() - - self.register_buffer("sparse_prompt_embeddings", sparse_embeddings) - self.register_buffer("dense_prompt_embeddings", dense_embeddings) + self.register_buffer("sparse_prompt", sparse) + self.register_buffer("dense_prompt", dense) self.register_buffer("image_pe", image_pe) - def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: - image_embeddings = self.image_encoder(pixel_values) - low_res_masks, _ = self.mask_decoder( - image_embeddings=image_embeddings, + def forward(self, image: torch.Tensor) -> torch.Tensor: + masks, _ = self.mask_decoder( + image_embeddings=self.image_encoder(image), image_pe=self.image_pe, - sparse_prompt_embeddings=self.sparse_prompt_embeddings, - dense_prompt_embeddings=self.dense_prompt_embeddings, + sparse_prompt_embeddings=self.sparse_prompt, + dense_prompt_embeddings=self.dense_prompt, multimask_output=False, ) - return low_res_masks - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Export fixed-prompt MobileSAM segmentation for Ethos-U." - ) - parser.add_argument( - "--checkpoint-path", - default=None, - help="Optional local MobileSAM checkpoint path.", - ) - parser.add_argument( - "--mobile-sam-source", - default=None, - help=( - "Optional MobileSAM checkout prepared by prepare_mobilesam.py. " - "Used before importing the mobile_sam package." - ), - ) - parser.add_argument( - "--local-files-only", - action="store_true", - help="Load the checkpoint from the local cache only; do not download.", - ) - parser.add_argument( - "--calibration-image", - action="append", - default=[], - help="Local RGB image used for PTQ calibration. Can be repeated.", - ) - parser.add_argument( - "--eval-image", - action="append", - default=[], - help="Local RGB image used for validation/debugging. Can be repeated.", - ) - parser.add_argument( - "--eval-mask", - action="append", - default=[], - help=( - "Optional binary reference mask used for validation/debugging. " - "Can be repeated and must match --eval-image count when provided." - ), - ) - parser.add_argument( - "--point", - type=float, - nargs=2, - action="append", - default=[], - metavar=("X", "Y"), - help=( - "Positive point prompt in the resized square input frame. " - "Can be repeated for multi-point prompts." - ), - ) - parser.add_argument( - "--mask-threshold", - type=float, - default=0.0, - help="Mask-logit threshold used for metrics and debug masks.", - ) - parser.add_argument( - "--output-path", - type=str, - required=True, - help="Path to save the exported ExecuTorch program.", - ) - parser.add_argument( - "--input-size", - type=int, - default=DEFAULT_INPUT_SIZE, - help="Square MobileSAM input size. Must be divisible by 16.", - ) - parser.add_argument( - "--num-calibration-samples", - type=int, - default=4, - help="Number of local samples used for PTQ calibration.", - ) - parser.add_argument( - "--num-eval-samples", - type=int, - default=4, - help="Number of local samples used for host-side validation.", - ) - parser.add_argument( - "--num-debug-samples", - type=int, - default=4, - help="Number of validation samples written as visual debug artifacts.", - ) - parser.add_argument( - "--minimum-fp32-quantized-iou", - type=float, - default=None, - help="Fail before lowering when host quantized/FP32 mask IoU is lower.", - ) - parser.add_argument( - "--target", - default="ethos-u85-256", - help="Ethos-U target passed to Vela.", - ) - parser.add_argument( - "--system-config", - default="Ethos_U85_SYS_DRAM_Mid", - help="Vela system configuration.", - ) - parser.add_argument( - "--memory-mode", - default="Dedicated_Sram_384KB", - help="Vela memory mode.", - ) - parser.add_argument( - "--extra-vela-flag", - action="append", - default=[], - help="Additional Vela flag. Can be provided multiple times.", - ) - parser.add_argument( - "--artifact-dir", - default=None, - help="Optional directory for intermediate TOSA/Vela artifacts.", - ) - parser.add_argument( - "--debug-output-dir", - default=None, - help="Optional directory for masks, overlays, and validation summaries.", - ) - return parser.parse_args() - - -def write_json(path: Path, payload: dict[str, Any]) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") - - -def validate_vela_version() -> str: - try: - installed_version = version("ethos-u-vela") - except PackageNotFoundError as error: - raise RuntimeError( - "MobileSAM export requires ethos-u-vela 5.1.0 or newer. " - "Run examples/arm/setup.sh and retry." - ) from error - - if Version(installed_version) < MINIMUM_VELA_VERSION: - raise RuntimeError( - "MobileSAM A16W8 attention requires ethos-u-vela 5.1.0 or newer; " - f"found {installed_version}. Run examples/arm/setup.sh and retry." - ) - return installed_version + return masks -def import_mobile_sam_module(mobile_sam_source: str | None) -> Any: - if mobile_sam_source is not None: - sys.path.insert(0, str(Path(mobile_sam_source).expanduser().resolve())) - try: - return importlib.import_module("mobile_sam") - except ImportError as error: - raise ImportError( - "Could not import the patched mobile_sam package. Run " - "model_export/prepare_mobilesam.py and pass its checkout with " - "--mobile-sam-source." - ) from error +def load_model() -> torch.nn.Module: + if not SOURCE_DIR.exists() or not CHECKPOINT.exists(): + raise RuntimeError("Run prepare_mobilesam.py first.") + sys.path.insert(0, str(SOURCE_DIR)) + from mobile_sam import sam_model_registry # type: ignore[import-not-found] + return sam_model_registry["vit_t"]( + checkpoint=str(CHECKPOINT), image_size=INPUT_SIZE + ).eval() -def find_module_type(module: torch.nn.Module, class_name: str) -> type[torch.nn.Module]: - for child in module.modules(): - if child.__class__.__name__ == class_name: - return child.__class__ - raise ValueError(f"Could not find module type {class_name} in {module.__class__}.") +def prepare_image(sam: torch.nn.Module) -> tuple[Image.Image, torch.Tensor]: + image = Image.open(IMAGE).convert("RGB") + scale = INPUT_SIZE / max(image.size) + resized_size = (round(image.width * scale), round(image.height * scale)) + resized = image.resize(resized_size, Image.Resampling.BILINEAR) + padded = Image.new("RGB", (INPUT_SIZE, INPUT_SIZE)) + padded.paste(resized) -def file_sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as file: - for chunk in iter(lambda: file.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def default_checkpoint_cache_dir() -> Path: - return ( - Path.home() / ".cache" / "executorch" / "mobilesam" / MOBILE_SAM_SOURCE_REVISION + sam = cast(Any, sam) + mean = sam.pixel_mean.detach().cpu().reshape(3).numpy() + std = sam.pixel_std.detach().cpu().reshape(3).numpy() + tensor = torch.from_numpy((np.asarray(resized, dtype=np.float32) - mean) / std) + tensor = tensor.permute(2, 0, 1).unsqueeze(0) + tensor = F.pad( + tensor, (0, INPUT_SIZE - resized.width, 0, INPUT_SIZE - resized.height) ) + return padded, tensor.contiguous() -def verify_checkpoint(path: Path, expected_sha256: str) -> None: - actual_sha256 = file_sha256(path) - if actual_sha256 != expected_sha256: - raise RuntimeError( - f"Checkpoint SHA256 mismatch for {path}: expected {expected_sha256}, " - f"got {actual_sha256}." - ) - - -def download_checkpoint(url: str, path: Path) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - temp_path = path.with_suffix(path.suffix + ".tmp") - try: - with ( - urllib.request.urlopen(url, timeout=60) as response, # nosec B310 - temp_path.open("wb") as file, - ): - for chunk in iter(lambda: response.read(1024 * 1024), b""): - file.write(chunk) - temp_path.replace(path) - except (OSError, urllib.error.URLError) as error: - temp_path.unlink(missing_ok=True) - raise RuntimeError( - f"Failed to download MobileSAM checkpoint from {url}." - ) from error - - -def resolve_checkpoint( - args: argparse.Namespace, -) -> tuple[str, str | None, str | None]: - if args.checkpoint_path is not None: - checkpoint_path = Path(args.checkpoint_path).expanduser().resolve() - if not checkpoint_path.is_file(): - raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") - return str(checkpoint_path), None, None - - checkpoint_path = default_checkpoint_cache_dir() / DEFAULT_CHECKPOINT_FILENAME - if not checkpoint_path.is_file(): - if args.local_files_only: - raise FileNotFoundError( - f"Checkpoint not found in local cache: {checkpoint_path}" - ) - download_checkpoint(DEFAULT_CHECKPOINT_URL, checkpoint_path) - - verify_checkpoint(checkpoint_path, DEFAULT_CHECKPOINT_SHA256) - return ( - str(checkpoint_path.resolve()), - DEFAULT_CHECKPOINT_URL, - DEFAULT_CHECKPOINT_SHA256, - ) +def mask(logits: torch.Tensor) -> np.ndarray: + return (logits.detach().cpu().squeeze().numpy() > 0).astype(np.uint8) -def preprocess_image( - image: Image.Image, - segmentation_map: Image.Image | None, - *, - name: str, - input_size: int, - pixel_mean: np.ndarray, - pixel_std: np.ndarray, - output_mask_size: tuple[int, int] | None, -) -> PreparedSample: - rgb_image = image.convert("RGB") - width, height = rgb_image.size - scale = input_size / max(height, width) - resized_size = (round(width * scale), round(height * scale)) - resized_image = rgb_image.resize(resized_size, Image.Resampling.BILINEAR) - - padded_image = Image.new("RGB", (input_size, input_size)) - padded_image.paste(resized_image, (0, 0)) - - image_np = np.asarray(resized_image, dtype=np.float32) - image_np = (image_np - pixel_mean) / pixel_std - pixel_values = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(0) - pixel_values = torch.nn.functional.pad( - pixel_values, - ( - 0, - input_size - resized_size[0], - 0, - input_size - resized_size[1], - ), - ) - pixel_values = pixel_values.contiguous() - - labels = None - if segmentation_map is not None: - if output_mask_size is None: - raise ValueError("An output mask size is required for reference masks.") - resized_mask = segmentation_map.convert("L").resize( - resized_size, - Image.Resampling.NEAREST, - ) - padded_mask = Image.new("L", (input_size, input_size)) - padded_mask.paste(resized_mask, (0, 0)) - resized_mask = padded_mask.resize(output_mask_size, Image.Resampling.NEAREST) - mask_np = np.asarray(resized_mask, dtype=np.uint8) - labels = torch.from_numpy((mask_np > 0).astype(np.uint8)).to(torch.long) - - return PreparedSample( - name=name, - image=padded_image, - pixel_values=pixel_values, - labels=labels, - ) +def iou(first: np.ndarray, second: np.ndarray) -> float: + intersection = np.logical_and(first, second).sum() + union = np.logical_or(first, second).sum() + return 1.0 if union == 0 else float(intersection / union) -def load_local_samples( - image_paths: list[str], - mask_paths: list[str], - limit: int, - *, - include_labels: bool, - input_size: int, - pixel_mean: np.ndarray, - pixel_std: np.ndarray, - output_mask_size: tuple[int, int] | None = None, -) -> list[PreparedSample]: - if include_labels and len(mask_paths) not in (0, len(image_paths)): - raise ValueError("--eval-mask must be omitted or match --eval-image count.") - - samples: list[PreparedSample] = [] - for index, image_path in enumerate(image_paths): - if len(samples) >= limit: - break - image = Image.open(image_path) - mask = None - if include_labels and len(mask_paths) > 0: - mask = Image.open(mask_paths[index]) - samples.append( - preprocess_image( - image, - mask, - name=Path(image_path).stem or f"local_sample_{index:04d}", - input_size=input_size, - pixel_mean=pixel_mean, - pixel_std=pixel_std, - output_mask_size=output_mask_size, - ) - ) - if len(samples) == 0: - raise ValueError("No local samples were loaded.") - return samples - - -def run_mask_logits(model: torch.nn.Module, input_tensor: torch.Tensor) -> torch.Tensor: - output = model(input_tensor) - if isinstance(output, (tuple, list)): - output = output[0] - if not isinstance(output, torch.Tensor): - raise TypeError(f"Expected tensor mask logits, got {type(output)}") - return output - - -def predict_mask(logits: torch.Tensor, threshold: float) -> np.ndarray: - mask = logits.detach().cpu().squeeze(0).squeeze(0).numpy() > threshold - return mask.astype(np.uint8) - - -def binary_iou(mask_a: np.ndarray, mask_b: np.ndarray) -> float: - intersection = np.logical_and(mask_a, mask_b).sum() - union = np.logical_or(mask_a, mask_b).sum() - if union == 0: - return 1.0 - return float(intersection / union) - - -def save_binary_mask(path: Path, mask: np.ndarray) -> Image.Image: - image = Image.fromarray((mask.astype(np.uint8) * 255), mode="L") - image.save(path) - return image.convert("RGB") - - -def save_mask_overlay( - path: Path, - image: Image.Image, - mask: np.ndarray, - color: tuple[int, int, int], -) -> None: - rgb_image = image.convert("RGB") - resized_mask = Image.fromarray(mask.astype(np.uint8) * 255, mode="L").resize( - rgb_image.size, - Image.Resampling.NEAREST, - ) - mask_np = np.asarray(resized_mask, dtype=np.uint8) > 0 - overlay_np = np.asarray(rgb_image, dtype=np.float32) - color_np = np.asarray(color, dtype=np.float32) - overlay_np[mask_np] = overlay_np[mask_np] * 0.55 + color_np * 0.45 - Image.fromarray(np.clip(overlay_np, 0, 255).astype(np.uint8), mode="RGB").save(path) - - -def write_debug_artifacts( - debug_dir: Path, - sample: PreparedSample, - fp32_mask: np.ndarray, - quantized_mask: np.ndarray, -) -> None: - sample_dir = debug_dir / sample.name - sample_dir.mkdir(parents=True, exist_ok=True) - - sample.image.save(sample_dir / "input.png") - save_binary_mask(sample_dir / "fp32_mask.png", fp32_mask) - save_binary_mask(sample_dir / "quantized_mask.png", quantized_mask) - save_mask_overlay( - sample_dir / "fp32_overlay.png", - sample.image, - fp32_mask, - (0, 220, 120), - ) - save_mask_overlay( - sample_dir / "quantized_overlay.png", - sample.image, - quantized_mask, - (0, 170, 255), - ) - if sample.labels is not None: - save_binary_mask( - sample_dir / "reference_mask.png", - sample.labels.detach().cpu().numpy().astype(np.uint8), - ) - - mismatch = fp32_mask != quantized_mask - heatmap = np.zeros((*fp32_mask.shape, 3), dtype=np.uint8) - heatmap[~mismatch] = [0, 128, 0] - heatmap[mismatch] = [255, 0, 0] - Image.fromarray(heatmap, mode="RGB").save(sample_dir / "mismatch_heatmap.png") - - write_json( - sample_dir / "mask_summary.json", - { - "foreground_pixels": int(quantized_mask.sum()), - "background_pixels": int(quantized_mask.size - quantized_mask.sum()), - "fp32_quantized_iou": binary_iou(fp32_mask, quantized_mask), - }, - ) - - -def evaluate_and_debug( - fp32_model: torch.nn.Module, - quantized_model: torch.nn.Module, - eval_samples: list[PreparedSample], - debug_dir: Path | None, - num_debug_samples: int, - threshold: float, -) -> dict[str, float]: - fp32_quantized_ious: list[float] = [] - fp32_quantized_pixel_agreements: list[float] = [] - reference_ious: list[float] = [] - - if debug_dir is not None: - debug_dir.mkdir(parents=True, exist_ok=True) - - print("\nEvaluating quantized MobileSAM on validation samples...") - for index, sample in enumerate(tqdm.tqdm(eval_samples)): - fp32_logits = run_mask_logits(fp32_model, sample.pixel_values) - quantized_logits = run_mask_logits(quantized_model, sample.pixel_values) - fp32_mask = predict_mask(fp32_logits, threshold) - quantized_mask = predict_mask(quantized_logits, threshold) - - fp32_quantized_ious.append(binary_iou(fp32_mask, quantized_mask)) - fp32_quantized_pixel_agreements.append( - float(np.mean(fp32_mask == quantized_mask)) - ) - if sample.labels is not None: - labels = sample.labels.detach().cpu().numpy().astype(np.uint8) - reference_ious.append(binary_iou(quantized_mask, labels)) - - if debug_dir is not None and index < num_debug_samples: - write_debug_artifacts(debug_dir, sample, fp32_mask, quantized_mask) - - metrics = { - "num_samples": float(len(eval_samples)), - "fp32_quantized_mean_iou": float(np.mean(fp32_quantized_ious)), - "fp32_quantized_pixel_agreement": float( - np.mean(fp32_quantized_pixel_agreements) - ), - } - if len(reference_ious) > 0: - metrics["reference_mean_iou"] = float(np.mean(reference_ious)) - return metrics - - -def quantize_model( - model: torch.nn.Module, - quantizer: EthosUQuantizer, - example_inputs: tuple[torch.Tensor], - calibration_samples: list[PreparedSample], +def quantize( + model: torch.nn.Module, image: torch.Tensor, quantizer: EthosUQuantizer ) -> torch.export.ExportedProgram: - exported = torch.export.export(model, example_inputs) + exported = torch.export.export(model, (image,)) prepared = prepare_pt2e(exported.module(), quantizer) - - print("\nCalibrating MobileSAM...") - for sample in tqdm.tqdm(calibration_samples): - prepared(sample.pixel_values) - - quantized = convert_pt2e(prepared) - return torch.export.export(quantized, example_inputs) - - -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 load_quantized_ops_library(library_path: Path) -> Path: - if not library_path.is_file(): - raise FileNotFoundError(f"Quantized ops library not found: {library_path}") - torch.ops.load_library(str(library_path)) - if has_quantized_out_variants(): - return library_path - raise RuntimeError( - f"Quantized ops library did not register required out variants: {library_path}" - ) - - -def ensure_quantized_ops_loaded() -> Path | None: - if has_quantized_out_variants(): - return None - - quantized_ops_library = os.environ.get("EXECUTORCH_QUANTIZED_OPS_AOT_LIBRARY") - if quantized_ops_library: - return load_quantized_ops_library( - Path(quantized_ops_library).expanduser().resolve() - ) - - try: - import executorch.kernels.quantized # noqa: F401 - except ImportError: - pass - else: - if 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.*", - "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 - return load_quantized_ops_library(candidate) - - raise RuntimeError( - "MobileSAM 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 write_delegation_report(edge_program_manager: Any, 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 resolve_point_prompts(args: argparse.Namespace) -> list[tuple[float, float]]: - if len(args.point) > 0: - point_prompts = [(float(x), float(y)) for x, y in args.point] - else: - point_prompts = [(args.input_size / 2, args.input_size / 2)] - - for point_x, point_y in point_prompts: - if not (0 <= point_x <= args.input_size and 0 <= point_y <= args.input_size): - raise ValueError("Point prompts must be inside the square input.") - return point_prompts - - -def validate_export_args(args: argparse.Namespace) -> None: - if args.input_size < 224 or args.input_size % MOBILE_SAM_INPUT_ALIGNMENT != 0: - raise ValueError("--input-size must be at least 224 and divisible by 16.") - if args.num_calibration_samples <= 0: - raise ValueError("--num-calibration-samples must be positive.") - if args.num_eval_samples <= 0: - raise ValueError("--num-eval-samples must be positive.") - if args.minimum_fp32_quantized_iou is not None and not ( - 0.0 <= args.minimum_fp32_quantized_iou <= 1.0 - ): - raise ValueError("--minimum-fp32-quantized-iou must be between 0 and 1.") - if len(args.calibration_image) == 0: - raise ValueError("At least one --calibration-image is required.") - if len(args.eval_image) == 0: - args.eval_image = list(args.calibration_image) - if len(args.eval_mask) not in (0, len(args.eval_image)): - raise ValueError("--eval-mask must be omitted or match --eval-image count.") + prepared(image) + return torch.export.export(convert_pt2e(prepared), (image,)) def main() -> None: - args = parse_args() - validate_export_args(args) - vela_version = validate_vela_version() - point_prompts = resolve_point_prompts(args) - quantized_ops_library = ensure_quantized_ops_loaded() - if quantized_ops_library is not None: - print(f"Loaded quantized ops library from {quantized_ops_library}") - - output_path = Path(args.output_path).resolve() - output_path.parent.mkdir(parents=True, exist_ok=True) - metadata_path = output_path.with_suffix(".json") - metrics_path = output_path.with_name(f"{output_path.stem}_metrics.json") - delegation_path = output_path.with_name(f"{output_path.stem}_delegation.txt") - debug_dir = Path(args.debug_output_dir).resolve() if args.debug_output_dir else None - - checkpoint_path, checkpoint_url, checkpoint_sha256 = resolve_checkpoint(args) - mobile_sam = load_mobile_sam( - checkpoint_path, - args.mobile_sam_source, - args.input_size, - ) - pixel_mean = cast(Any, mobile_sam).pixel_mean.detach().cpu().reshape(-1).numpy() - pixel_std = cast(Any, mobile_sam).pixel_std.detach().cpu().reshape(-1).numpy() - if pixel_mean.shape != (3,) or pixel_std.shape != (3,): - raise ValueError("MobileSAM preprocessing must provide three RGB values.") - wrapped_model = MobileSAMFixedPrompt(mobile_sam, point_prompts).eval() - - calibration_samples = load_local_samples( - args.calibration_image, - [], - args.num_calibration_samples, - include_labels=False, - input_size=args.input_size, - pixel_mean=pixel_mean, - pixel_std=pixel_std, - ) - example_inputs = (calibration_samples[0].pixel_values,) - with torch.no_grad(): - output_shape = list(run_mask_logits(wrapped_model, example_inputs[0]).shape) - if len(output_shape) != 4 or output_shape[:2] != [1, 1]: - raise ValueError( - f"Expected MobileSAM output shape [1, 1, height, width], got {output_shape}." - ) - output_mask_size = (output_shape[3], output_shape[2]) - - eval_samples = load_local_samples( - args.eval_image, - args.eval_mask, - args.num_eval_samples, - include_labels=True, - input_size=args.input_size, - pixel_mean=pixel_mean, - pixel_std=pixel_std, - output_mask_size=output_mask_size, - ) + export_dir = WORK_DIR / "export" + export_dir.mkdir(parents=True, exist_ok=True) + + sam = load_model() + input_image, example_input = prepare_image(sam) + model = MobileSAMFixedPrompt(sam).eval() compile_spec = EthosUCompileSpec( - target=args.target, - system_config=args.system_config, - memory_mode=args.memory_mode, - extra_flags=args.extra_vela_flag, + "ethos-u85-256", memory_mode="Dedicated_Sram_384KB" ) compile_spec.set_pass_pipeline_config( ArmPassPipelineConfig(softmax=SoftmaxDecompositionConfig.STABLE) ) - 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)) + compile_spec.dump_intermediate_artifacts_to(str(export_dir / "artifacts")) quantizer = EthosUQuantizer(compile_spec) quantizer.set_global(get_symmetric_quantization_config()) - attention_module_type = find_module_type(wrapped_model.image_encoder, "Attention") - quantizer.set_module_type( - attention_module_type, - get_symmetric_a16w8_quantization_config(), + attention_type = next( + type(module) + for module in model.image_encoder.modules() + if type(module).__name__ == "Attention" ) + # Int16 attention activations preserve the segmentation mask quality. + quantizer.set_module_type(attention_type, get_symmetric_a16w8_quantization_config()) with torch.no_grad(): - quantized_program = quantize_model( - wrapped_model, - quantizer, - example_inputs, - calibration_samples, - ) - quantized_module = quantized_program.module() - metrics = evaluate_and_debug( - wrapped_model, - quantized_module, - eval_samples, - debug_dir, - args.num_debug_samples, - args.mask_threshold, - ) - write_json(metrics_path, metrics) - print( - "Validation metrics: " - f"fp32_quantized_mean_iou={metrics['fp32_quantized_mean_iou']:.4f} " - "fp32_quantized_pixel_agreement=" - f"{metrics['fp32_quantized_pixel_agreement']:.4f}" - ) - if ( - args.minimum_fp32_quantized_iou is not None - and metrics["fp32_quantized_mean_iou"] < args.minimum_fp32_quantized_iou - ): - raise RuntimeError( - "Host quantized/FP32 mask IoU " - f"{metrics['fp32_quantized_mean_iou']:.4f} is below " - f"{args.minimum_fp32_quantized_iou:.4f}." - ) + fp32_mask = mask(model(example_input)) + quantized = quantize(model, example_input, quantizer) + quantized_mask = mask(quantized.module()(example_input)) + + host_iou = iou(fp32_mask, quantized_mask) + if host_iou < MINIMUM_IOU: + raise RuntimeError(f"FP32/quantized mask IoU is too low: {host_iou:.4f}") - edge_program_manager = to_edge_transform_and_lower( - programs=quantized_program, + edge = to_edge_transform_and_lower( + quantized, partitioner=[EthosUPartitioner(compile_spec)], compile_config=EdgeCompileConfig(_check_ir_validity=False), ) - write_delegation_report(edge_program_manager, delegation_path) + delegation = get_delegation_info(edge.exported_program().graph_module) + if delegation.num_delegated_subgraphs != 1: + raise RuntimeError("Expected one Ethos-U delegate.") - executorch_program_manager = edge_program_manager.to_executorch( + program = edge.to_executorch( config=ExecutorchBackendConfig(extract_delegate_segments=False) ) - save_pte_program( - executorch_program_manager, - str(output_path), - output_dir=str(output_path.parent), - ) + save_pte_program(program, "mobilesam", output_dir=str(export_dir)) - write_json( - metadata_path, - { - "model_name": "MobileSAM vit_t", - "checkpoint_filename": DEFAULT_CHECKPOINT_FILENAME, - "checkpoint_path": checkpoint_path, - "checkpoint_url": checkpoint_url, - "checkpoint_sha256": checkpoint_sha256, - "mobile_sam_source_license": MOBILE_SAM_SOURCE_LICENSE, - "mobile_sam_source_url": MOBILE_SAM_SOURCE_URL, - "mobile_sam_source_revision": MOBILE_SAM_SOURCE_REVISION, - "mobile_sam_patch": MOBILE_SAM_PATCH, - "input_shape": list(example_inputs[0].shape), - "output_shape": output_shape, - "input_size": args.input_size, - "preprocessing": { - "pixel_mean": pixel_mean.tolist(), - "pixel_std": pixel_std.tolist(), - "resize": "longest_side_then_zero_pad", - }, - "point_prompts_xy": point_prompts, - "mask_threshold": args.mask_threshold, - "target": args.target, - "vela_version": vela_version, - "system_config": args.system_config, - "memory_mode": args.memory_mode, - "extra_vela_flags": args.extra_vela_flag, - "quantization": { - "global": "int8 activations and int8 weights", - "tinyvit_attention": "int16 activations and int8 weights", - }, - "num_calibration_samples": len(calibration_samples), - "num_eval_samples": len(eval_samples), - "calibration_images": args.calibration_image, - "eval_images": args.eval_image, - "eval_masks": args.eval_mask, - "output_path": str(output_path), - "metrics_path": str(metrics_path), - "delegation_path": str(delegation_path), - "debug_output_dir": str(debug_dir) if debug_dir is not None else None, - }, + input_image.save(export_dir / "input.png") + Image.fromarray(fp32_mask * 255).save(export_dir / "fp32_mask.png") + Image.fromarray(quantized_mask * 255).save(export_dir / "quantized_mask.png") + example_input.numpy().astype(np.float32).tofile(export_dir / "input.bin") + (export_dir / "delegation.txt").write_text(delegation.get_summary() + "\n") + (export_dir / "metrics.json").write_text( + json.dumps({"fp32_quantized_iou": host_iou}, indent=2) + "\n" ) - print(f"\nExported model saved to {output_path}") - print(f"Metadata saved to {metadata_path}") - print(f"Metrics saved to {metrics_path}") - print(f"Delegation summary saved to {delegation_path}") - if debug_dir is not None: - print(f"Debug artifacts saved to {debug_dir}") + print(f"FP32/quantized mask IoU: {host_iou:.4f}") + print(f"Saved {export_dir / 'mobilesam.pte'}") if __name__ == "__main__": diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py index 73dd6119c39..bcb922ca052 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py @@ -3,119 +3,86 @@ # 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 hashlib import subprocess # nosec B404 +import urllib.request from pathlib import Path -MOBILE_SAM_SOURCE_URL = "https://github.com/ChaoningZhang/MobileSAM.git" -MOBILE_SAM_SOURCE_REVISION = "f706ad9c4eb7f219c00d9050e46328518ffb65d2" -PATCH_DIR = Path(__file__).resolve().parent / "patches" / "mobile_sam" - - -def run(command: list[str], *, cwd: Path | None = None) -> None: +REVISION = "f706ad9c4eb7f219c00d9050e46328518ffb65d2" +SOURCE_URL = "https://github.com/ChaoningZhang/MobileSAM.git" +CHECKPOINT_URL = ( + f"https://github.com/ChaoningZhang/MobileSAM/raw/{REVISION}/weights/mobile_sam.pt" +) +CHECKPOINT_SHA256 = "6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f" +CACHE_DIR = Path.home() / ".cache" / "executorch" / "mobilesam" / REVISION +SOURCE_DIR = CACHE_DIR / "source" +CHECKPOINT = CACHE_DIR / "mobile_sam.pt" +PATCH = ( + Path(__file__).parent + / "patches" + / "mobile_sam" + / ("0001-Make-TinyViT-image-size-configurable.patch") +) + + +def run(*command: str, cwd: Path | None = None) -> None: subprocess.run(command, cwd=cwd, check=True) # nosec B603 -def default_source_dir() -> Path: - return ( - Path.home() - / ".cache" - / "executorch" - / "mobilesam" - / MOBILE_SAM_SOURCE_REVISION - / "source" - ) - +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() -def prepare_source(source_dir: Path, *, local_files_only: bool) -> None: - source_dir = source_dir.expanduser().resolve() - marker = source_dir.parent / f".{source_dir.name}.executorch-managed" - if source_dir.exists() and not marker.is_file(): - raise RuntimeError( - f"Refusing to modify unmanaged MobileSAM directory: {source_dir}" - ) +def prepare_source() -> None: + marker = SOURCE_DIR.parent / ".source.executorch-managed" + if SOURCE_DIR.exists() and not marker.exists(): + raise RuntimeError(f"Refusing to modify unmanaged directory: {SOURCE_DIR}") - if not source_dir.exists(): - if local_files_only: - raise FileNotFoundError( - f"Managed MobileSAM checkout not found: {source_dir}" - ) - source_dir.parent.mkdir(parents=True, exist_ok=True) - run( - [ - "git", - "clone", - "--filter=blob:none", - "--no-checkout", - MOBILE_SAM_SOURCE_URL, - str(source_dir), - ] - ) - marker.write_text(MOBILE_SAM_SOURCE_REVISION + "\n") + if not SOURCE_DIR.exists(): + SOURCE_DIR.parent.mkdir(parents=True, exist_ok=True) run( - [ - "git", - "sparse-checkout", - "set", - "mobile_sam", - ], - cwd=source_dir, + "git", + "clone", + "--filter=blob:none", + "--no-checkout", + SOURCE_URL, + str(SOURCE_DIR), ) - - if not local_files_only: - run( - ["git", "fetch", "--quiet", "origin", MOBILE_SAM_SOURCE_REVISION], - cwd=source_dir, - ) - - try: - run( - ["git", "cat-file", "-e", f"{MOBILE_SAM_SOURCE_REVISION}^{{commit}}"], - cwd=source_dir, - ) - except subprocess.CalledProcessError as error: + marker.write_text(REVISION + "\n") + run("git", "sparse-checkout", "set", "mobile_sam", cwd=SOURCE_DIR) + + run("git", "fetch", "--quiet", "origin", REVISION, cwd=SOURCE_DIR) + run("git", "checkout", "--detach", "--force", REVISION, cwd=SOURCE_DIR) + run("git", "reset", "--hard", REVISION, cwd=SOURCE_DIR) + run("git", "apply", str(PATCH), cwd=SOURCE_DIR) + + +def prepare_checkpoint() -> None: + if not CHECKPOINT.exists(): + CHECKPOINT.parent.mkdir(parents=True, exist_ok=True) + with ( + urllib.request.urlopen( + CHECKPOINT_URL, timeout=60 + ) as response, # nosec B310 + CHECKPOINT.open("wb") as file, + ): + while chunk := response.read(1024 * 1024): + file.write(chunk) + + actual_sha256 = sha256(CHECKPOINT) + if actual_sha256 != CHECKPOINT_SHA256: raise RuntimeError( - f"MobileSAM revision {MOBILE_SAM_SOURCE_REVISION} is unavailable locally." - ) from error - - run( - ["git", "checkout", "--detach", "--force", MOBILE_SAM_SOURCE_REVISION], - cwd=source_dir, - ) - run(["git", "reset", "--hard", MOBILE_SAM_SOURCE_REVISION], cwd=source_dir) - - patches = sorted(PATCH_DIR.glob("*.patch")) - if not patches: - raise FileNotFoundError(f"No MobileSAM patches found in {PATCH_DIR}") - for patch in patches: - run(["git", "apply", "--check", str(patch)], cwd=source_dir) - run(["git", "apply", str(patch)], cwd=source_dir) - - print(f"Prepared patched MobileSAM source at {source_dir}") - - -def main() -> None: - parser = argparse.ArgumentParser( - description="Prepare the pinned MobileSAM source with ExecuTorch patches." - ) - parser.add_argument( - "--source-dir", - type=Path, - default=default_source_dir(), - help="Managed checkout destination.", - ) - parser.add_argument( - "--local-files-only", - action="store_true", - help="Reuse an existing managed checkout without network access.", - ) - args = parser.parse_args() - prepare_source(args.source_dir, local_files_only=args.local_files_only) + f"Checkpoint SHA256 mismatch: expected {CHECKPOINT_SHA256}, " + f"got {actual_sha256}" + ) if __name__ == "__main__": - main() + prepare_source() + prepare_checkpoint() + print(f"MobileSAM ready in {CACHE_DIR}") diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/run.sh b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/run.sh new file mode 100755 index 00000000000..511284d2ba7 --- /dev/null +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/run.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +# 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 + +if (($#)); then + echo "This example has one supported configuration; run it without arguments." + exit 2 +fi + +example_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +repo_root=$(cd "${example_dir}/../../.." && pwd) +work_dir="${repo_root}/arm_test/mobilesam" +export_dir="${work_dir}/export" +io_dir="${work_dir}/io" +runner_dir="${work_dir}/runner" + +cd "${repo_root}" +[[ -f examples/arm/arm-scratch/setup_path.sh ]] || { + echo "Arm tools are missing. Run ./examples/arm/setup.sh first." + exit 1 +} +source examples/arm/arm-scratch/setup_path.sh +mkdir -p "${io_dir}" + +if [[ "$(uname -s)" == "Darwin" ]]; then + export FVP_MOUNT_DIR="${FVP_MOUNT_DIR:-${repo_root}}" + export FVP_WORKDIR="${FVP_WORKDIR:-${repo_root}}" +fi + +echo "[1/4] Prepare MobileSAM" +python3 "${example_dir}/model_export/prepare_mobilesam.py" + +echo "[2/4] Export, quantize, and lower to Ethos-U" +python3 "${example_dir}/model_export/export_mobilesam.py" + +echo "[3/4] Build the standard Arm executor runner" +backends/arm/scripts/build_executor_runner.sh \ + --pte="${export_dir}/mobilesam.pte" \ + --target=ethos-u85-256 \ + --output="${runner_dir}" \ + '--extra_build_flags=-DSEMIHOSTING=ON -DET_COMPILED_PTE=ON' + +cp "${export_dir}/input.bin" "${io_dir}/input.bin" +rm -f "${io_dir}/output-0.bin" + +echo "[4/4] Run on FVP and validate the output" +backends/arm/scripts/run_fvp.sh \ + --elf="${runner_dir}/arm_executor_runner" \ + --target=ethos-u85-256 \ + --timeout=300 \ + --semihosting-cwd="${io_dir}" \ + '--semihosting-cmd-line=executor_runner -i input.bin -o output' \ + --fast | tee "${work_dir}/fvp.log" +python3 "${example_dir}/runtime/visualize_fvp_output.py" + +echo "MobileSAM example: PASS" diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/CMakeLists.txt b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/CMakeLists.txt deleted file mode 100644 index 6f30110b95c..00000000000 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/CMakeLists.txt +++ /dev/null @@ -1,223 +0,0 @@ -# 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(mobilesam_prompt_segmentation_ethos_u_application) - -set(ET_DIR_PATH - "${CMAKE_CURRENT_SOURCE_DIR}/../../../.." - CACHE PATH "Path to ExecuTorch dir" -) -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(MODEL_METADATA_PATH - "" - CACHE PATH "Path to MobileSAM exporter metadata" -) -set(IMAGE_PATH - "" - CACHE PATH "Path to an RGB image to use for the application" -) -set(MASK_THRESHOLD - "0.0" - CACHE STRING "MobileSAM mask logit threshold" -) -set(SYSTEM_CONFIG - "Ethos_U85_SYS_DRAM_Mid" - CACHE STRING "Vela system configuration" -) -set(MEMORY_MODE - "Dedicated_Sram_384KB" - CACHE STRING "Vela memory mode" -) -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" -) -option(ET_SEGMENTATION_DUMP_MASK - "Dump the predicted segmentation mask as run-length encoded log chunks" - OFF -) -option(ET_SEGMENTATION_SEMIHOSTING_OUTPUT - "Emit validation logs through Arm semihosting for FVP smoke runs" ON -) - -if(NOT EXISTS "${IMAGE_PATH}") - message( - FATAL_ERROR - "Image not provided. Please provide -DIMAGE_PATH= and retry." - ) -endif() -if(NOT EXISTS "${ET_PTE_FILE_PATH}") - message( - FATAL_ERROR - "PTE file not provided. Please provide -DET_PTE_FILE_PATH= and retry." - ) -endif() -if(NOT MODEL_METADATA_PATH) - get_filename_component(model_directory "${ET_PTE_FILE_PATH}" DIRECTORY) - get_filename_component(model_name "${ET_PTE_FILE_PATH}" NAME_WE) - set(MODEL_METADATA_PATH "${model_directory}/${model_name}.json") -endif() -if(NOT EXISTS "${MODEL_METADATA_PATH}") - message( - FATAL_ERROR - "Model metadata not found. Provide -DMODEL_METADATA_PATH=." - ) -endif() -if(NOT SYSTEM_CONFIG MATCHES "Ethos_U85") - message(FATAL_ERROR "This example currently supports Corstone-320/Ethos-U85.") -endif() - -include(${ET_DIR_PATH}/backends/arm/scripts/corstone_utils.cmake) -fetch_ethos_u_content(${ETHOS_SDK_PATH} ${ET_DIR_PATH}) - -if(NOT EXISTS "${ETHOS_SDK_PATH}") - message( - FATAL_ERROR - "The ${ETHOS_SDK_PATH} directory does not exist. Please run examples/arm/setup.sh and retry." - ) -endif() - -find_package( - executorch REQUIRED HINTS "${ET_BUILD_DIR_PATH}/lib/cmake/ExecuTorch" -) - -add_corstone_subdirectory(${SYSTEM_CONFIG} ${ETHOS_SDK_PATH}) -configure_timing_adapters(${SYSTEM_CONFIG} ${MEMORY_MODE}) - -add_executable(mobilesam_prompt_segmentation_example main.cpp) -target_sources( - mobilesam_prompt_segmentation_example - PRIVATE main.cpp ${ET_DIR_PATH}/examples/arm/common/arm_memory_allocator.cpp -) -target_link_libraries( - mobilesam_prompt_segmentation_example - PUBLIC executorch - ethosu_target_init - extension_runner_util - quantized_ops_lib - portable_kernels - cortex_m_kernels - cortex_m_ops_lib -) - -include(${ET_DIR_PATH}/tools/cmake/Utils.cmake) -executorch_target_link_options_shared_lib(executorch_delegate_ethos_u) -target_link_libraries( - mobilesam_prompt_segmentation_example PUBLIC executorch_delegate_ethos_u -) - -if(MEMORY_MODE MATCHES "^Dedicated_Sram($|_)") - set(ETHOSU_ARENA "1") - if(NOT DEFINED ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE) - set(ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE 0x1000000) - endif() - if(NOT DEFINED ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE) - set_ethosu_dedicated_sram_fast_scratch_size( - ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE "${MEMORY_MODE}" - ) - endif() -else() - set(ETHOSU_ARENA "0") - if(NOT DEFINED ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE) - set(ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE 0x400000) - endif() -endif() -if(NOT DEFINED ET_SEGMENTATION_METHOD_ALLOCATOR_POOL_SIZE) - set(ET_SEGMENTATION_METHOD_ALLOCATOR_POOL_SIZE 0x4000000) -endif() - -target_compile_definitions( - mobilesam_prompt_segmentation_example - PRIVATE - ET_SEGMENTATION_METHOD_ALLOCATOR_POOL_SIZE=${ET_SEGMENTATION_METHOD_ALLOCATOR_POOL_SIZE} - ET_SEGMENTATION_MASK_THRESHOLD=${MASK_THRESHOLD} - ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE=${ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE} -) -if(DEFINED ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE) - target_compile_definitions( - mobilesam_prompt_segmentation_example - PRIVATE - ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE=${ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE} - ) -endif() -if(ET_SEGMENTATION_DUMP_MASK) - target_compile_definitions( - mobilesam_prompt_segmentation_example PRIVATE ET_SEGMENTATION_DUMP_MASK - ) -endif() -if(ET_SEGMENTATION_SEMIHOSTING_OUTPUT) - target_compile_definitions( - mobilesam_prompt_segmentation_example - PRIVATE ET_SEGMENTATION_SEMIHOSTING_OUTPUT - ) -endif() - -if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") - set(LINK_FILE_EXT ld) - set(COMPILER_PREPROCESSOR_OPTIONS -E -x c -P) -endif() - -set(LINK_FILE_OUT_BASE "platform_linker_script") -set(LINK_FILE_IN - "${ET_DIR_PATH}/backends/arm/cmake/linker_scripts/Corstone-320.ld" -) -set(LINK_FILE_OUT - ${CMAKE_CURRENT_BINARY_DIR}/${LINK_FILE_OUT_BASE}.${LINK_FILE_EXT} -) -execute_process( - COMMAND ${CMAKE_C_COMPILER} ${COMPILER_PREPROCESSOR_OPTIONS} -DETHOSU_MODEL=1 - -DETHOSU_ARENA=${ETHOSU_ARENA} -o ${LINK_FILE_OUT} ${LINK_FILE_IN} -) -target_link_options( - mobilesam_prompt_segmentation_example PRIVATE "-T" "${LINK_FILE_OUT}" -) - -set(MODEL_PTE_HEADER "${CMAKE_CURRENT_BINARY_DIR}/model_pte.h") -add_custom_command( - OUTPUT "${MODEL_PTE_HEADER}" - COMMAND - ${PYTHON_EXECUTABLE} ${ET_DIR_PATH}/examples/arm/common/pte_to_header.py - --pte ${ET_PTE_FILE_PATH} --outdir ${CMAKE_CURRENT_BINARY_DIR} - DEPENDS ${ET_PTE_FILE_PATH} - ${ET_DIR_PATH}/examples/arm/common/pte_to_header.py - VERBATIM -) -set(IMAGE_HEADER "${CMAKE_CURRENT_BINARY_DIR}/image.h") -add_custom_command( - OUTPUT "${IMAGE_HEADER}" - COMMAND - ${PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/image_to_array.py --image - ${IMAGE_PATH} --metadata ${MODEL_METADATA_PATH} --output ${IMAGE_HEADER} - DEPENDS ${IMAGE_PATH} ${MODEL_METADATA_PATH} - ${CMAKE_SOURCE_DIR}/image_to_array.py - VERBATIM -) -target_sources( - mobilesam_prompt_segmentation_example PRIVATE ${MODEL_PTE_HEADER} - ${IMAGE_HEADER} -) - -target_include_directories( - mobilesam_prompt_segmentation_example - PRIVATE ${ET_INCLUDE_PATH} ${ET_DIR_PATH}/runtime/core/portable_type/c10 - ${ET_DIR_PATH}/examples/arm/common ${CMAKE_CURRENT_BINARY_DIR} -) diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/README.md b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/README.md index f422edc6ad9..5473faa16a2 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/README.md +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/README.md @@ -1,110 +1,16 @@ -# MobileSAM Runtime Example +# MobileSAM Runtime -This directory builds a bare-metal Corstone-320 application for a MobileSAM -fixed-prompt `.pte` generated by `model_export/export_mobilesam.py`. +MobileSAM uses the standard ExecuTorch Arm executor runner. The example has no +model-specific C++ runtime or CMake project. -## Build ExecuTorch for Arm +The top-level [`run.sh`](../run.sh) embeds the exported `.pte` in the runner, +builds it for Ethos-U85-256, and launches it on Corstone-320. Semihosting passes +the raw input and output tensors between the host and FVP. -Run the Arm setup first if it has not already been run: +After inference, [`visualize_fvp_output.py`](visualize_fvp_output.py) thresholds +the raw output tensor, checks it against the host quantized mask, and writes: -```bash -./examples/arm/setup.sh --i-agree-to-the-contained-eula -source examples/arm/arm-scratch/setup_path.sh -``` +`arm_test/mobilesam/result/fvp_comparison.png` -Build and install the Arm bare-metal ExecuTorch libraries from -`examples/arm`: - -```bash -cmake --preset arm-baremetal \ - -DCMAKE_BUILD_TYPE=Release \ - -B../../cmake-out-arm ../.. -cmake --build ../../cmake-out-arm --target install -j$(nproc) -``` - -The Arm bare-metal preset installs into the build directory. If you use a -different `-B` path, pass the same path as `-DET_BUILD_DIR_PATH` when -configuring the runtime app. - -## Configure the Runtime App - -Use the `.pte` from the export step and any RGB image: - -```bash -cmake \ - -DCMAKE_TOOLCHAIN_FILE=$(pwd)/ethos-u-setup/arm-none-eabi-gcc.cmake \ - -DET_BUILD_DIR_PATH=../../cmake-out-arm \ - -DET_PTE_FILE_PATH= \ - -DMODEL_METADATA_PATH= \ - -DIMAGE_PATH= \ - -DMASK_THRESHOLD=0.0 \ - -DSYSTEM_CONFIG=Ethos_U85_SYS_DRAM_Mid \ - -DMEMORY_MODE=Dedicated_Sram_384KB \ - -Bmobilesam_point_runtime \ - mobilesam_prompt_segmentation_example_ethos_u/runtime -``` - -The generated image header reads the input shape and normalization values from -the exporter metadata, resizes the longest side, and pads the shorter side with -zeros. `MODEL_METADATA_PATH` defaults to the `.json` file beside the `.pte`. -The generated `448x448` float input is placed in the Corstone-320 DDR input -section so it does not consume BRAM. - -The `MEMORY_MODE` value must match the value used during export. The default -`Dedicated_Sram_384KB` places the Ethos-U tensor arena in DDR and the fast -scratch buffer in SRAM, matching the Corstone-320 linker script. The default -method allocator pool is sized for MobileSAM's reduced `448x448` input tensor. - -The runtime emits validation logs through Arm semihosting by default so the FVP -smoke run below prints the segmentation summary. Configure with -`-DET_SEGMENTATION_SEMIHOSTING_OUTPUT=OFF` when targeting an environment -without semihosting support. - -The default target-side mask threshold is `0.0`, matching MobileSAM's mask-logit -threshold. The runtime logs foreground counts for several thresholds so users -can inspect and retune post-processing without re-exporting the graph. - -## Compile - -```bash -cmake --build mobilesam_point_runtime -j$(nproc) -- mobilesam_prompt_segmentation_example -``` - -## Run on Corstone-320 FVP - -```bash -../../backends/arm/scripts/run_fvp.sh \ - --elf=mobilesam_point_runtime/mobilesam_prompt_segmentation_example \ - --target=ethos-u85-256 \ - --timeout=300 \ - --semihosting-cwd=$(pwd)/mobilesam_point_runtime \ - --fast 2>&1 | tee mobilesam_fvp.log -``` - -Expected logs include: - -- `MobileSAM Ethos-U example started`. -- Input and output tensor shapes. -- `Mask threshold`. -- `Segmentation mask hash`. -- Foreground/background pixel counts. -- `Model executed successfully.` - -To dump the predicted binary mask as run-length encoded log chunks, configure -with `-DET_SEGMENTATION_DUMP_MASK=ON`. This is useful for debugging but makes -the UART output larger. - -Capture that FVP output and reconstruct the target mask and overlay with: - -```bash -python mobilesam_prompt_segmentation_example_ethos_u/runtime/visualize_fvp_output.py \ - --fvp-log=mobilesam_fvp.log \ - --input-image= \ - --metadata= \ - --reference-mask= \ - --minimum-iou=0.9 \ - --output-dir=mobilesam_fvp_visual -``` - -The tool writes the decoded target mask, an input/prompt image, a colored FVP -overlay, a side-by-side comparison, and target/reference agreement metrics. +The runtime stage passes when the runner reports successful execution and the +FVP/reference mask IoU is at least `0.9`. diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/image_to_array.py b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/image_to_array.py deleted file mode 100644 index a3d169907ec..00000000000 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/image_to_array.py +++ /dev/null @@ -1,123 +0,0 @@ -# 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. -# -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import json -import os -from argparse import ArgumentParser -from pathlib import Path - -import numpy as np -from PIL import Image - - -def convert_image_to_c_array( - image_path: str, - output_path: str, - image_size: tuple[int, int], - pixel_mean: tuple[float, float, float], - pixel_std: tuple[float, float, float], - array_name: str = "image_data", -) -> None: - image = Image.open(image_path).convert("RGB") - width, height = image.size - target_width, target_height = image_size - if target_width != target_height: - raise ValueError("MobileSAM runtime preprocessing expects a square input.") - - scale = target_width / max(height, width) - resized_size = (round(width * scale), round(height * scale)) - image = image.resize(resized_size, resample=Image.Resampling.BILINEAR) - - data = np.asarray(image, dtype=np.float32) - data = (data - np.asarray(pixel_mean, dtype=np.float32)) / np.asarray( - pixel_std, dtype=np.float32 - ) - padded_data = np.zeros((target_height, target_width, 3), dtype=np.float32) - padded_data[: resized_size[1], : resized_size[0], :] = data - data = np.transpose(padded_data, (2, 0, 1)).flatten() - - array_lines = [] - for i in range(0, len(data), 12): - line = ", ".join(f"{value:.8f}" for value in data[i : i + 12]) - array_lines.append(" " + line + ",") - - c_array = f"""#include -#include - -const size_t image_width = {image_size[0]}; -const size_t image_height = {image_size[1]}; -const size_t image_channels = 3; -__attribute__((section("input_data_sec"), aligned(16))) float {array_name}[{len(data)}] = {{ -{os.linesep.join(array_lines)} -}}; -""" - with open(output_path, "w") as output_file: - output_file.write(c_array) - print(f"Converted '{image_path}' to '{output_path}' ({len(data)} floats)") - - -def load_model_metadata( - metadata_path: str, -) -> tuple[tuple[int, int], tuple[float, float, float], tuple[float, float, float]]: - metadata = json.loads(Path(metadata_path).read_text()) - input_shape = metadata.get("input_shape") - if not isinstance(input_shape, list) or len(input_shape) != 4: - raise ValueError("Model metadata must contain a four-dimensional input_shape.") - if input_shape[0] != 1 or input_shape[1] != 3: - raise ValueError("MobileSAM runtime expects input shape [1, 3, H, W].") - - preprocessing = metadata.get("preprocessing") - if not isinstance(preprocessing, dict): - raise ValueError("Model metadata does not contain preprocessing values.") - pixel_mean_values = preprocessing.get("pixel_mean") - pixel_std_values = preprocessing.get("pixel_std") - if not isinstance(pixel_mean_values, list) or len(pixel_mean_values) != 3: - raise ValueError("MobileSAM pixel_mean must contain three RGB values.") - if not isinstance(pixel_std_values, list) or len(pixel_std_values) != 3: - raise ValueError("MobileSAM pixel_std must contain three RGB values.") - - image_size = (int(input_shape[3]), int(input_shape[2])) - pixel_mean = tuple(float(value) for value in pixel_mean_values) - pixel_std = tuple(float(value) for value in pixel_std_values) - return ( - image_size, - (pixel_mean[0], pixel_mean[1], pixel_mean[2]), - (pixel_std[0], pixel_std[1], pixel_std[2]), - ) - - -def main() -> None: - parser = ArgumentParser() - parser.add_argument("--image", required=True, help="Path to an RGB image.") - parser.add_argument( - "--output", required=True, help="Output path for the generated C array." - ) - parser.add_argument( - "--metadata", - required=True, - help="Exporter metadata containing input shape and preprocessing values.", - ) - args = parser.parse_args() - - image_size, pixel_mean, pixel_std = load_model_metadata(args.metadata) - convert_image_to_c_array( - args.image, - args.output, - image_size, - pixel_mean, - pixel_std, - ) - - -if __name__ == "__main__": - main() diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/main.cpp b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/main.cpp deleted file mode 100644 index 9acf74ec860..00000000000 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/main.cpp +++ /dev/null @@ -1,413 +0,0 @@ -/* - * 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 -#include -#include -#include -#include -#include -#include - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -#include "arm_memory_allocator.h" -#include "image.h" -#include "model_pte.h" - -using executorch::aten::ScalarType; -using executorch::aten::Tensor; -using executorch::extension::BufferDataLoader; -using executorch::runtime::Error; -using executorch::runtime::EValue; -using executorch::runtime::HierarchicalAllocator; -using executorch::runtime::MemoryAllocator; -using executorch::runtime::MemoryManager; -using executorch::runtime::Method; -using executorch::runtime::MethodMeta; -using executorch::runtime::Program; -using executorch::runtime::Result; -using executorch::runtime::Span; - -const size_t method_allocation_pool_size = - ET_SEGMENTATION_METHOD_ALLOCATOR_POOL_SIZE; -unsigned char __attribute__(( - section("method_allocator_sec"), - aligned(16))) method_allocation_pool[method_allocation_pool_size]; - -const size_t temp_allocation_pool_size = - ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE; -unsigned char __attribute__(( - section(".bss.tensor_arena"), - aligned(16))) temp_allocation_pool[temp_allocation_pool_size]; - -#if defined(ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE) -extern "C" { -size_t ethosu_fast_scratch_size = - ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE; -unsigned char __attribute__((section(".bss.ethosu_scratch"), aligned(16))) -dedicated_sram[ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE]; -unsigned char* ethosu_fast_scratch = dedicated_sram; -} -#endif - -namespace { - -#if defined(ET_SEGMENTATION_MASK_THRESHOLD) -constexpr float kMaskThreshold = ET_SEGMENTATION_MASK_THRESHOLD; -#else -constexpr float kMaskThreshold = 0.0f; -#endif - -#if defined(ET_SEGMENTATION_SEMIHOSTING_OUTPUT) && \ - (defined(__arm__) || defined(__thumb__)) -constexpr uint32_t kSemihostingSysWrite0 = 0x04; -constexpr uint32_t kSemihostingSysExitExtended = 0x20; -constexpr uint32_t kAdpStoppedApplicationExit = 0x20026; - -uint32_t semihosting_call(uint32_t operation_code, const void* argument) { - uint32_t result; - asm volatile( - "mov r0, %[operation]\n" - "mov r1, %[argument]\n" - "bkpt 0xab\n" - "mov %[result], r0\n" - : [result] "=r"(result) - : [operation] "r"(operation_code), [argument] "r"(argument) - : "r0", "r1", "memory"); - return result; -} - -void semihosting_write0(const char* message) { - semihosting_call(kSemihostingSysWrite0, message); -} -#endif - -void write_runtime_line(const char* message) { -#if defined(ET_SEGMENTATION_SEMIHOSTING_OUTPUT) && \ - (defined(__arm__) || defined(__thumb__)) - semihosting_write0(message); - semihosting_write0("\n"); -#else - (void)message; -#endif -} - -void write_runtime_format(const char* format, ...) { - char line[768]; - va_list args; - va_start(args, format); - vsnprintf(line, sizeof(line), format, args); - va_end(args); - write_runtime_line(line); -} - -void request_runtime_exit(int code) { -#if defined(ET_SEGMENTATION_SEMIHOSTING_OUTPUT) && \ - (defined(__arm__) || defined(__thumb__)) - const uint32_t exit_block[2] = { - kAdpStoppedApplicationExit, - static_cast(code), - }; - semihosting_call(kSemihostingSysExitExtended, exit_block); -#else - (void)code; -#endif -} - -uint32_t update_hash(uint32_t hash, uint8_t value) { - hash ^= value; - hash *= 16777619u; - return hash; -} - -#if defined(ET_SEGMENTATION_DUMP_MASK) -void dump_mask_rle(const std::vector& mask) { - ET_LOG(Info, "Segmentation mask RLE begin"); - write_runtime_line("Segmentation mask RLE begin"); - char line[512]; - size_t line_len = 0; - line[0] = '\0'; - - size_t index = 0; - while (index < mask.size()) { - const uint8_t class_id = mask[index]; - size_t run_len = 1; - while (index + run_len < mask.size() && mask[index + run_len] == class_id) { - ++run_len; - } - - char entry[32]; - const int entry_len = snprintf( - entry, - sizeof(entry), - "%u:%zu,", - static_cast(class_id), - run_len); - if (entry_len <= 0) { - break; - } - if (line_len + static_cast(entry_len) >= sizeof(line)) { - ET_LOG(Info, "Segmentation mask RLE chunk %s", line); - write_runtime_format("Segmentation mask RLE chunk %s", line); - line_len = 0; - line[0] = '\0'; - } - line_len += static_cast( - snprintf(line + line_len, sizeof(line) - line_len, "%s", entry)); - index += run_len; - } - if (line_len > 0) { - ET_LOG(Info, "Segmentation mask RLE chunk %s", line); - write_runtime_format("Segmentation mask RLE chunk %s", line); - } - ET_LOG(Info, "Segmentation mask RLE end"); - write_runtime_line("Segmentation mask RLE end"); -} -#endif - -void summarize_segmentation_output(const Tensor& out) { - ET_CHECK_MSG( - out.dim() == 4, - "Expected mask logits with shape [1, 1, height, width], got rank %zd", - out.dim()); - ET_CHECK_MSG( - out.scalar_type() == ScalarType::Float, - "Expected float mask logits, got dtype %d", - out.scalar_type()); - ET_CHECK_MSG(out.size(0) == 1, "Only batch size 1 is supported."); - ET_CHECK_MSG( - out.size(1) == 1, - "MobileSAM fixed-prompt export expects one mask channel, got %zd", - out.size(1)); - - const size_t height = static_cast(out.size(2)); - const size_t width = static_cast(out.size(3)); - const auto strides = out.strides(); - const float* data = out.const_data_ptr(); - - size_t foreground_pixels = 0; -#if defined(ET_SEGMENTATION_DUMP_MASK) - std::vector mask(height * width, 0); -#endif - uint32_t mask_hash = 2166136261u; - float min_score = data[0]; - float max_score = data[0]; - double score_sum = 0.0; - const float threshold_sweep[] = { - 0.0f, - -2.0f, - -4.0f, - -5.0f, - -6.0f, - -7.0f, - -8.0f, - -10.0f, - -12.0f, - -14.0f, - }; - size_t threshold_sweep_counts - [sizeof(threshold_sweep) / sizeof(threshold_sweep[0])] = {}; - - for (size_t y = 0; y < height; ++y) { - for (size_t x = 0; x < width; ++x) { - const float score = data[y * strides[2] + x * strides[3]]; - min_score = std::min(min_score, score); - max_score = std::max(max_score, score); - score_sum += score; - for (size_t i = 0; - i < sizeof(threshold_sweep) / sizeof(threshold_sweep[0]); - ++i) { - threshold_sweep_counts[i] += score > threshold_sweep[i] ? 1 : 0; - } - const uint8_t mask_value = score > kMaskThreshold ? 1 : 0; - foreground_pixels += mask_value; -#if defined(ET_SEGMENTATION_DUMP_MASK) - mask[y * width + x] = mask_value; -#endif - mask_hash = update_hash(mask_hash, mask_value); - } - } - - ET_LOG(Info, "Output mask logits shape = [1, 1, %zu, %zu]", height, width); - write_runtime_format( - "Output mask logits shape = [1, 1, %zu, %zu]", height, width); - ET_LOG( - Info, - "Segmentation input image = %zu x %zu x %zu", - image_width, - image_height, - image_channels); - write_runtime_format( - "Segmentation input image = %zu x %zu x %zu", - image_width, - image_height, - image_channels); - ET_LOG(Info, "Mask threshold = %.4f", static_cast(kMaskThreshold)); - write_runtime_format( - "Mask threshold = %.4f", static_cast(kMaskThreshold)); - ET_LOG( - Info, - "Mask logits min/max/mean = %.6f / %.6f / %.6f", - static_cast(min_score), - static_cast(max_score), - score_sum / static_cast(height * width)); - write_runtime_format( - "Mask logits min/max/mean = %.6f / %.6f / %.6f", - static_cast(min_score), - static_cast(max_score), - score_sum / static_cast(height * width)); - for (size_t i = 0; i < sizeof(threshold_sweep) / sizeof(threshold_sweep[0]); - ++i) { - write_runtime_format( - "Threshold %.1f foreground pixels = %zu", - static_cast(threshold_sweep[i]), - threshold_sweep_counts[i]); - } - ET_LOG(Info, "Segmentation mask hash = 0x%08" PRIx32, mask_hash); - write_runtime_format("Segmentation mask hash = 0x%08" PRIx32, mask_hash); - ET_LOG(Info, "Mask foreground pixels = %zu", foreground_pixels); - write_runtime_format("Mask foreground pixels = %zu", foreground_pixels); - ET_LOG( - Info, "Mask background pixels = %zu", height * width - foreground_pixels); - write_runtime_format( - "Mask background pixels = %zu", height * width - foreground_pixels); - -#if defined(ET_SEGMENTATION_DUMP_MASK) - dump_mask_rle(mask); -#endif -} - -} // namespace - -int main() { - executorch::runtime::runtime_init(); - ET_LOG(Info, "Runtime initialized"); - write_runtime_line("MobileSAM Ethos-U example started"); - BufferDataLoader loader(model_pte, sizeof(model_pte)); - ET_LOG(Info, "Size of the model = %zu", sizeof(model_pte)); - write_runtime_format("Model size = %zu bytes", sizeof(model_pte)); - write_runtime_line("Loading ExecuTorch program"); - Result program = Program::load(&loader); - ET_CHECK_MSG(program.ok(), "Program::load failed: 0x%x", program.error()); - write_runtime_line("Program loaded"); - - const auto method_name_result = program->get_method_name(0); - ET_CHECK_MSG(method_name_result.ok(), "Program has no methods"); - const char* method_name = *method_name_result; - ET_LOG(Info, "Running method %s", method_name); - write_runtime_format("Running method %s", method_name); - - Result method_meta_result = program->method_meta(method_name); - ET_CHECK_MSG( - method_meta_result.ok(), - "method_meta lookup failed: 0x%x", - method_meta_result.error()); - - ArmMemoryAllocator method_allocator( - method_allocation_pool_size, method_allocation_pool); - ArmMemoryAllocator temp_allocator( - temp_allocation_pool_size, temp_allocation_pool); - - std::vector planned_buffers; - std::vector> planned_spans; - const size_t num_memory_planned_buffers = - method_meta_result->num_memory_planned_buffers(); - ET_LOG(Info, "num_memory_planned_buffers = %zu", num_memory_planned_buffers); - for (size_t id = 0; id < num_memory_planned_buffers; ++id) { - const size_t buffer_size = - method_meta_result->memory_planned_buffer_size(id).get(); - ET_LOG(Info, "Planned memory buffer_size %zu %zu bytes", id, buffer_size); - - uint8_t* buffer = reinterpret_cast( - method_allocator.allocate(buffer_size, 16UL)); - ET_CHECK_MSG( - buffer != nullptr, - "Could not allocate memory for memory planned buffer size %zu", - buffer_size); - planned_buffers.push_back(buffer); - planned_spans.push_back({planned_buffers.back(), buffer_size}); - } - HierarchicalAllocator planned_memory( - {planned_spans.data(), planned_spans.size()}); - - MemoryManager memory_manager( - &method_allocator, &planned_memory, &temp_allocator); - write_runtime_line("Loading method"); - Result method = program->load_method(method_name, &memory_manager); - ET_CHECK_MSG(method.ok(), "load_method failed: 0x%x", method.error()); - write_runtime_line("Method loaded"); - - const size_t num_inputs = method->inputs_size(); - ET_LOG(Info, "Number of input tensors = %zu", num_inputs); - ET_CHECK_MSG( - num_inputs == 1, - "The segmentation model has a single input tensor, but the provided model has %zu input tensors", - num_inputs); - - EValue* input_evalues = method_allocator.allocateList(num_inputs); - Error err = method->get_inputs(input_evalues, num_inputs); - ET_CHECK_MSG(err == Error::Ok, "get_inputs failed"); - Tensor& input_tensor = input_evalues[0].toTensor(); - const size_t expected_elems = input_tensor.numel(); - const size_t image_elements = sizeof(image_data) / sizeof(image_data[0]); - ET_CHECK_MSG( - expected_elems == image_elements, - "Input tensor expects %zu elements, but image_data has %zu elements", - expected_elems, - image_elements); - ET_CHECK_MSG( - input_tensor.scalar_type() == ScalarType::Float, - "Expected float input tensor, got dtype %d", - input_tensor.scalar_type()); - - float* input_data = input_tensor.mutable_data_ptr(); - write_runtime_format("Copying %zu input elements", expected_elems); - for (size_t i = 0; i < expected_elems; ++i) { - input_data[i] = image_data[i]; - } - - write_runtime_line("Running model execution"); - Error status_inference = method->execute(); - ET_CHECK_MSG( - status_inference == Error::Ok, - "Inference failed 0x%" PRIx32, - status_inference); - write_runtime_line("Inference finished"); - - const size_t num_outputs = method->outputs_size(); - std::vector outputs(num_outputs); - Error status_outputs = method->get_outputs(outputs.data(), outputs.size()); - ET_CHECK_MSG( - status_outputs == Error::Ok, - "get_outputs failed 0x%" PRIx32, - status_outputs); - - for (size_t i = 0; i < outputs.size(); ++i) { - if (outputs[i].isTensor()) { - summarize_segmentation_output(outputs[i].toTensor()); - ET_LOG(Info, "Model executed successfully."); - write_runtime_line("Model executed successfully."); - request_runtime_exit(0); - return 0; - } - } - - ET_CHECK_MSG(false, "No tensor output found."); - request_runtime_exit(1); - return 1; -} diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/visualize_fvp_output.py b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/visualize_fvp_output.py index bb176c3adf4..df8f73a0a1f 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/visualize_fvp_output.py +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/visualize_fvp_output.py @@ -3,192 +3,80 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import argparse import json -import re from pathlib import Path -from typing import Any +import numpy as np from PIL import Image, ImageDraw -RLE_PREFIX = "Segmentation mask RLE chunk " +ROOT = Path(__file__).resolve().parents[4] +WORK_DIR = ROOT / "arm_test" / "mobilesam" +EXPORT_DIR = WORK_DIR / "export" +RESULT_DIR = WORK_DIR / "result" +OUTPUT_SIZE = 112 +POINT = (219, 193) -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Visualize and validate a MobileSAM mask dumped by the FVP." - ) - parser.add_argument("--fvp-log", required=True, type=Path) - parser.add_argument("--input-image", required=True, type=Path) - parser.add_argument("--metadata", required=True, type=Path) - parser.add_argument("--output-dir", required=True, type=Path) - parser.add_argument("--reference-mask", type=Path) - parser.add_argument("--minimum-iou", type=float) - return parser.parse_args() - - -def parse_rle_mask(log_path: Path, expected_pixels: int) -> list[int]: - mask: list[int] = [] - in_dump = False - for line in log_path.read_text().splitlines(): - if "executorch:main.cpp:" in line: - continue - if "Segmentation mask RLE begin" in line: - in_dump = True - continue - if "Segmentation mask RLE end" in line: - break - if not in_dump or RLE_PREFIX not in line: - continue - payload = line.split(RLE_PREFIX, maxsplit=1)[1] - for value, count in re.findall(r"([01]):([0-9]+),", payload): - mask.extend([int(value)] * int(count)) - - if len(mask) != expected_pixels: - raise ValueError( - f"FVP RLE contains {len(mask)} pixels; expected {expected_pixels}." - ) - return mask - - -def prepare_input_image(image_path: Path, input_size: int) -> Image.Image: - image = Image.open(image_path).convert("RGB") - width, height = image.size - scale = input_size / max(width, height) - resized = image.resize( - (round(width * scale), round(height * scale)), - Image.Resampling.BILINEAR, - ) - padded = Image.new("RGB", (input_size, input_size)) - padded.paste(resized, (0, 0)) - return padded +def load_mask(path: Path) -> np.ndarray: + logits = np.fromfile(path, dtype=np.float32) + if logits.size != OUTPUT_SIZE * OUTPUT_SIZE: + raise ValueError(f"Expected {OUTPUT_SIZE**2} logits, got {logits.size}.") + return (logits.reshape(OUTPUT_SIZE, OUTPUT_SIZE) > 0).astype(np.uint8) + +def iou(first: np.ndarray, second: np.ndarray) -> float: + intersection = np.logical_and(first, second).sum() + union = np.logical_or(first, second).sum() + return 1.0 if union == 0 else float(intersection / union) -def create_overlay( - image: Image.Image, mask: Image.Image, color: tuple[int, int, int] + +def overlay( + image: Image.Image, mask: np.ndarray, color: tuple[int, int, int] ) -> Image.Image: - resized_mask = mask.resize(image.size, Image.Resampling.NEAREST) - color_layer = Image.new("RGB", image.size, color) - blended = Image.blend(image, color_layer, 0.45) - overlay = image.copy() - overlay.paste(blended, mask=resized_mask) - return overlay - - -def draw_prompts(image: Image.Image, metadata: dict[str, Any]) -> Image.Image: - result = image.copy() - draw = ImageDraw.Draw(result) - for point_x, point_y in metadata["point_prompts_xy"]: - radius = max(4, metadata["input_size"] // 80) - draw.ellipse( - ( - point_x - radius, - point_y - radius, - point_x + radius, - point_y + radius, - ), - fill=(255, 48, 48), - outline=(255, 255, 255), - width=2, - ) - return result - - -def add_title(image: Image.Image, title: str) -> Image.Image: - title_height = 30 - panel = Image.new("RGB", (image.width, image.height + title_height), "white") - panel.paste(image, (0, title_height)) - ImageDraw.Draw(panel).text((10, 8), title, fill="black") - return panel - - -def binary_metrics(mask: list[int], reference: list[int]) -> tuple[float, float]: - intersection = sum(a == 1 and b == 1 for a, b in zip(mask, reference)) - union = sum(a == 1 or b == 1 for a, b in zip(mask, reference)) - iou = intersection / union if union else 1.0 - agreement = sum(a == b for a, b in zip(mask, reference)) / len(mask) - return iou, agreement + resized = Image.fromarray(mask * 255).resize(image.size, Image.Resampling.NEAREST) + pixels = np.asarray(image, dtype=np.float32).copy() + selected = np.asarray(resized) > 0 + pixels[selected] = pixels[selected] * 0.55 + np.asarray(color) * 0.45 + return Image.fromarray(pixels.astype(np.uint8)) def main() -> None: - args = parse_args() - metadata = json.loads(args.metadata.read_text()) - _, _, mask_height, mask_width = metadata["output_shape"] - mask = parse_rle_mask(args.fvp_log, mask_width * mask_height) - mask_image = Image.new("L", (mask_width, mask_height)) - mask_image.putdata([value * 255 for value in mask]) - - args.output_dir.mkdir(parents=True, exist_ok=True) - mask_image.save(args.output_dir / "fvp_mask.png") - input_image = prepare_input_image(args.input_image, metadata["input_size"]) - prompted_input = draw_prompts(input_image, metadata) - prompted_input.save(args.output_dir / "input_with_prompts.png") - fvp_overlay = draw_prompts( - create_overlay(input_image, mask_image, (0, 220, 120)), metadata + RESULT_DIR.mkdir(parents=True, exist_ok=True) + image = Image.open(EXPORT_DIR / "input.png").convert("RGB") + reference = (np.asarray(Image.open(EXPORT_DIR / "quantized_mask.png")) > 0).astype( + np.uint8 ) - fvp_overlay.save(args.output_dir / "fvp_overlay.png") - - panels = [add_title(prompted_input, "Input and positive prompt")] - metrics: dict[str, Any] = { - "background_pixels": mask.count(0), - "foreground_pixels": mask.count(1), - "output_mask_size": [mask_width, mask_height], - } - if metrics["foreground_pixels"] in (0, len(mask)): - raise RuntimeError( - "FVP produced a degenerate mask with " - f"{metrics['foreground_pixels']} foreground pixels." - ) - - below_minimum_iou = False - if args.reference_mask is not None: - reference_image = ( - Image.open(args.reference_mask) - .convert("L") - .resize((mask_width, mask_height), Image.Resampling.NEAREST) - ) - reference = [int(value > 0) for value in reference_image.tobytes()] - iou, agreement = binary_metrics(mask, reference) - metrics["fvp_reference_iou"] = iou - metrics["fvp_reference_pixel_agreement"] = agreement - reference_overlay = draw_prompts( - create_overlay(input_image, reference_image, (0, 170, 255)), metadata - ) - panels.append(add_title(reference_overlay, "Host quantized mask")) - below_minimum_iou = args.minimum_iou is not None and iou < args.minimum_iou - elif args.minimum_iou is not None: - raise ValueError("--minimum-iou requires --reference-mask.") - - panels.append(add_title(fvp_overlay, "FVP mask")) - comparison = Image.new( - "RGB", - (sum(panel.width for panel in panels), max(panel.height for panel in panels)), - "white", + fvp_mask = load_mask(WORK_DIR / "io" / "output-0.bin") + score = iou(fvp_mask, reference) + if score < 0.9: + raise RuntimeError(f"FVP/reference mask IoU is too low: {score:.4f}") + + prompted = image.copy() + ImageDraw.Draw(prompted).ellipse( + (POINT[0] - 6, POINT[1] - 6, POINT[0] + 6, POINT[1] + 6), fill="red" ) - offset = 0 - for panel in panels: - comparison.paste(panel, (offset, 0)) - offset += panel.width - comparison.save(args.output_dir / "fvp_comparison.png") - (args.output_dir / "metrics.json").write_text( - json.dumps(metrics, indent=2, sort_keys=True) + "\n" + reference_overlay = overlay(image, reference, (0, 170, 255)) + fvp_overlay = overlay(image, fvp_mask, (0, 220, 120)) + panels = ( + ("Input and prompt", prompted), + ("Host quantized mask", reference_overlay), + ("FVP output", fvp_overlay), ) - - print( - f"FVP mask: {metrics['foreground_pixels']} foreground pixels, " - f"artifacts saved to {args.output_dir}" + comparison = Image.new("RGB", (image.width * 3, image.height + 32), "white") + drawing = ImageDraw.Draw(comparison) + for index, (label, panel) in enumerate(panels): + x = index * image.width + drawing.text((x + 10, 10), label, fill="black") + comparison.paste(panel, (x, 32)) + + Image.fromarray(fvp_mask * 255).save(RESULT_DIR / "fvp_mask.png") + comparison.save(RESULT_DIR / "fvp_comparison.png") + (RESULT_DIR / "metrics.json").write_text( + json.dumps({"fvp_reference_iou": score}, indent=2) + "\n" ) - if "fvp_reference_iou" in metrics: - print( - f"FVP/reference IoU={metrics['fvp_reference_iou']:.4f} " - f"agreement={metrics['fvp_reference_pixel_agreement']:.4f}" - ) - if below_minimum_iou: - raise RuntimeError( - f"FVP/reference IoU {metrics['fvp_reference_iou']:.4f} is below " - f"{args.minimum_iou:.4f}." - ) + print(f"FVP/reference mask IoU: {score:.4f}") + print(f"Saved {RESULT_DIR / 'fvp_comparison.png'}") if __name__ == "__main__": diff --git a/examples/arm/model-explorer.md b/examples/arm/model-explorer.md new file mode 100644 index 00000000000..d42a0d04f24 --- /dev/null +++ b/examples/arm/model-explorer.md @@ -0,0 +1,137 @@ + + +# Visualize Arm models and Ethos-U performance + +The Arm example scripts can open ExecuTorch PTE or TOSA graphs in +[Model Explorer](https://github.com/google-ai-edge/model-explorer). For models +run on an Ethos-U FVP, they can also overlay the measured number of cycles on +the corresponding TOSA operators. + +The performance overlay helps identify expensive operators and compare the +NPU work performed by different versions of a model. It is generated from a +trace of an actual FVP execution, rather than from Vela's compile-time cycle +estimates. + +## Visualize a graph + +Use `--visualize_pte` to inspect the ExecuTorch program, including its delegate +calls: + +```bash +./examples/arm/run.sh \ + --model_name=mv2 \ + --target=ethos-u85-256 \ + --model_explorer \ + --visualize_pte +``` + +Use `--visualize_tosa` to inspect the TOSA graph passed to Vela: + +```bash +./examples/arm/run.sh \ + --model_name=mv2 \ + --target=ethos-u85-256 \ + --model_explorer \ + --visualize_tosa +``` + +These views serve different purposes. The PTE view shows the ExecuTorch +program around each delegate call. The TOSA view exposes the operators inside +the Ethos-U delegate and is the graph to which performance data can be mapped. + +## Overlay measured Ethos-U cycles + +Add `--perf_overlay` to a TOSA visualization: + +```bash +./examples/arm/run.sh \ + --model_name=mv2 \ + --target=ethos-u85-256 \ + --model_explorer \ + --visualize_tosa \ + --perf_overlay +``` + +`run.sh` performs the following additional steps: + +1. Enables compiler debug output so that Vela emits tables mapping command + stream offsets back to TOSA operators. +2. Enables PMU tracing when it runs the model on the FVP. +3. Combines the trace timestamps with the Vela mapping tables. +4. Adds the resulting per-operator duration data to Model Explorer as + `Duration (Cycles)`. + +With the default build root, the relevant MobileNetV2 artifacts are: + +```text +arm_test/mv2/pmu_trace.gz +arm_test/mv2/output/out_debug.xml +``` + +`pmu_trace.gz` contains the FVP trace events. `out_debug.xml` contains the Vela +debug tables required to attribute those events to TOSA operators. The two +files must come from the same compilation and execution. + +Model Explorer color-codes operators by duration. Its node-data panel also +provides aggregate values for a selected layer, which can be used to find the +parts of the delegated graph that consume the most cycles. + +## Open existing artifacts + +To reopen an existing overlay without rebuilding and rerunning the model, pass +the generated files directly to `visualize.py`: + +```bash +python3 examples/arm/visualize.py \ + --model_dir arm_test/mv2 \ + --tosa \ + --trace arm_test/mv2/pmu_trace.gz \ + --tables arm_test/mv2/output/out_debug.xml +``` + +Both `--trace` and `--tables` are required when either one is specified. + +## Compare model versions + +The overlay can expose performance improvements or regressions caused by +changes to model lowering, quantization, fusion, or Vela scheduling. Compare +the aggregate cycle count first. If the graph structure is unchanged, the +per-operator values can then show where the difference originated. + +For a meaningful comparison, keep the following fixed between runs: + +- Ethos-U target and MAC configuration +- Vela and FVP versions +- Vela system configuration and memory mode +- Input shapes and compiler options +- Quantization configuration + +When a change adds, removes, or fuses operators, node identifiers may no longer +correspond between graphs. In that case, compare totals and groups of +semantically equivalent operators instead of matching nodes only by ID. Record +the PTE or source revision and verify model outputs or accuracy alongside the +cycle results. + +The overlay measures work in the Ethos-U command stream. It is not an +end-to-end latency measurement and does not account for portable CPU operators +or all ExecuTorch and delegate overhead. Use +[ETDump](https://docs.pytorch.org/executorch/stable/etdump.html) when the total +runtime behavior is the metric of interest. + +## Current limitations + +- Performance overlays are supported for the TOSA view, not the PTE view. +- The trace parser expects the gzip-compressed JSON trace generated by the + Ethos-U FVP. +- The model directory should contain the TOSA and Vela artifacts from the same + build as the trace. +- Cycle counts are comparable only when the target and memory timing + configuration are equivalent. + +For general information about displaying node data, see the +[Model Explorer custom node data documentation](https://github.com/google-ai-edge/model-explorer/wiki/2.-User-Guide#custom-node-data). diff --git a/examples/arm/run.sh b/examples/arm/run.sh index b69f8f1c4a7..007c7c63cd2 100755 --- a/examples/arm/run.sh +++ b/examples/arm/run.sh @@ -144,6 +144,10 @@ if [ "$perf_overlay" = true ] && [ "$model_explorer" != true ]; then echo "Error: --perf_overlay requires --model_explorer" >&2 exit 1 fi +if [ "$perf_overlay" = true ] && [ "$visualize_tosa" != true ]; then + echo "Error: --perf_overlay requires --visualize_tosa" >&2 + exit 1 +fi # Cortex-M backend is an operator-library, not a delegate; force-disable # --delegate when targeting cortex-m so users don't need --no_delegate. diff --git a/examples/arm/setup.sh b/examples/arm/setup.sh index 33bae7c13e1..55e7a873cbf 100755 --- a/examples/arm/setup.sh +++ b/examples/arm/setup.sh @@ -328,6 +328,7 @@ function create_setup_path(){ if [[ $is_script_sourced -eq 0 ]]; then set -e + ARM_SETUP_CURL_PROGRESS_ARGS=(--progress-bar) if [[ -n "$("${et_dir}/.ci/scripts/detect_ci.sh" --and-not-debug)" ]]; then ARM_SETUP_CURL_PROGRESS_ARGS=(--no-progress-meter) export PIP_PROGRESS_BAR=off @@ -376,9 +377,12 @@ if [[ $is_script_sourced -eq 0 ]]; then # Setup FVP if [[ "${enable_fvps}" -eq 1 ]]; then log_step "fvp" "Setting up Arm Fixed Virtual Platforms" - check_fvp_eula - setup_fvp - install_fvp + if [[ "${OS}" == "Linux" ]]; then + check_fvp_eula + install_fvp + else + setup_fvp + fi fi warn_if_mlsdk_python_is_untested diff --git a/examples/arm/smollm2_example_ethos_u/README.md b/examples/arm/smollm2_example_ethos_u/README.md index 25403fb0e7d..6571264af2b 100644 --- a/examples/arm/smollm2_example_ethos_u/README.md +++ b/examples/arm/smollm2_example_ethos_u/README.md @@ -234,6 +234,45 @@ How to interpret the main options: - `--repetition-penalty 1.1` still matters in greedy mode because it modifies the logits before `argmax`. +### 5.1 Profile prompt processing and decoding + +The runner exposes Ethos-U85 PMU counters for every server-mode inference. +Capture those counters with FVP fast mode disabled: + +```bash +python examples/arm/smollm2_example_ethos_u/generate_sampled.py \ + --fvp examples/arm/arm-scratch/FVP-corstone320/models/Linux64_GCC-9.3/FVP_Corstone_SSE-320 \ + --runner smollm2_ethosu_static_kvq_seq64_w8a16_wikitext/cmake-out/arm_executor_runner \ + --embedded-pte \ + --tokenizer data/tokenizers/smollm2/tokenizer.json \ + --prompt "Once upon a time in a small village," \ + --window 64 \ + --max-context-length 64 \ + --use-kv-cache \ + --max-new-tokens 2 \ + --temperature 0 \ + --no-topk-print \ + --profile-output outputs/ethosu_u85_profile.json \ + --timeout 24000 +``` + +`--profile-output` accepts `.json` or `.csv` and automatically removes the +Ethos-U `--fast` FVP option. The report contains one NPU PMU sample per model +execution, split into `prefill` and `decode` phases. Passing +`--npu-frequency-mhz` also reports an estimated NPU-only token rate. + +This KV-cache demo processes the prompt one token at a time; its prefill result +is therefore the aggregate and average of those token executions, not a +batched-prefill measurement. The final prompt execution supplies the logits for +the first generated token, so steady-state decode executions start with the +second generated token. + +CS-320 provides useful Ethos-U85 NPU cycle and event estimates when its timing +adapters match the target configuration. Cortex-M85 CPU timing and simulator +wall time are not cycle accurate, so the report must not be interpreted as +end-to-end latency or measured device tokens/s. Use FPGA or hardware for those +measurements. + ## 6. Optional: evaluate Wikitext perplexity The KV-cache generation artifact can also be used for step-wise perplexity scoring over the same 64-token context. diff --git a/examples/arm/smollm2_example_ethos_u/generate_sampled.py b/examples/arm/smollm2_example_ethos_u/generate_sampled.py index 67877d21fc3..fe4311a9dec 100644 --- a/examples/arm/smollm2_example_ethos_u/generate_sampled.py +++ b/examples/arm/smollm2_example_ethos_u/generate_sampled.py @@ -4,6 +4,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 argparse +import csv +import json import re import secrets import select @@ -13,8 +15,9 @@ import time from collections import deque +from dataclasses import asdict, dataclass from pathlib import Path -from typing import Deque, List, Optional, Sequence +from typing import Deque, Dict, List, Optional, Sequence import numpy as np from pytorch_tokenizers import ( # type: ignore[import-not-found, import-untyped] @@ -26,6 +29,142 @@ re.MULTILINE, ) +ETHOSU_PMU_CYCLE_PATTERN = re.compile(r"ethosu_pmu_cycle_cntr\s*:\s*(\d+)") +ETHOSU_PMU_COUNTER_PATTERN = re.compile(r"ethosu_pmu_cntr(\d+)\s*:\s*(\d+)") +ETHOSU_DELEGATIONS_PATTERN = re.compile(r"NPU delegations:\s*(\d+)") +ETHOSU85_EVENT_NAMES = [ + "sram_read_beats", + "sram_write_beats", + "external_read_beats", + "external_write_beats", + "npu_idle", + "mac_active", + "weight_decoder_active", +] + + +@dataclass +class EthosUPmuMeasurement: + npu_cycles: int + delegations: int + events: Dict[str, int] + + +@dataclass +class ProfileSample: + prompt_no: int + phase: str + input_pos: int + npu_cycles: int + delegations: int + events: Dict[str, int] + + +def parse_ethosu_pmu(lines: Sequence[str]) -> EthosUPmuMeasurement: + text = "".join(lines) + cycle_match = ETHOSU_PMU_CYCLE_PATTERN.search(text) + if cycle_match is None: + raise RuntimeError("Ethos-U PMU cycle count was not found in FVP output") + + counter_values = { + int(index): int(value) + for index, value in ETHOSU_PMU_COUNTER_PATTERN.findall(text) + } + events = { + name: counter_values.get(index, 0) + for index, name in enumerate(ETHOSU85_EVENT_NAMES) + } + delegations_match = ETHOSU_DELEGATIONS_PATTERN.search(text) + return EthosUPmuMeasurement( + npu_cycles=int(cycle_match.group(1)), + delegations=( + int(delegations_match.group(1)) if delegations_match is not None else 0 + ), + events=events, + ) + + +def summarize_profile( + samples: Sequence[ProfileSample], npu_frequency_mhz: Optional[float] +) -> Dict[str, Dict[str, float]]: + summary: Dict[str, Dict[str, float]] = {} + for phase in ("prefill", "decode"): + phase_samples = [sample for sample in samples if sample.phase == phase] + if not phase_samples: + continue + total_cycles = sum(sample.npu_cycles for sample in phase_samples) + mean_cycles = total_cycles / len(phase_samples) + values = { + "executions": float(len(phase_samples)), + "total_npu_cycles": float(total_cycles), + "mean_npu_cycles": mean_cycles, + } + if npu_frequency_mhz is not None: + values["estimated_npu_tokens_per_second"] = ( + npu_frequency_mhz * 1_000_000 / mean_cycles + ) + summary[phase] = values + return summary + + +def write_profile( + path: Path, + samples: Sequence[ProfileSample], + npu_frequency_mhz: Optional[float], +) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + summary = summarize_profile(samples, npu_frequency_mhz) + if path.suffix.lower() == ".json": + path.write_text( + json.dumps( + { + "metadata": { + "ethosu_fast": False, + "npu_frequency_mhz": npu_frequency_mhz, + "timing_scope": "Ethos-U85 NPU only", + }, + "samples": [asdict(sample) for sample in samples], + "summary": summary, + }, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + elif path.suffix.lower() == ".csv": + fieldnames = [ + "prompt_no", + "phase", + "input_pos", + "npu_cycles", + "delegations", + *ETHOSU85_EVENT_NAMES, + ] + with path.open("w", encoding="utf-8", newline="") as output: + writer = csv.DictWriter(output, fieldnames=fieldnames) + writer.writeheader() + for sample in samples: + row = asdict(sample) + row.update(row.pop("events")) + writer.writerow(row) + else: + raise ValueError("--profile-output must end in .json or .csv") + + print("\nEthos-U85 NPU profile (FVP estimate):") + for phase, values in summary.items(): + line = ( + f" {phase}: executions={int(values['executions'])} " + f"total_cycles={int(values['total_npu_cycles'])} " + f"mean_cycles={values['mean_npu_cycles']:.2f}" + ) + if "estimated_npu_tokens_per_second" in values: + line += ( + " estimated_npu_tokens/s=" + f"{values['estimated_npu_tokens_per_second']:.3f}" + ) + print(line) + print(f"Profile written to {path}") + def prepare_input( ids: List[int], @@ -220,12 +359,17 @@ def __init__( timeout: int, input_names: Optional[Sequence[str]] = None, server_mode: bool = False, + ethosu_fast: bool = True, + collect_profile: bool = False, ) -> None: self._fvp = fvp self._runner = runner self._pte = pte self._timeout = timeout self._server_mode = server_mode + self._ethosu_fast = ethosu_fast + self._collect_profile = collect_profile + self.last_pmu: Optional[EthosUPmuMeasurement] = None self._proc: Optional[subprocess.Popen[str]] = None self._recent_stdout: Deque[str] = deque(maxlen=400) self._tmpdir: Optional[tempfile.TemporaryDirectory[str]] = None @@ -246,7 +390,7 @@ def _init_paths(self, input_names: Sequence[str]) -> None: def _build_command(self, cmd_line: str) -> List[str]: assert self._tmpdir_path is not None - return [ + command = [ self._fvp, "-C", "mps4_board.subsystem.ethosu.num_macs=256", @@ -271,14 +415,19 @@ def _build_command(self, cmd_line: str) -> List[str]: "-C", f"mps4_board.subsystem.cpu0.semihosting-cwd={self._tmpdir_path}", "-C", - "mps4_board.subsystem.ethosu.extra_args='--fast'", - "-C", f"mps4_board.subsystem.cpu0.semihosting-cmd_line='{cmd_line}'", "-a", self._runner, "--timelimit", str(self._timeout), ] + if self._ethosu_fast: + insert_at = command.index("-a") + command[insert_at:insert_at] = [ + "-C", + "mps4_board.subsystem.ethosu.extra_args='--fast'", + ] + return command def close(self) -> None: if self._proc is not None: @@ -347,6 +496,9 @@ def _run_server_once(self, output_path: Path) -> np.ndarray: self._proc.stdin.write("go\n") self._proc.stdin.flush() + self.last_pmu = None + profile_lines: List[str] = [] + deadline = time.monotonic() + self._timeout while time.monotonic() < deadline: self._check_proc() @@ -362,7 +514,11 @@ def _run_server_once(self, output_path: Path) -> np.ndarray: f"\n\n[FVP stdout tail]\n{''.join(self._recent_stdout)}" ) self._recent_stdout.append(line) + if self._collect_profile: + profile_lines.append(line) if "SERVER_INFERENCE_DONE" in line: + if self._collect_profile: + self.last_pmu = parse_ethosu_pmu(profile_lines) if output_path.exists() and output_path.stat().st_size > 0: return np.fromfile(output_path, dtype=np.float32) raise RuntimeError( @@ -434,6 +590,8 @@ def __init__( runner: str, pte: Optional[str], timeout: int, + ethosu_fast: bool = True, + collect_profile: bool = False, ) -> None: self._runner = FvpRunnerSession( fvp, @@ -442,7 +600,10 @@ def __init__( timeout, input_names=["i0.bin", "i1.bin"], server_mode=True, + ethosu_fast=ethosu_fast, + collect_profile=collect_profile, ) + self.samples: List[ProfileSample] = [] def close(self) -> None: self._runner.close() @@ -453,13 +614,41 @@ def __enter__(self) -> "KvFvpRunnerSession": def __exit__(self, exc_type, exc, tb) -> None: # type: ignore[no-untyped-def] self.close() - def run(self, token_id: int, input_pos: int) -> np.ndarray: + def run( + self, + token_id: int, + input_pos: int, + *, + phase: str, + prompt_no: int, + ) -> np.ndarray: logits = self._runner.run_inputs( [ np.array([[token_id]], dtype=np.int32), np.array([input_pos], dtype=np.int32), ] ) + if self._runner.last_pmu is not None: + measurement = self._runner.last_pmu + if measurement.npu_cycles <= 0: + raise RuntimeError( + "Ethos-U PMU returned no cycles; ensure FVP fast mode is disabled" + ) + self.samples.append( + ProfileSample( + prompt_no, + phase, + input_pos, + measurement.npu_cycles, + measurement.delegations, + measurement.events, + ) + ) + print( + f"\n[Ethos-U profile phase={phase} input_pos={input_pos} " + f"npu_cycles={measurement.npu_cycles}]", + flush=True, + ) return logits.reshape(1, -1)[0] @@ -573,7 +762,7 @@ def run_one_prompt_kv( logits = None for pos, token_id in enumerate(ids): - logits = runner.run(token_id, pos) + logits = runner.run(token_id, pos, phase="prefill", prompt_no=prompt_no) if topk_print: token_text = tokenizer.decode_token(int(token_id)) print( @@ -601,7 +790,12 @@ def run_one_prompt_kv( print(tokenizer.decode_token(next_id), end="", flush=True) if next_id == eos_id or step == max_new_tokens - 1: break - logits = runner.run(next_id, len(ids) - 1) + logits = runner.run( + next_id, + len(ids) - 1, + phase="decode", + prompt_no=prompt_no, + ) print("\n=== Generation complete ===") decoded = tokenizer.decode(ids) @@ -728,6 +922,23 @@ def main() -> None: default=120, help="FVP time limit in seconds for each runner call.", ) + parser.add_argument( + "--no-ethosu-fast", + action="store_true", + help="Disable Ethos-U FVP fast mode. Required for NPU PMU profiling.", + ) + parser.add_argument( + "--profile-output", + type=Path, + default=None, + help="Write per-token Ethos-U85 PMU samples and summaries to .json or .csv.", + ) + parser.add_argument( + "--npu-frequency-mhz", + type=float, + default=None, + help="Optional assumed NPU frequency for NPU-only token/s estimates.", + ) parser.add_argument( "--full-logits", action="store_true", @@ -757,11 +968,25 @@ def main() -> None: pte_path = None if args.embedded_pte else args.pte if not args.embedded_pte and pte_path is None: raise ValueError("--pte is required unless --embedded-pte is set") + if args.profile_output is not None and not args.use_kv_cache: + raise ValueError("--profile-output requires --use-kv-cache") + if args.profile_output is not None and args.profile_output.suffix.lower() not in { + ".csv", + ".json", + }: + raise ValueError("--profile-output must end in .json or .csv") + if args.npu_frequency_mhz is not None and args.npu_frequency_mhz <= 0: + raise ValueError("--npu-frequency-mhz must be greater than zero") max_context_length = args.max_context_length or args.window if args.use_kv_cache: with KvFvpRunnerSession( - args.fvp, args.runner, pte_path, args.timeout + args.fvp, + args.runner, + pte_path, + args.timeout, + ethosu_fast=not (args.no_ethosu_fast or args.profile_output is not None), + collect_profile=args.profile_output is not None, ) as runner: for i, prompt in enumerate(prompts): run_one_prompt_kv( @@ -779,6 +1004,12 @@ def main() -> None: save_generations_path=args.save_generations, topk_print=not args.no_topk_print, ) + if args.profile_output is not None: + write_profile( + args.profile_output, + runner.samples, + args.npu_frequency_mhz, + ) else: with FvpRunnerSession(args.fvp, args.runner, pte_path, args.timeout) as runner: for i, prompt in enumerate(prompts): diff --git a/examples/cuda/README.md b/examples/cuda/README.md index a5421edb035..76831505be5 100644 --- a/examples/cuda/README.md +++ b/examples/cuda/README.md @@ -35,3 +35,42 @@ installed PyTorch build does not list in `torch.cuda.get_arch_list()`. The example emits `amd_triton.pte` and `aoti_cuda_blob.ptd`. It uses a fresh Inductor cache and fails unless it finds generated Triton source there, and it checks that the `.pte` embeds a code object for the architecture it compiled for. + +## Merge native NVIDIA GPU exports + +CUDA AOTI exports record their compiled target SM. Exports of the same program +and weights can be combined so the runtime selects an exactly matching native +AOTI library, or uses one explicitly designated PTX fallback. + +```bash +python -m executorch.backends.cuda.merge_ptes \ + --input-pte a100/model.pte \ + --input-pte rtx5090/model.pte \ + --input-ptd a100/aoti_cuda_blob.ptd \ + --input-ptd rtx5090/aoti_cuda_blob.ptd \ + --fallback-pte portable/model.pte \ + --fallback-ptd portable/aoti_cuda_blob.ptd \ + --output-pte merged/model.pte \ + --output-ptd merged/aoti_cuda_blob.ptd +``` + +The inputs must come from the same ExecuTorch program and contain identical +weights. Each regular `--input-pte` contributes only exact-SM native cubins; +any PTX capability in a regular input is ignored. At most one +`--fallback-pte` may be provided, and it must contain exactly one PTX-capable +variant. The runtime uses it only when no regular input provides a native cubin +for the current SM. The output PTD reuses one validated copy of the weights. +After merging, the tool prints every native SM and PTX fallback together with +its source PTE. + +Export every regular input with PTX disabled: + +```python +CompileSpec("cuda_include_ptx", b"OFF") +``` + +Export the fallback with PTX enabled and with any portability constraints, such +as a shared-memory limit, required by its target GPU set. AOTI host code and its +CUDA fatbin are linked into one shared library, so the merge step does not +rewrite ELF sections. Instead, the merged metadata makes regular libraries +native-only and marks the fallback library as PTX-only for runtime selection. diff --git a/examples/cuda/scripts/export.py b/examples/cuda/scripts/export.py index ee4390bf938..8eb9c651b24 100644 --- a/examples/cuda/scripts/export.py +++ b/examples/cuda/scripts/export.py @@ -19,6 +19,7 @@ from executorch.examples.models.model_factory import EagerModelFactory from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower +from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.extension.export_util.utils import save_pte_program @@ -51,6 +52,16 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument("--generate_etrecord", action=argparse.BooleanOptionalAction) parser.add_argument("--save_processed_bytes", action=argparse.BooleanOptionalAction) + parser.add_argument( + "--cuda_include_ptx", + choices=("ON", "OFF"), + help="Explicitly enable or disable PTX in the exported CUDA AOTI library", + ) + parser.add_argument( + "--seed", + type=int, + help="Seed model initialization and example input generation", + ) args = parser.parse_args() return args @@ -73,6 +84,9 @@ def main(): f"Available models are {list(MODEL_NAME_TO_MODEL.keys())}." ) + if args.seed is not None: + torch.manual_seed(args.seed) + ( model, example_args, @@ -87,9 +101,12 @@ def main(): dynamic_shapes=dynamic_shapes, ) - partitioner = CudaPartitioner( - [CudaBackend.generate_method_name_compile_spec(args.model_name)] - ) + compile_specs = [CudaBackend.generate_method_name_compile_spec(args.model_name)] + if args.cuda_include_ptx is not None: + compile_specs.append( + CompileSpec("cuda_include_ptx", args.cuda_include_ptx.encode()) + ) + partitioner = CudaPartitioner(compile_specs) et_prog = to_edge_transform_and_lower( exported_programs, diff --git a/examples/mediatek/mtk_build_examples.sh b/examples/mediatek/mtk_build_examples.sh index afdd9f16d51..c8ea0e927c1 100755 --- a/examples/mediatek/mtk_build_examples.sh +++ b/examples/mediatek/mtk_build_examples.sh @@ -43,7 +43,7 @@ main() { -B"${example_build_dir}" \ $EXECUTORCH_ROOT/$example_dir - cmake --build "${example_build_dir}" -j5 + cmake --build "${example_build_dir}" -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) # Switch back to the original directory cd - > /dev/null diff --git a/examples/models/llama/README.md b/examples/models/llama/README.md index 5a65eaa0cb4..6cbdd8559fd 100644 --- a/examples/models/llama/README.md +++ b/examples/models/llama/README.md @@ -288,7 +288,7 @@ cmake -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake \ -DEXECUTORCH_BUILD_KERNELS_LLM=ON \ -Bcmake-out-android . -cmake --build cmake-out-android -j16 --target install --config Release +cmake --build cmake-out-android -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install --config Release ``` **1.2 Build llama runner for android** @@ -307,7 +307,7 @@ cmake -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake \ -Bcmake-out-android/examples/models/llama \ examples/models/llama -cmake --build cmake-out-android/examples/models/llama -j16 --config Release +cmake --build cmake-out-android/examples/models/llama -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release ``` **2. Run on Android via adb shell** @@ -400,7 +400,7 @@ cmake -DPYTHON_EXECUTABLE=python \ -DEXECUTORCH_BUILD_EXTENSION_LLM=ON \ -DEXECUTORCH_BUILD_KERNELS_LLM=ON \ -Bcmake-out . -cmake --build cmake-out -j16 --config Release --target install +cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release --target install ``` Next install the llama runner with torchao kernels enabled (similar to step 3.2 above): @@ -410,7 +410,7 @@ cmake -DPYTHON_EXECUTABLE=python \ -DCMAKE_BUILD_TYPE=Release \ -Bcmake-out/examples/models/llama \ examples/models/llama -cmake --build cmake-out/examples/models/llama -j16 --config Release +cmake --build cmake-out/examples/models/llama -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release ``` Finally run your model (similar to step 3.3 above): @@ -561,15 +561,14 @@ registered in `executorch.extension.llm.custom_ops.custom_ops`. The runtime kernel ships in `extension/llm/custom_ops/op_moe.cpp`. It always compiles with a portable reference fallback (unpack + dequant + -`cpublas::gemm`) that works on any platform. `ENABLE_QUANTIZED_MOE_FFN` -is an **optimization gate**, not a correctness requirement — when -defined, the kernel uses torchao's fused `linear_operator` (NEON -i8mm/dotprod on aarch64) instead of the reference path. +`cpublas::gemm`) that works on any platform. The optimized build option +uses torchao's fused `linear_operator` (NEON i8mm/dotprod on aarch64) +instead of the reference path. In CMake, opt in to the optimized path with: ```cmake --DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON +-DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED=ON ``` In Buck, `_get_quantized_moe_deps()` in `targets.bzl` wires: diff --git a/examples/models/llama/evaluate/eager_eval.py b/examples/models/llama/evaluate/eager_eval.py index cefb951281c..b094b2245cf 100644 --- a/examples/models/llama/evaluate/eager_eval.py +++ b/examples/models/llama/evaluate/eager_eval.py @@ -44,14 +44,16 @@ def eot_token_id(self): """ The stories model does not have an EOT token, so we use the EOS token instead. """ - if hasattr(self._tokenizer, "eot_id"): - return self._tokenizer.eot_id + eot_id = getattr(self._tokenizer, "eot_id", None) + if eot_id is not None: + return eot_id return self._tokenizer.eos_id @property def prefix_token_id(self): - if hasattr(self._tokenizer, "bos_id"): - return self._tokenizer.bos_id + bos_id = getattr(self._tokenizer, "bos_id", None) + if bos_id is not None: + return bos_id return self.eot_token_id @property diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index ffe9a89a570..07240b11d8c 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -115,6 +115,7 @@ "qwen3_5_4b", "phi_4_mini", "smollm2", + "smollm2_360m", "lfm2_350m", # hybrid "lfm2_700m", # hybrid "lfm2_1_2b", # hybrid @@ -128,6 +129,7 @@ "qwen2_5_coder_32b": "Qwen/Qwen2.5-Coder-32B-Instruct", "phi_4_mini": "microsoft/Phi-4-mini-instruct", "smollm2": "HuggingFaceTB/SmolLM2-135M", + "smollm2_360m": "HuggingFaceTB/SmolLM2-360M", "qwen3_0_6b": "Qwen/Qwen3-0.6B", "qwen3_1_7b": "Qwen/Qwen3-1.7B", "qwen3_4b": "Qwen/Qwen3-4B", @@ -592,8 +594,8 @@ def build_args_parser() -> argparse.ArgumentParser: "Replace eager MoE feed-forward modules with the " "`llama::quantized_moe_ffn` portable-runtime custom op (INT4 " "weights, INT8 dyn-quant activations via torchao). On aarch64 " - "with ENABLE_QUANTIZED_MOE_FFN the optimized torchao kernel is " - "used; otherwise a portable reference fallback runs." + "an optimized runtime build uses the torchao kernel; otherwise " + "a portable reference fallback runs." ), ) @@ -712,7 +714,7 @@ def export_llama( # noqa: C901 from executorch.examples.models.qwen3 import convert_weights elif model_name == "phi_4_mini": from executorch.examples.models.phi_4_mini import convert_weights - elif model_name == "smollm2": + elif model_name in ("smollm2", "smollm2_360m"): from executorch.examples.models.smollm2 import convert_weights elif model_name.startswith("lfm2"): from executorch.examples.models.lfm2 import convert_weights diff --git a/examples/models/llama/source_transformation/custom_kv_cache.py b/examples/models/llama/source_transformation/custom_kv_cache.py index dbaac9accf4..71cfe33d753 100644 --- a/examples/models/llama/source_transformation/custom_kv_cache.py +++ b/examples/models/llama/source_transformation/custom_kv_cache.py @@ -265,6 +265,7 @@ def __init__( scale: float = 1.0 / 127.0, use_custom_update_cache_op: bool = True, return_float_values: bool = True, + use_per_channel: bool = True, dtype: torch.dtype = torch.float32, ): super().__init__() @@ -276,9 +277,12 @@ def __init__( self.quantized_cache_dtype = torch.int8 self.return_float_values = return_float_values self.max_context_length = max_context_length + self.use_per_channel = use_per_channel + self.k_cache_scale = scale + self.v_cache_scale = scale self.calibration_enabled = False cache_shape = (max_batch_size, max_context_length, n_heads, head_dim) - scale_shape = (1, 1, 1, head_dim) + scale_shape = (1, 1, 1, head_dim) if use_per_channel else (1,) self.register_buffer( "k_cache", torch.zeros(cache_shape, dtype=self.quantized_cache_dtype), @@ -324,9 +328,11 @@ def finalize_calibration(self): k_scales = self.k_observed_max.to(self.k_cache_scales.dtype) / 127.0 v_scales = self.v_observed_max.to(self.v_cache_scales.dtype) / 127.0 if torch.any(k_scales == 0) or torch.any(v_scales == 0): + qparam_scope = "channel" if self.use_per_channel else "cache" logging.warning( - "Static KV cache calibration observed an all-zero K/V channel; " - "using the smallest positive scale for that channel." + "Static KV cache calibration observed an all-zero K/V %s; " + "using the smallest positive scale.", + qparam_scope, ) # This floor prevents division by zero; it is not an accuracy threshold. self.k_cache_scales.copy_( @@ -335,6 +341,9 @@ def finalize_calibration(self): self.v_cache_scales.copy_( v_scales.clamp_min(torch.finfo(self.v_cache_scales.dtype).tiny) ) + if not self.use_per_channel: + self.k_cache_scale = self.k_cache_scales.item() + self.v_cache_scale = self.v_cache_scales.item() self.calibration_enabled = False self.k_calibration_cache = None self.v_calibration_cache = None @@ -347,30 +356,54 @@ def _observe_and_update(self, input_pos, k_val, v_val): self.k_observed_max.copy_( torch.maximum( self.k_observed_max, - k_val.detach().abs().amax(dim=(0, 1, 2), keepdim=True), + ( + k_val.detach().abs().amax(dim=(0, 1, 2), keepdim=True) + if self.use_per_channel + else k_val.detach().abs().amax().reshape_as(self.k_observed_max) + ), ) ) self.v_observed_max.copy_( torch.maximum( self.v_observed_max, - v_val.detach().abs().amax(dim=(0, 1, 2), keepdim=True), + ( + v_val.detach().abs().amax(dim=(0, 1, 2), keepdim=True) + if self.use_per_channel + else v_val.detach().abs().amax().reshape_as(self.v_observed_max) + ), ) ) self.k_calibration_cache[:, input_pos] = k_val self.v_calibration_cache[:, input_pos] = v_val return self.k_calibration_cache, self.v_calibration_cache - def _quantize(self, value, scales): - # torchao affine custom ops do not yet have the required Arm/TOSA - # lowering and ExecuTorch out-variant runtime support. - qmin = torch.iinfo(self.quantized_cache_dtype).min - qmax = torch.iinfo(self.quantized_cache_dtype).max - return torch.clamp(torch.round(value / scales), qmin, qmax).to( - self.quantized_cache_dtype + def _quantize(self, value, scale): + if self.use_per_channel: + qmin = torch.iinfo(self.quantized_cache_dtype).min + qmax = torch.iinfo(self.quantized_cache_dtype).max + return torch.clamp(torch.round(value / scale), qmin, qmax).to( + self.quantized_cache_dtype + ) + return torch.ops.quantized_decomposed.quantize_per_tensor.default( + value, + scale, + 0, + torch.iinfo(self.quantized_cache_dtype).min, + torch.iinfo(self.quantized_cache_dtype).max, + self.quantized_cache_dtype, ) - def _dequantize(self, value, scales, dtype): - return value.to(dtype) * scales.to(dtype) + def _dequantize(self, value, scale, dtype): + if self.use_per_channel: + return value.to(dtype) * scale.to(dtype) + return torch.ops.quantized_decomposed.dequantize_per_tensor.default( + value, + scale, + 0, + torch.iinfo(self.quantized_cache_dtype).min, + torch.iinfo(self.quantized_cache_dtype).max, + self.quantized_cache_dtype, + ).to(dtype) def _update_cache(self, value, cache, input_pos, indices=None): start_pos = input_pos[0].item() @@ -386,8 +419,10 @@ def _update_cache(self, value, cache, input_pos, indices=None): cache[:, input_pos] = value def _quantize_and_update(self, input_pos, k_val, v_val, indices=None): - quantized_k_val = self._quantize(k_val, self.k_cache_scales) - quantized_v_val = self._quantize(v_val, self.v_cache_scales) + k_scale = self.k_cache_scales if self.use_per_channel else self.k_cache_scale + v_scale = self.v_cache_scales if self.use_per_channel else self.v_cache_scale + quantized_k_val = self._quantize(k_val, k_scale) + quantized_v_val = self._quantize(v_val, v_scale) self._update_cache(quantized_k_val, self.k_cache, input_pos, indices) self._update_cache(quantized_v_val, self.v_cache, input_pos, indices) @@ -395,8 +430,10 @@ def _quantize_and_update(self, input_pos, k_val, v_val, indices=None): def _update_and_return_float_values(self, input_pos, k_val, v_val, indices=None): self._quantize_and_update(input_pos, k_val, v_val, indices) - k_out = self._dequantize(self.k_cache, self.k_cache_scales, k_val.dtype) - v_out = self._dequantize(self.v_cache, self.v_cache_scales, v_val.dtype) + k_scale = self.k_cache_scales if self.use_per_channel else self.k_cache_scale + v_scale = self.v_cache_scales if self.use_per_channel else self.v_cache_scale + k_out = self._dequantize(self.k_cache, k_scale, k_val.dtype) + v_out = self._dequantize(self.v_cache, v_scale, v_val.dtype) self._update_cache(k_val, k_out, input_pos, indices) self._update_cache(v_val, v_out, input_pos, indices) @@ -414,7 +451,7 @@ def update(self, input_pos, k_val, v_val, indices=None): """ k_val, v_val: [B, H, S, D] return: [B, H, S, D] - Storage is [B, S, H, D], with static per-head-dim qparams. + Storage is [B, S, H, D], with static per-head-dim or per-tensor qparams. """ k_val = k_val.transpose(1, 2) @@ -440,6 +477,7 @@ def from_float( kv_cache, scale: float = 1.0 / 127.0, use_custom_update_cache_op: bool = True, + use_per_channel: bool = True, ): if isinstance(kv_cache, CustomKVCache): max_batch_size, max_context_length, n_heads, head_dim = ( @@ -456,6 +494,7 @@ def from_float( head_dim, scale=scale, use_custom_update_cache_op=use_custom_update_cache_op, + use_per_channel=use_per_channel, dtype=kv_cache.k_cache.dtype, ) diff --git a/examples/models/llama/source_transformation/quantize.py b/examples/models/llama/source_transformation/quantize.py index 6bcf35b2a69..65981c57f8c 100644 --- a/examples/models/llama/source_transformation/quantize.py +++ b/examples/models/llama/source_transformation/quantize.py @@ -776,6 +776,7 @@ def get_quant_embedding_transform( embedding_quantize: str, use_shared_embedding: bool = False, quantize_with_hqq: bool = True, + range_learning: bool = False, ): if embedding_quantize.startswith("torchao:"): from torchao.prototype.quantization.embedding.api import ( @@ -819,6 +820,7 @@ def _torchao_embedding_quantizer(model): weight_dtype=weight_dtype, granularity=granularity, mapping_type=mapping_type, + range_learning=range_learning, ).quantize(model) return model diff --git a/examples/models/llama/tests/BUCK b/examples/models/llama/tests/BUCK index 74430d9e306..dc3401ced83 100644 --- a/examples/models/llama/tests/BUCK +++ b/examples/models/llama/tests/BUCK @@ -3,6 +3,17 @@ load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") oncall("executorch") +fbcode_target(_kind = python_unittest, + name = "test_eager_eval", + srcs = [ + "test_eager_eval.py", + ], + deps = [ + "//executorch/examples/models/llama:eval_library", + "fbsource//third-party/pypi/pytest:pytest", + ], +) + fbcode_target(_kind = python_unittest, name = "test_simple_sdpa", srcs = [ diff --git a/examples/models/llama/tests/test_eager_eval.py b/examples/models/llama/tests/test_eager_eval.py new file mode 100644 index 00000000000..bb53292b206 --- /dev/null +++ b/examples/models/llama/tests/test_eager_eval.py @@ -0,0 +1,42 @@ +# 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. + +import unittest +from types import SimpleNamespace +from typing import Optional + +import pytest + +pytest.importorskip("lm_eval", reason="requires lm-evaluation-harness") + +from executorch.examples.models.llama.evaluate.eager_eval import ( # noqa: E402 + EagerEvalWrapper, +) + + +class TestEagerEvalWrapperTokenIds(unittest.TestCase): + @staticmethod + def _wrapper(**token_ids: Optional[int]) -> EagerEvalWrapper: + # HFLM initialization loads a model and is unrelated to these properties. + wrapper = object.__new__(EagerEvalWrapper) + wrapper._tokenizer = SimpleNamespace(**token_ids) # pyre-ignore[8] + return wrapper + + def test_token_id_fallbacks(self): + cases = ( + ({"bos_id": 1, "eot_id": 2, "eos_id": 3}, 2, 1), + ({"bos_id": 0, "eot_id": 2, "eos_id": 3}, 2, 0), + ({"bos_id": None, "eot_id": 2, "eos_id": 3}, 2, 2), + ({"bos_id": None, "eot_id": 0, "eos_id": 3}, 0, 0), + ({"bos_id": None, "eot_id": None, "eos_id": 0}, 0, 0), + ({"eos_id": 0}, 0, 0), + ) + + for token_ids, expected_eot, expected_prefix in cases: + with self.subTest(token_ids=token_ids): + wrapper = self._wrapper(**token_ids) + self.assertEqual(wrapper.eot_token_id, expected_eot) + self.assertEqual(wrapper.prefix_token_id, expected_prefix) diff --git a/examples/models/phi-3-mini-lora/README.md b/examples/models/phi-3-mini-lora/README.md index 62efda6c3dc..fa00e17fd3c 100644 --- a/examples/models/phi-3-mini-lora/README.md +++ b/examples/models/phi-3-mini-lora/README.md @@ -24,7 +24,7 @@ python export_model.py (mkdir cmake-out && cd cmake-out && cmake ..) # Build the executor_runner target -cmake --build cmake-out --target executor_runner -j9 +cmake --build cmake-out --target executor_runner -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) # Run the model for inference. ./cmake-out/executor_runner --model_path phi3_mini_lora.pte diff --git a/examples/models/phi-3-mini/CMakeLists.txt b/examples/models/phi-3-mini/CMakeLists.txt index 3c7ed6a4acb..c5c5eae30ac 100644 --- a/examples/models/phi-3-mini/CMakeLists.txt +++ b/examples/models/phi-3-mini/CMakeLists.txt @@ -14,7 +14,11 @@ # cmake_minimum_required(VERSION 3.24) -cmake_policy(SET CMP0144 NEW) +# CMP0144 arrives in 3.27, and cmake_policy(SET) on a policy the running CMake +# does not know is a hard error, so the request has to be guarded. +if(POLICY CMP0144) + cmake_policy(SET CMP0144 NEW) +endif() project(phi_3_mini_runner) set(CMAKE_CXX_STANDARD 17) diff --git a/examples/models/phi-3-mini/README.md b/examples/models/phi-3-mini/README.md index dac378213d8..1df5844db9c 100644 --- a/examples/models/phi-3-mini/README.md +++ b/examples/models/phi-3-mini/README.md @@ -39,7 +39,7 @@ cmake -DCMAKE_PREFIX_PATH=cmake-out \ -Bcmake-out/examples/models/phi-3-mini \ examples/models/phi-3-mini -cmake --build cmake-out/examples/models/phi-3-mini -j16 --config Release +cmake --build cmake-out/examples/models/phi-3-mini -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release ``` - Run model. Options available [here](https://github.com/pytorch/executorch/blob/main/examples/models/phi-3-mini/main.cpp#L16-L33) ``` diff --git a/examples/models/smollm2/360M_config.json b/examples/models/smollm2/360M_config.json new file mode 100644 index 00000000000..332c4c12337 --- /dev/null +++ b/examples/models/smollm2/360M_config.json @@ -0,0 +1,16 @@ +{ + "dim": 960, + "ffn_dim_multiplier": 1, + "hidden_dim": 2560, + "n_heads": 15, + "n_kv_heads": 5, + "n_layers": 32, + "norm_eps": 1e-05, + "rope_theta": 100000.0, + "use_scaled_rope": false, + "vocab_size": 49152, + "use_hf_rope": false, + "attention_qkv_bias": false, + "bos_idx": 0, + "eos_idx": 0 +} diff --git a/examples/models/smollm2/BUCK b/examples/models/smollm2/BUCK index 6d81065373b..45173b5cbf5 100644 --- a/examples/models/smollm2/BUCK +++ b/examples/models/smollm2/BUCK @@ -14,6 +14,7 @@ fbcode_target(_kind = runtime.python_library, base_module = "executorch.examples.models.smollm2", resources = { "135M_config.json": "135M_config.json", + "360M_config.json": "360M_config.json", }, deps = [ "//caffe2:torch", diff --git a/examples/models/supertonic/README.md b/examples/models/supertonic/README.md index d7089d2802c..2508faf02f9 100644 --- a/examples/models/supertonic/README.md +++ b/examples/models/supertonic/README.md @@ -110,7 +110,7 @@ exit cleanly; closing stdin also exits with status zero. ## Platform and model limits - This workflow requires an Apple silicon Mac, macOS, Xcode command-line - tools, CMake 3.24 or newer, and an ExecuTorch Python environment with the MLX + tools, CMake 3.26 or newer, and an ExecuTorch Python environment with the MLX backend and custom operations available. The native runner supports only arm64 Darwin and uses MLX GPU delegation with FP16 activations. - Exported programs use dynamic sequence lengths, five flow-matching steps, diff --git a/examples/models/whisper/CMakeLists.txt b/examples/models/whisper/CMakeLists.txt index 6a1c2902977..08baf3b825f 100644 --- a/examples/models/whisper/CMakeLists.txt +++ b/examples/models/whisper/CMakeLists.txt @@ -4,7 +4,7 @@ # 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.29) +cmake_minimum_required(VERSION 3.26) project(whisper_runner) set(CMAKE_CXX_STANDARD 17) diff --git a/examples/nxp/analyzing_with_inspector.py b/examples/nxp/analyzing_with_inspector.py index b339af79d6e..e5f6d97a8c1 100644 --- a/examples/nxp/analyzing_with_inspector.py +++ b/examples/nxp/analyzing_with_inspector.py @@ -7,8 +7,17 @@ from typing import Any, Union +from executorch.backends.nxp.tests.profiling_utils import ( + get_neutron_compiler_version, + get_neutron_driver_version, + get_neutron_kernel_kinds, +) + from executorch.devtools import Inspector +# Global mapping of Neutron kernel IDs to names used by the delegate metadata parser. +kernel_kinds = {} + def parse_delegate_metadata( delegate_metadatas: list[bytes], @@ -26,7 +35,13 @@ def parse_delegate_metadata( if function_code == 0: metadata_list.append("Profiling dump") else: - metadata_list.append("Neutron kernel " + str(function_code)) + metadata_list.append( + kernel_kinds.get( + function_code, "Neutron kernel " + str(function_code) + ) + ) + elif len(metadata_bytes) == 2: + metadata_list.append("Profiling dump") else: metadata_list.append("Invalid metadata size") return metadata_list @@ -37,6 +52,12 @@ def parse_delegate_metadata( try: etrecord_path = "etrecord/etrecord.bin" etdump_path = "etdump/trace.etdump" + + driver_version = get_neutron_driver_version(etdump_path) + compiler_version = get_neutron_compiler_version() + if driver_version and driver_version == compiler_version: + kernel_kinds = get_neutron_kernel_kinds() + inspector = Inspector( etdump_path=etdump_path, etrecord=etrecord_path, diff --git a/examples/nxp/aot_neutron_compile.py b/examples/nxp/aot_neutron_compile.py index 697953f7946..b9e3298c26d 100644 --- a/examples/nxp/aot_neutron_compile.py +++ b/examples/nxp/aot_neutron_compile.py @@ -27,6 +27,7 @@ from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner from executorch.backends.nxp.nxp_backend import ( core_aten_ops_exception_list, + default_preserve_ops, generate_neutron_compile_spec, ) from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer @@ -45,6 +46,9 @@ from executorch.examples.nxp.models.mlperf_tiny.image_classification.mlperf_tiny_image_classification import ( MLPerfTinyImageClassification, ) +from executorch.examples.nxp.models.mlperf_tiny.keyword_spotting.mlperf_tiny_keyword_spotting import ( + MLPerfTinyKeywordSpotting, +) from executorch.examples.nxp.models.mobilenet_v2 import MobilenetV2 from executorch.exir import ( EdgeCompileConfig, @@ -64,6 +68,7 @@ "cifar10": CifarNet, "mobilenetv2": MobilenetV2, "mlperf_tiny_image_classification": MLPerfTinyImageClassification, + "mlperf_tiny_keyword_spotting": MLPerfTinyKeywordSpotting, } FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" @@ -99,7 +104,10 @@ def _print_ops_in_edge_program(edge_program): def _get_model_info_from_name( - model_name: str, dataset_path: str | None, use_random_dataset: bool + model_name: str, + dataset_path: str | None, + use_random_dataset: bool, + num_samples: int | None, ): """Given the name of an example pytorch model and args, return the model, its class instance (can be None), example inputs and calibration inputs (can be None). @@ -118,9 +126,11 @@ def _get_model_info_from_name( ) model_cls_inst = model_cls() - elif model_cls is MLPerfTinyImageClassification: + elif model_cls in (MLPerfTinyImageClassification, MLPerfTinyKeywordSpotting): model_cls_inst = model_cls( - dataset_path=dataset_path, use_random_dataset=use_random_dataset + dataset_path=dataset_path, + use_random_dataset=use_random_dataset, + num_samples=num_samples, ) else: @@ -245,7 +255,7 @@ def _get_arg_parser(): required=False, default=False, action="store_true", - help="During conversion to Neutron microcode by Neutron Converter, a kernel selection file will be dumped in " + help="During compilation to Neutron microcode by Neutron Compiler, a kernel selection file will be dumped in " "the working directory. This file can be used for reduction of Neutron Firmware size in the built app." "See `docs/source/backends/nxp/nxp-kernel-selection.md` for details.", ) @@ -256,6 +266,13 @@ def _get_arg_parser(): action="store_true", help="The calibration and testing datasets will be generated randomly instead of being downloaded.", ) + parser.add_argument( + "--num_random_samples", + required=False, + default=None, + type=int, + help="Number of random samples to generate, required when `--use_random_dataset` flag is set.", + ) parser.add_argument( "-dst", "--dataset_path", @@ -285,7 +302,10 @@ def _get_arg_parser(): # 1. pick model from one of the supported lists model, example_inputs, calibration_inputs, model_cls_inst = ( _get_model_info_from_name( - args.model_name, args.dataset_path, args.use_random_dataset + args.model_name, + args.dataset_path, + args.use_random_dataset, + args.num_random_samples, ) ) model = model.eval() @@ -316,7 +336,8 @@ def _get_arg_parser(): quantizer = NeutronQuantizer(neutron_target_spec, is_qat=args.use_qat) if args.use_qat: if not isinstance( - model_cls_inst, (CifarNet, MLPerfTinyImageClassification) + model_cls_inst, + (CifarNet, MLPerfTinyImageClassification, MLPerfTinyKeywordSpotting), ): raise ValueError( f"QAT training is not supported for model '{args.model_name}'" @@ -384,6 +405,7 @@ def _get_arg_parser(): compile_spec, neutron_target_spec, post_quantization_state_dict=module.state_dict(), + preserve_ops=default_preserve_ops, ) ] if args.delegate diff --git a/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/CMakeLists.txt b/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/CMakeLists.txt new file mode 100644 index 00000000000..0b0dd3f8849 --- /dev/null +++ b/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/CMakeLists.txt @@ -0,0 +1,144 @@ +# Copyright 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. + +cmake_minimum_required(VERSION 3.24) + +set(CMAKE_EXECUTABLE_LIBRARY_PREFIX) +set(CMAKE_EXECUTABLE_LIBRARY_SUFFIX) + +# CURRENT DIRECTORY +set(ProjDirPath ${CMAKE_CURRENT_SOURCE_DIR}) + +set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/${CMAKE_BUILD_TYPE}) +set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/${CMAKE_BUILD_TYPE}) + +# Skip link step during compiler check (bare-metal cross-compilation). +set(CMAKE_TRY_COMPILE_TARGET_TYPE STATIC_LIBRARY) + +project(executorch_cifarnet) + +enable_language(ASM) + +set(MCUX_SDK_PROJECT_NAME executorch_cifarnet.elf) + +set(EXECUTORCH_ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../../../../..) + +# CPU and FPU flags required for Cortex-M33 with single-precision FPU. +set(CPU_FLAGS "-mcpu=cortex-m33 -mthumb -mfloat-abi=hard -mfpu=fpv5-sp-d16") +set(CPU_DEFINES + "-DCPU_MIMXRT798SGFOA_cm33_core0 -DCPU_MIMXRT798SGFOB_cm33_core0 \ + -DMIMXRT798S_cm33_core0_SERIES -DMCUXPRESSO_SDK \ + -D__STARTUP_INITIALIZE_NONCACHEDATA -D__STARTUP_CLEAR_BSS \ + -DDSP_IMAGE_COPY_TO_RAM=1 -DBOOT_HEADER_ENABLE=1 \ + -DEIQ_EXAMPLE_HSRUN_CLOCK -DMCUX_META_BUILD \ + -DPRINTF_ADVANCED_ENABLE=1 -DPRINTF_FLOAT_ENABLE=1 -DNO_HEAP_USAGE=1 \ + -DSDK_DEBUGCONSOLE=1 -DSDK_I2C_BASED_COMPONENT_USED=1" +) +set(CMAKE_C_FLAGS + "${CMAKE_C_FLAGS} ${CPU_FLAGS} ${CPU_DEFINES} -fno-common -ffunction-sections -fdata-sections -fno-builtin -mapcs -std=gnu99" +) +set(CMAKE_CXX_FLAGS + "${CMAKE_CXX_FLAGS} ${CPU_FLAGS} ${CPU_DEFINES} -fno-common -ffunction-sections -fdata-sections -fno-builtin -mapcs -fno-rtti -fno-exceptions" +) +set(CMAKE_ASM_FLAGS "${CMAKE_ASM_FLAGS} ${CPU_FLAGS} ${CPU_DEFINES}") +set(CMAKE_EXE_LINKER_FLAGS + "${CMAKE_EXE_LINKER_FLAGS} ${CPU_FLAGS} -fno-common -ffunction-sections -fdata-sections -fno-builtin -mapcs -Wl,--gc-sections -Wl,-static -specs=nano.specs -specs=nosys.specs -T\"${SdkRootDirPath}/examples/_boards/mimxrt700evk/eiq_examples/executorch_cifarnet/cm33_core0/gcc/MIMXRT798Sxxxx_cm33_core0_flash.ld\" -static" +) + +add_executable( + ${MCUX_SDK_PROJECT_NAME} + ${SdkRootDirPath}/examples/_boards/mimxrt700evk/flash_config/flash_config.c + ${SdkRootDirPath}/examples/_boards/mimxrt700evk/eiq_examples/executorch_cifarnet/cm33_core0/hardware_init.c + ${SdkRootDirPath}/examples/_boards/mimxrt700evk/eiq_examples/executorch_cifarnet/cm33_core0/pin_mux.c + ${SdkRootDirPath}/examples/_boards/mimxrt700evk/board.c + ${SdkRootDirPath}/examples/_boards/mimxrt700evk/pmic_support.c + ${SdkRootDirPath}/examples/_boards/mimxrt700evk/common/clock/cm33_core0/clock_config.c + ${SdkRootDirPath}/examples/eiq_examples/executorch_cifarnet/main.cpp + ${SdkRootDirPath}/examples/eiq_examples/executorch_cifarnet/RegisterKernels.cpp + ${SdkRootDirPath}/examples/eiq_examples/common/timer.c + ${SdkRootDirPath}/devices/RT/RT700/MIMXRT798S/startup_MIMXRT798S_cm33_core0.c + ${EXECUTORCH_ROOT_DIR}/backends/nxp/runtime/NeutronBackend.cpp + ${SdkRootDirPath}/middleware/tfm/tf-m/platform/ext/common/syscalls_stub.c + # Device-level drivers (clock, power, reset, system init). + ${SdkRootDirPath}/devices/RT/RT700/MIMXRT798S/drivers/fsl_clock.c + ${SdkRootDirPath}/devices/RT/RT700/MIMXRT798S/drivers/fsl_power.c + ${SdkRootDirPath}/devices/RT/RT700/MIMXRT798S/drivers/fsl_reset.c + ${SdkRootDirPath}/devices/RT/RT700/MIMXRT798S/system_MIMXRT798S_cm33_core0.c + # Common ARM driver (provides SDK_DelayAtLeastUs). + ${SdkRootDirPath}/drivers/common/fsl_common_arm.c + # Peripheral drivers. + ${SdkRootDirPath}/drivers/cache/xcache/fsl_cache.c + ${SdkRootDirPath}/drivers/lpflexcomm/fsl_lpflexcomm.c + ${SdkRootDirPath}/drivers/lpflexcomm/lpi2c/fsl_lpi2c.c + ${SdkRootDirPath}/drivers/lpflexcomm/lpuart/fsl_lpuart.c + ${SdkRootDirPath}/drivers/gpio/fsl_gpio.c + ${SdkRootDirPath}/drivers/glikey/fsl_glikey.c + # UART HAL adapter (provides HAL_UartInit etc.). + ${SdkRootDirPath}/components/uart/fsl_adapter_lpuart.c + # PMIC driver. + ${SdkRootDirPath}/components/pmic/pca9422/fsl_pca9422.c + # Debug console (provides DbgConsole_Init/Printf). + ${SdkRootDirPath}/components/debug_console_lite/fsl_debug_console.c +) + +target_include_directories( + ${MCUX_SDK_PROJECT_NAME} + PRIVATE + ${MODEL_DIR} + ${SdkRootDirPath}/arch/arm/CMSIS/Core/Include + ${SdkRootDirPath}/devices/RT/RT700/MIMXRT798S + ${SdkRootDirPath}/devices/RT/RT700/MIMXRT798S/cm33_core0 + ${SdkRootDirPath}/devices/RT/RT700/MIMXRT798S/drivers + ${SdkRootDirPath}/devices/RT/RT700/periph + ${SdkRootDirPath}/drivers/common + ${SdkRootDirPath}/components/pmic/pca9422 + ${SdkRootDirPath}/components/uart + ${SdkRootDirPath}/components/debug_console_lite + ${SdkRootDirPath}/examples/_boards/mimxrt700evk + ${SdkRootDirPath}/examples/_boards/mimxrt700evk/flash_config + ${SdkRootDirPath}/examples/_boards/mimxrt700evk/eiq_examples/executorch_cifarnet/cm33_core0 + ${SdkRootDirPath}/examples/_boards/mimxrt700evk/common/clock/cm33_core0 + ${SdkRootDirPath}/examples/eiq_examples/executorch_cifarnet + ${SdkRootDirPath}/examples/eiq_examples/common + ${SdkRootDirPath}/examples/_boards/mimxrt700evk/eiq_examples/executorch_cifarnet/npu + ${SdkRootDirPath}/drivers/cache/xcache + ${SdkRootDirPath}/drivers/gpio + ${SdkRootDirPath}/drivers/lpflexcomm + ${SdkRootDirPath}/drivers/lpflexcomm/lpuart + ${SdkRootDirPath}/drivers/lpflexcomm/lpi2c + ${SdkRootDirPath}/drivers/xspi + ${SdkRootDirPath}/drivers/reset + ${SdkRootDirPath}/drivers/clock + ${SdkRootDirPath}/drivers/glikey + ${SdkRootDirPath}/drivers/mu1 + ${SdkRootDirPath}/drivers/power + ${SdkRootDirPath}/drivers/iopctl + ${SdkRootDirPath}/components/str +) + +set(EXECUTORCH_BUILD_PYBIND OFF) +set(EXECUTORCH_BUILD_TESTS OFF) +set(EXECUTORCH_BUILD_DEVTOOLS OFF) +set(EXECUTORCH_BUILD_EXECUTOR_RUNNER OFF) +set(EXECUTORCH_BUILD_CPUINFO OFF) +set(EXECUTORCH_BUILD_PTHREADPOOL OFF) +set(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL ON) +set(EXECUTORCH_BUILD_PORTABLE_OPS ON) +set(EXECUTORCH_BUILD_KERNELS_QUANTIZED ON) +set(CMAKE_POSITION_INDEPENDENT_CODE OFF) +add_subdirectory(${EXECUTORCH_ROOT_DIR} EXCLUDE_FROM_ALL executorch) + +target_link_libraries( + ${MCUX_SDK_PROJECT_NAME} + PRIVATE -Wl,--start-group + executorch + executorch_core + extension_runner_util + quantized_kernels + portable_kernels + ${NEUTRON_LIB_DIR}/libNeutronDriver.a + ${NEUTRON_LIB_DIR}/libNeutronFirmware.a + -Wl,--end-group +) diff --git a/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/build_example.sh b/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/build_example.sh new file mode 100755 index 00000000000..f8a528ac7ee --- /dev/null +++ b/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/build_example.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# Copyright 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. + +if [ -z ${ARMGCC_DIR+x} ]; then + echo "ARMGCC_DIR needs to be set in the environment!" + exit 1; +fi + +if [ -z ${SdkRootDirPath+x} ]; then + echo "SdkRootDirPath needs to be set in the environment!" + exit 1; +fi + +if [ ! -f model_pte.h ]; then + echo "Cannot find model_pte.h!" + exit 1; +fi + +if [ -z ${NEUTRON_LIB_DIR+x} ]; then + echo "NEUTRON_LIB_DIR needs to be set in the environment!" + exit 1; +fi + +if [ ! -f ${NEUTRON_LIB_DIR}/libNeutronDriver.a ]; then + echo "Neutron driver not found in ${NEUTRON_LIB_DIR}!" + exit 1; +fi + +if [ ! -f ${NEUTRON_LIB_DIR}/libNeutronFirmware.a ]; then + echo "Neutron firmware not found in ${NEUTRON_LIB_DIR}!" + exit 1; +fi + +rm -rf cmake-out && mkdir -p cmake-out + +cmake -DSdkRootDirPath=${SdkRootDirPath} \ + -DCMAKE_TOOLCHAIN_FILE=${SdkRootDirPath}/cmake/toolchain/armgcc.cmake \ + -DMODEL_DIR=$(pwd) \ + -DNEUTRON_LIB_DIR=${NEUTRON_LIB_DIR} \ + -DCMAKE_BUILD_TYPE=flash_release \ + -G "Unix Makefiles" \ + -B cmake-out \ + "$(dirname "$0")" + +make -C cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) executorch_cifarnet.elf diff --git a/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/prepare_model.sh b/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/prepare_model.sh new file mode 100755 index 00000000000..9a5c29022f5 --- /dev/null +++ b/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/prepare_model.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# Copyright 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. + +set -ue + +pushd "$(dirname "$0")/../../../../.." + +./install_executorch.sh +./devtools/install_requirements.sh + +pip install -r backends/nxp/requirements-eiq.txt + +python3 -m examples.nxp.aot_neutron_compile -m cifar10 -d -q --use_channels_last_dim_order --remove-quant-io-ops +mv cifar10_nxp_delegate.pte model.pte + +popd + +cat > model_pte.h <<'EOF' +#ifdef __MCUXPRESSO +#define __PLACEMENT __attribute__((section(".data.$modeldata"))) +#else +#define __PLACEMENT __attribute__((section(".modeldata"))) +#endif + +static const uint8_t model_pte[] __ALIGNED(16) __PLACEMENT = { +EOF + + +xxd -i "$(dirname "$0")/../../../../../model.pte" | grep -v unsigned >> model_pte.h diff --git a/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/test_build_from_scratch.sh b/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/test_build_from_scratch.sh new file mode 100755 index 00000000000..f734d1b79a1 --- /dev/null +++ b/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/test_build_from_scratch.sh @@ -0,0 +1,58 @@ +#!/bin/bash +# Copyright 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. + +set -ue +ARM_TOOLCHAIN_URL="${ARM_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}" + +# Get arm gcc +echo Downloading ARM GCC toolchain +if [ ! -d arm-toolchain ]; then + mkdir -p arm-toolchain + pushd arm-toolchain + wget $ARM_TOOLCHAIN_URL + tar -xvf *.tar.xz + rm *.tar.xz + popd +fi +export ARMGCC_DIR=$(pwd)/$(find arm-toolchain -maxdepth 1 -type d | tail -1) + +# Prepare model +# Side effect: the neutron SDK is installed +echo Preparing model and installing neutron SDK +$(dirname $0)/prepare_model.sh +# Check the model exists +if [ ! -f model_pte.h ]; then + echo "Cannot create the model_pte.h!" + exit 1; +fi + +# Locate Neutron SDK +NEUTRON_LIB_DIR=$(python3 -c "import eiq_neutron_sdk; print(eiq_neutron_sdk.__path__[0])") +export NEUTRON_LIB_DIR=${NEUTRON_LIB_DIR}/target/imxrt700/rt700/cm33 + +# Get MCUX SDK +echo Downloading MCUXpresso SDK +if [ ! -d mcuxpresso-sdk ]; then + pip install west + west init -m https://github.com/nxp-mcuxpresso/mcuxsdk-manifests.git mcuxpresso-sdk + pushd mcuxpresso-sdk + west update_board --set board mimxrt700evk + popd +fi +export SdkRootDirPath=$(pwd)/mcuxpresso-sdk/mcuxsdk + +# Build now +echo Building the example +$(dirname $0)/build_example.sh + +# Test the result +if [ ! -f cmake-out/flash_release/executorch_cifarnet.elf ]; then + echo "Build not successful!" + exit 1; +else + echo "Build successful." + exit 0; +fi diff --git a/examples/nxp/models/mlperf_tiny/image_classification/mlperf_tiny_image_classification.py b/examples/nxp/models/mlperf_tiny/image_classification/mlperf_tiny_image_classification.py index 60b590fa9a9..c8fdcd6226d 100644 --- a/examples/nxp/models/mlperf_tiny/image_classification/mlperf_tiny_image_classification.py +++ b/examples/nxp/models/mlperf_tiny/image_classification/mlperf_tiny_image_classification.py @@ -4,105 +4,49 @@ # LICENSE file in the root directory of this source tree. import logging -from pathlib import Path import torch -from executorch.backends.nxp.tests.calibration_dataset import ( - CalibrationDataset, - RandomCalibrationDataset, -) - from executorch.examples.models.mlperf_tiny import ResNet8 from executorch.examples.nxp.models.mlperf_tiny.mlperf_tiny_model import MLPerfTinyModel -from torch.utils.data import Dataset -from torchao.quantization.pt2e import disable_observer -from tqdm import tqdm log = logging.getLogger(__name__) -INPUT_SHAPE = (1, 3, 32, 32) -IDX_TO_LABEL = { - 0: "airplane", - 1: "automobile", - 2: "bird", - 3: "cat", - 4: "deer", - 5: "dog", - 6: "frog", - 7: "horse", - 8: "ship", - 9: "truck", -} - class MLPerfTinyImageClassification(MLPerfTinyModel): """MLPerf Tiny image classification model (ResNet-8).""" - def __init__( - self, - num_samples: int = 200, - dataset_path: Path | str | None = None, - use_random_dataset: bool = False, - ): - self._num_samples = num_samples - self._use_random_dataset = use_random_dataset - self._dataset_path = dataset_path - - super().__init__() + # ResNet-8 specific QAT training hyperparameters. + TRAIN_HYPERPARAMETERS = { + "num_epochs": 15, + "batch_size": 20, + "lr": 1e-5, + "eps": 1e-8, + "weight_decay": 1e-4, + } + + INPUT_SHAPE = (1, 3, 32, 32) + IDX_TO_LABEL = { + 0: "airplane", + 1: "automobile", + 2: "bird", + 3: "cat", + 4: "deer", + 5: "dog", + 6: "frog", + 7: "horse", + 8: "ship", + 9: "truck", + } @property def input_shape(self): - return INPUT_SHAPE + return self.INPUT_SHAPE @property def labels(self): - return IDX_TO_LABEL - - def _init_dataset(self) -> Dataset: - if self._use_random_dataset: - num_classes = len(self.labels) - sample_shape = tuple(self.input_shape)[1:] - return RandomCalibrationDataset( - self._num_samples, sample_shape, num_classes - ) - else: - if self._dataset_path is None: - raise ValueError( - "Path to dataset data cannot be empty. If you want to use random data, set `use_random_dataset = True`" - ) - return CalibrationDataset(self._dataset_path) + return self.IDX_TO_LABEL def _init_eager_model(self) -> torch.nn.Module: num_classes = len(self.labels) - return ResNet8(num_classes) - - def train_model_fn(self, model, num_epochs=15, batch_size=20, channels_last=False): - torch.manual_seed(42) - torch.use_deterministic_algorithms(True) - - optimizer = torch.optim.Adam( - params=model.parameters(), - lr=1e-5, - weight_decay=1e-4, - ) - loss_fn = torch.nn.CrossEntropyLoss() - - logging.warning("Starting training...") - - data = self.get_qat_train_inputs(batch_size=batch_size) - for nepoch in range(num_epochs): - for images, labels in tqdm(data): - if channels_last: - images = images.to(memory_format=torch.channels_last) - - optimizer.zero_grad() - outputs = model(images) - loss = loss_fn(outputs, labels) - loss.backward() - optimizer.step() - - if nepoch >= num_epochs / 3: - model.apply(disable_observer) - - return model + return ResNet8(num_classes).eval() diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/requirements.txt b/examples/nxp/models/mlperf_tiny/keyword_spotting/__init__.py similarity index 66% rename from examples/arm/mobilesam_prompt_segmentation_example_ethos_u/requirements.txt rename to examples/nxp/models/mlperf_tiny/keyword_spotting/__init__.py index 6eb309f00de..55dc5fccf45 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/requirements.txt +++ b/examples/nxp/models/mlperf_tiny/keyword_spotting/__init__.py @@ -1,6 +1,4 @@ -# Copyright 2026 Arm Limited and/or its affiliates. +# Copyright 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. - -tqdm == 4.67.1 diff --git a/examples/nxp/models/mlperf_tiny/keyword_spotting/mlperf_tiny_keyword_spotting.py b/examples/nxp/models/mlperf_tiny/keyword_spotting/mlperf_tiny_keyword_spotting.py new file mode 100644 index 00000000000..3bc16a5283d --- /dev/null +++ b/examples/nxp/models/mlperf_tiny/keyword_spotting/mlperf_tiny_keyword_spotting.py @@ -0,0 +1,60 @@ +# Copyright 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 logging + +import torch + +from executorch.examples.models.mlperf_tiny import DSCNNKWS +from executorch.examples.nxp.models.mlperf_tiny.mlperf_tiny_model import MLPerfTinyModel + +log = logging.getLogger(__name__) + + +class MLPerfTinyKeywordSpotting(MLPerfTinyModel): + """MLPerf Tiny keyword spotting model (DS-CNN).""" + + INPUT_SHAPE = (1, 1, 49, 10) + # Because of the architecture of the model, + # non-scaled random weights tend to produce zero tensors, + # making it hard to compute numerical accuracy of the delegated model. + # Scaling the random weights makes the model produce reasonable results. + WEIGHT_INIT_SCALE = 2.0 + + IDX_TO_LABEL = { + 0: "Down", + 1: "Go", + 2: "Left", + 3: "No", + 4: "Off", + 5: "On", + 6: "Right", + 7: "Stop", + 8: "Up", + 9: "Yes", + 10: "Silence", + 11: "Unknown", + } + + @property + def input_shape(self): + return self.INPUT_SHAPE + + @property + def labels(self): + return self.IDX_TO_LABEL + + def _init_weights(self, model: torch.nn.Module): + with torch.no_grad(): + for module in model.modules(): + if isinstance(module, (torch.nn.Conv2d, torch.nn.Linear)): + module.weight *= self.WEIGHT_INIT_SCALE + + def _init_eager_model(self) -> torch.nn.Module: + num_classes = len(self.labels) + model = DSCNNKWS(num_classes) + self._init_weights(model) + + return model.eval() diff --git a/examples/nxp/models/mlperf_tiny/mlperf_tiny_model.py b/examples/nxp/models/mlperf_tiny/mlperf_tiny_model.py index d72a2ad273b..b33a8b02057 100644 --- a/examples/nxp/models/mlperf_tiny/mlperf_tiny_model.py +++ b/examples/nxp/models/mlperf_tiny/mlperf_tiny_model.py @@ -4,20 +4,56 @@ # LICENSE file in the root directory of this source tree. import itertools +import logging +import os from abc import abstractmethod +from pathlib import Path from typing import Iterator import torch +from executorch.backends.nxp.tests.calibration_dataset import ( + CalibrationDataset, + RandomCalibrationDataset, +) from executorch.examples.models import model_base from torch.utils.data import DataLoader, Dataset +from torchao.quantization.pt2e import disable_observer +from tqdm import tqdm + +log = logging.getLogger(__name__) class MLPerfTinyModel(model_base.EagerModelBase): """Base class of the MLPerf Tiny models.""" - def __init__(self): - """Create the model wrapper along with the dataset it owns.""" - self._num_workers = 4 + # Default QAT training hyperparameters. Subclasses may override them. + TRAIN_HYPERPARAMETERS = { + "num_epochs": 15, + "batch_size": 64, + "lr": 5e-6, + "eps": 1e-7, + "weight_decay": 1e-4, + } + + def __init__( + self, + dataset_path: Path | str | None = None, + use_random_dataset: bool = False, + num_samples: int | None = None, + num_workers: int = 4, + ): + """ + Create the model wrapper along with the dataset it owns. + If `use_random_dataset = True`, then `num_samples` must be set. + If `use_random_dataset = False`, then `dataset_path` must be set. + """ + self._num_workers = num_workers + self._num_samples = num_samples + self._use_random_dataset = use_random_dataset + self._dataset_path = dataset_path + + # throws ValueError if validation fails + self._validate_args() self._eager_model = self._init_eager_model() self.dataset = self._init_dataset() @@ -27,10 +63,6 @@ def _collate_fn(data: list[tuple]): data, labels = zip(*data) return torch.stack(list(data)), torch.tensor(list(labels)) - @abstractmethod - def _init_dataset(self) -> Dataset: - pass - @abstractmethod def _init_eager_model(self) -> torch.nn.Module: pass @@ -45,6 +77,26 @@ def input_shape(self): def labels(self): pass + def _validate_args(self): + valid_num_samples = isinstance(self._num_samples, int) and self._num_samples > 0 + + if self._use_random_dataset: + if not valid_num_samples: + raise ValueError( + f"Invalid number of samples to randomly generate. Got {self._num_samples}." + ) + + else: + if valid_num_samples: + raise ValueError( + "Num samples was supplied, but it is omitted because `use_random_dataset=False`." + ) + + if self._dataset_path is None or not os.path.exists(self._dataset_path): + raise ValueError( + f"Invalid dataset path for loading the data. Got {self._dataset_path}." + ) + def get_qat_train_inputs( self, batch_size: int = 5, dataset_portion: float = 0.1 ) -> Iterator[tuple[torch.Tensor]]: @@ -79,3 +131,54 @@ def get_eager_model(self): def get_example_inputs(self) -> tuple[torch.Tensor]: return (torch.randn(self.input_shape, dtype=torch.float32),) + + def train_model_fn( + self, model, num_epochs=None, batch_size=None, channels_last=False + ): + hyperparameters = self.TRAIN_HYPERPARAMETERS + num_epochs = ( + num_epochs if num_epochs is not None else hyperparameters["num_epochs"] + ) + batch_size = ( + batch_size if batch_size is not None else hyperparameters["batch_size"] + ) + + torch.manual_seed(42) + torch.use_deterministic_algorithms(True) + + optimizer = torch.optim.Adam( + params=model.parameters(), + lr=hyperparameters["lr"], + eps=hyperparameters["eps"], + weight_decay=hyperparameters["weight_decay"], + ) + loss_fn = torch.nn.CrossEntropyLoss() + + log.warning("Starting training...") + + data = self.get_qat_train_inputs(batch_size=batch_size) + for nepoch in range(num_epochs): + for samples, labels in tqdm(data): + if channels_last: + samples = samples.to(memory_format=torch.channels_last) + + optimizer.zero_grad() + outputs = model(samples) + loss = loss_fn(outputs, labels) + loss.backward() + optimizer.step() + + if nepoch >= num_epochs / 3: + model.apply(disable_observer) + + return model + + def _init_dataset(self) -> Dataset: + if self._use_random_dataset: + num_classes = len(self.labels) + sample_shape = tuple(self.input_shape)[1:] + return RandomCalibrationDataset( + self._num_samples, sample_shape, num_classes + ) + else: + return CalibrationDataset(self._dataset_path) diff --git a/examples/nxp/run.sh b/examples/nxp/run.sh index b8cc87a5964..dc73bb474a2 100755 --- a/examples/nxp/run.sh +++ b/examples/nxp/run.sh @@ -19,7 +19,7 @@ rm -rf ${SCRIPT_DIR}/executor_runner/build/* pushd ${SCRIPT_DIR}/executor_runner/build cmake -DCMAKE_BUILD_TYPE=Release .. -make -j8 nxp_executor_runner +make -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) nxp_executor_runner popd echo "** Export cifar10 model to executorch" diff --git a/examples/portable/README.md b/examples/portable/README.md index ef9b44a48a3..f3b097881dd 100644 --- a/examples/portable/README.md +++ b/examples/portable/README.md @@ -49,7 +49,7 @@ Use `-h` (or `--help`) to see all the supported models. (mkdir cmake-out \ && cd cmake-out \ && cmake -DEXECUTORCH_PAL_DEFAULT=posix ..) \ - && cmake --build cmake-out -j32 --target executor_runner + && cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target executor_runner # Run the tool on the generated model. ./cmake-out/executor_runner --model_path mv2.pte diff --git a/examples/portable/custom_ops/test_custom_ops.sh b/examples/portable/custom_ops/test_custom_ops.sh index 58a7de3a5f2..a54761aefc7 100644 --- a/examples/portable/custom_ops/test_custom_ops.sh +++ b/examples/portable/custom_ops/test_custom_ops.sh @@ -30,7 +30,7 @@ test_cmake_custom_op_1() { ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config Release + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release echo 'Running custom_ops_executor_runner' ${build_dir}/custom_ops_executor_runner --model_path="./${model_name}.pte" @@ -66,7 +66,7 @@ test_cmake_custom_op_2() { ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config Release + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release EXT=$(get_shared_lib_ext) echo "Exporting ${model_name}.pte" diff --git a/examples/portable/scripts/test_demo_backend_delegation.sh b/examples/portable/scripts/test_demo_backend_delegation.sh index d1ecf9150f9..2aeb1f90689 100644 --- a/examples/portable/scripts/test_demo_backend_delegation.sh +++ b/examples/portable/scripts/test_demo_backend_delegation.sh @@ -24,7 +24,7 @@ build_cmake_executor_runner() { && cd ${CMAKE_OUTPUT_DIR} \ && retry cmake -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" ..) - cmake --build ${CMAKE_OUTPUT_DIR} -j4 + cmake --build ${CMAKE_OUTPUT_DIR} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) } test_demo_backend_delegation() { diff --git a/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py b/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py index 29eba63a07d..521dba38ef7 100644 --- a/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py +++ b/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py @@ -697,11 +697,7 @@ def __init__(self, verbose: bool = False): class Smollm2QATQuantRecipe(StaticLLMQATRecipe): - default_quant_dtype = QuantDtype.use_16a8w - frozen_param_patterns: List[str] = [ - r"tok_embedding", # Freeze token embeddings to prevent drift in the token space. - r"output\.conv", # Freeze lm head to prevent drift in the token space. - ] + default_quant_dtype = QuantDtype.use_16a4w def __init__(self, verbose: bool = False): super().__init__() @@ -725,14 +721,7 @@ def __init__(self, verbose: bool = False): ) .add_regex( {r"tok_embeddings"}, - QuantDtype.use_16a8w, - True, - act_observer=MovingAverageMinMaxObserver, - granularity=QuantGranularity.PER_TENSOR, - ) - .add_regex( - {r"output\.conv"}, - QuantDtype.use_16a8w, + QuantDtype.use_16a4w, True, act_observer=MovingAverageMinMaxObserver, granularity=QuantGranularity.PER_CHANNEL, diff --git a/examples/qualcomm/test_qualcomm.sh b/examples/qualcomm/test_qualcomm.sh index 51a563863f3..24a9c224061 100644 --- a/examples/qualcomm/test_qualcomm.sh +++ b/examples/qualcomm/test_qualcomm.sh @@ -21,7 +21,7 @@ cmake_install_executorch_qnn_lib() { -DEXECUTORCH_BUILD_QNN=ON \ -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \ -Bcmake-out . - cmake --build cmake-out -j9 --target install --config Release + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install --config Release } test_cmake_qualcomm() { @@ -45,7 +45,7 @@ test_cmake_qualcomm() { ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config Release + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release # Need to run on device # ${build_dir}/qnn_executor_runner --model_path="./mv2_qnn.pte" } diff --git a/examples/riscv/README.md b/examples/riscv/README.md index 563ff4913fd..02a55351a8a 100644 --- a/examples/riscv/README.md +++ b/examples/riscv/README.md @@ -35,7 +35,7 @@ The driver does three steps: ## CI -`.github/workflows/_test_riscv_qemu.yml` is a reusable `workflow_call` -job (mirroring `_test_cortex_m_e2e.yml`) invoked from `pull.yml` to run on -every PR. It runs on the standard `linux.2xlarge` x86_64 runner using the -`executorch-ubuntu-22.04-gcc11` docker image. +`.github/workflows/_test_riscv.yml` is a reusable `workflow_call` +job (mirroring `_test_cortex_m_e2e.yml`) invoked from `riscv64.yml`. It runs on +the `mt-l-x86iavx512-8-64` x86_64 runner using the +`executorch-ubuntu-24.04-gcc14` docker image. diff --git a/examples/samsung/CMakeLists.txt b/examples/samsung/CMakeLists.txt new file mode 100644 index 00000000000..28a6dcc6660 --- /dev/null +++ b/examples/samsung/CMakeLists.txt @@ -0,0 +1,61 @@ +# Copyright (c) 2025 Samsung Electronics Co. LTD +# 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. + +cmake_minimum_required(VERSION 3.15) +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +project(samsung_example) + +# Source root directory for executorch. +if(NOT EXECUTORCH_ROOT) + set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..) +endif() + +include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) +include(${EXECUTORCH_ROOT}/tools/cmake/Codegen.cmake) + +if(NOT PYTHON_EXECUTABLE) + resolve_python_executable() +endif() + +if(NOT CMAKE_BUILD_TYPE) + set(CMAKE_BUILD_TYPE RelWithDebInfo) +endif() + +if(CMAKE_TOOLCHAIN_FILE MATCHES ".*(iOS|ios\.toolchain)\.cmake$") + message(FATAL_ERROR "ios is not supported by Samsung AI system.") +endif() + +# prebuilt libraries. executorch package should contain portable_ops_lib, +# etdump, bundled_program. +find_package(executorch CONFIG REQUIRED) +target_compile_options(executorch INTERFACE -DET_EVENT_TRACER_ENABLED) +find_package(gflags REQUIRED) + +set(_common_compile_options -Wno-deprecated-declarations -fPIC) + +add_compile_options(-Wall -Werror -fPIC) + +message("Build Samsung Android Examples") + +set(__enn_executor_runner__srcs + ${CMAKE_CURRENT_LIST_DIR}/executor_runner/enn_executor_runner.cpp +) + +add_executable(enn_executor_runner ${__enn_executor_runner__srcs}) + +target_include_directories(enn_executor_runner PRIVATE ${EXECUTORCH_ROOT}/..) + +target_compile_options(enn_executor_runner PRIVATE ${_common_compile_options}) + +target_link_libraries( + enn_executor_runner PRIVATE enn_logging enn_backend gflags executorch + extension_data_loader portable_ops_lib +) + +set_target_properties( + enn_executor_runner PROPERTIES CXX_VISIBILITY_PRESET hidden +) diff --git a/examples/samsung/README.md b/examples/samsung/README.md index 8b21c48a34f..a64169bd9a1 100644 --- a/examples/samsung/README.md +++ b/examples/samsung/README.md @@ -1,9 +1,9 @@ -# Exynos backend Examples +# Exynos Backend examples This directory contains examples for some AI models. -Please make sure you have built the library and executable before -you start, if you have no idea how to build, please refer to [backend README](../../backends/samsung/README.md). +Please make sure you have built the library before you start, +if you have no idea how to build, please refer to [backend README](../../backends/samsung/README.md). ## Environment We set up `PYTHONPATH` because it's easier to develop and import executorch Python APIs. @@ -42,15 +42,26 @@ Examples use "PerformanceMode.HIGH_PERFORMANCE" mode, this mode is experimental. If you want to use this mode on your model, verify your model on devicefarm which can use samsung developer society site firstly for checking stability. (https://soc-developer.semiconductor.samsung.com/) +## Building Executable +### Prerequisites +Please set up the backend before building the executable, See the [backend README](../../backends/samsung/README.md) for details. +### Building 'enn_executor_runner' for Android +```bash +export EXYNOS_AI_LITECORE_ROOT=/path/to/enn_sdk +export ANDROID_NDK_ROOT=/path/to/android_ndk +${EXECUTORCH_ROOT}/examples/samsung/build.sh +``` +After the build completes, `enn_executor_runner` can be found at `${EXECUTORCH_ROOT}/build_samsung_android/examples/samsung/` ## Execution + After lowering, we could get a pte model and then run it on mobile phone. #### Step 1: Push required ENN libraries and executor runner to device ```bash DEVICE_DIR=/data/local/tmp/executorch adb shell mkdir ${DEVICE_DIR} -adb push ${EXECUTORCH_ROOT}/cmake-android-out/backends/samsung/enn_executor_runner ${DEVICE_DIR} +adb push ${EXECUTORCH_ROOT}/build_samsung_android/examples/samsung/enn_executor_runner ${DEVICE_DIR} ``` #### Step 2: Indicate dynamic linkers and execute model diff --git a/examples/samsung/build.sh b/examples/samsung/build.sh new file mode 100755 index 00000000000..4568180777a --- /dev/null +++ b/examples/samsung/build.sh @@ -0,0 +1,45 @@ +#!/bin/bash + +set -e + +BASE_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" +PROJECT_DIR=$(realpath ${BASE_DIR}/../../) + +echo PROJECT_DIR=${PROJECT_DIR} + +if [[ -z ${ANDROID_NDK_ROOT} ]]; then + echo "Please export ANDROID_NDK_ROOT" + exit 1 +fi + +ANDROID_ABI=arm64-v8a +ANDROID_PLATFORM=android-28 # Trace requires over android-23 + +echo ANDROID_NDK_ROOT=${ANDROID_NDK_ROOT} +echo ANDROID_ABI=${ANDROID_ABI} +echo ANDROID_PLATFORM=${ANDROID_PLATFORM} + +main() { + cd "$PROJECT_DIR" + local build_dir_root="build_samsung_android" + local example_root="examples/samsung" + local build_dir_example="$PROJECT_DIR/${build_dir_root}/${example_root}" + local cmake_prefix_path="$PROJECT_DIR/${build_dir_root}/lib/cmake/ExecuTorch;$PROJECT_DIR/${build_dir_root}/third-party/gflags;$PROJECT_DIR/${build_dir_root}/lib/cmake/tokenizers;$PROJECT_DIR/${build_dir_root}/lib/cmake/re2;$PROJECT_DIR/${build_dir_root}/lib/cmake/absl;" + + echo build_dir=${build_dir_root} + echo build_dir_example=${build_dir_example} + echo cmake_prefix_path=${cmake_prefix_path} + + cmake -DCMAKE_PREFIX_PATH=${cmake_prefix_path} \ + -DCMAKE_TOOLCHAIN_FILE="$ANDROID_NDK_ROOT/build/cmake/android.toolchain.cmake" \ + -DANDROID_NDK=$ANDROID_NDK \ + -DANDROID_ABI="$ANDROID_ABI" \ + -DANDROID_PLATFORM=$ANDROID_PLATFORM \ + -DCMAKE_FIND_ROOT_PATH_MODE_PACKAGE=BOTH \ + -DCMAKE_BUILD_TYPE=Release \ + -B"${build_dir_example}" \ + "$PROJECT_DIR/${example_root}" + cmake --build build_samsung_android/examples/samsung/ --config Release +} + +main "$@" diff --git a/examples/samsung/executor_runner/enn_executor_runner.cpp b/examples/samsung/executor_runner/enn_executor_runner.cpp index de168be9d7c..14dfbf58ee2 100644 --- a/examples/samsung/executor_runner/enn_executor_runner.cpp +++ b/examples/samsung/executor_runner/enn_executor_runner.cpp @@ -17,8 +17,8 @@ */ #include +#include #include -#include #include #include #include @@ -27,6 +27,7 @@ #include #include +#include #include #include #include @@ -45,8 +46,8 @@ DEFINE_bool(dump_statistics, false, "Dump inference statistics."); DEFINE_string(output_path, "", "Output Execution results to target directory."); using namespace torch::executor; -using torch::executor::util::FileDataLoader; using namespace torch::executor::enn; +using executorch::backends::enn::ExynosFileDataLoader; std::vector split(std::string str, char delimiter = ' ') { std::vector result; @@ -96,10 +97,26 @@ class DataReader { return data_set_[index].size(); } + // Some ops rewrite their input buffer in place during execute(), so a + // repeated execution would otherwise run on whatever the previous + // execution left behind. Snapshot the pristine bytes once inputs are set, + // then restore() before every execution. + void snapshot() { + pristine_ = data_set_; + } + + void restore() { + for (size_t i = 0; i < data_set_.size(); ++i) { + std::copy( + pristine_[i].cbegin(), pristine_[i].cend(), data_set_[i].begin()); + } + } + ~DataReader() = default; private: std::vector data_set_; + std::vector pristine_; int32_t index_ = 0; }; @@ -118,27 +135,11 @@ void saveOutput(const exec_aten::Tensor& tensor, int32_t output_index) { fout.close(); } -struct EnnApiDeinit { - void operator()(EnnApi* ptr) const { - if (ptr == nullptr) { - return; - } - - auto ret = ptr->EnnDeinitialize(); - ET_CHECK_MSG(ret == ENN_RET_SUCCESS, "Enn Deinitialize failed."); - } -}; - -std::unique_ptr exynos_npu_init() { - EnnApi* enn_api_inst = EnnApi::getEnnApiInstance(); - auto ret = enn_api_inst->EnnInitialize(); - ET_CHECK_MSG(ret == ENN_RET_SUCCESS, "Enn initialize failed."); - return std::unique_ptr(enn_api_inst); -} - int main(int argc, char** argv) { auto before_init = std::chrono::high_resolution_clock::now(); - std::unique_ptr instance = exynos_npu_init(); + // The EnnApi singleton initializes the NPU on construction and deinitializes + // it on process teardown. + EnnApi::getEnnApiInstance(); auto after_init = std::chrono::high_resolution_clock::now(); double interval_init = std::chrono::duration_cast( after_init - before_init) @@ -159,10 +160,10 @@ int main(int argc, char** argv) { // DataLoaders that use mmap() or point to data that's already in memory, and // users can create their own DataLoaders to load from arbitrary sources. const char* model_path = FLAGS_model.c_str(); - Result loader = FileDataLoader::from(model_path); + Result loader = ExynosFileDataLoader::from(model_path); ET_CHECK_MSG( loader.ok(), - "FileDataLoader::from() failed: 0x%" PRIx32, + "ExynosFileDataLoader::from() failed: 0x%" PRIx32, (uint32_t)loader.error()); // Parse the program file. This is immutable, and can also be reused between @@ -312,46 +313,37 @@ int main(int argc, char** argv) { ET_CHECK_MSG(ret == Error::Ok, "Failed to set input tensor: %d", ret); } EXYNOS_ATRACE_END(); + input_data_reader.snapshot(); // Warm up ET_LOG(Info, "Perform %d inference for warming up", FLAGS_warm_up); Error status; for (int i = 0; i < FLAGS_warm_up; ++i) { + input_data_reader.restore(); status = method->execute(); } - // Run the model. - ET_LOG(Info, "Start 1st inference."); - auto before_exec = std::chrono::high_resolution_clock::now(); - status = method->execute(); - auto after_exec = std::chrono::high_resolution_clock::now(); - double interval_1st_infs = - std::chrono::duration_cast( - after_exec - before_exec) - .count() / - 1000.0; - ET_LOG(Info, "Start inference."); - before_exec = std::chrono::high_resolution_clock::now(); + std::chrono::microseconds infs_duration{0}; for (int i = 0; i < FLAGS_num_executions; ++i) { + // Restored outside the timed section so it measures execute() alone, + // not the cost of undoing the previous iteration's input mutation. + input_data_reader.restore(); + auto before_exec = std::chrono::high_resolution_clock::now(); status = method->execute(); + auto after_exec = std::chrono::high_resolution_clock::now(); + infs_duration += std::chrono::duration_cast( + after_exec - before_exec); } - after_exec = std::chrono::high_resolution_clock::now(); - double interval_infs = std::chrono::duration_cast( - after_exec - before_exec) - .count() / - 1000.0; + double interval_infs = infs_duration.count() / 1000.0; if (FLAGS_dump_statistics) { auto output_file_name = "statistics.txt"; std::ofstream fout(output_file_name); fout << "init: " + std::to_string(interval_init) << "\nload: " + std::to_string(interval_load) - << "\n1st: " + std::to_string(interval_1st_infs) << "\navg: " + - std::to_string( - (interval_infs + interval_1st_infs) / - ((float)FLAGS_num_executions + 1.f)) + std::to_string(interval_infs / (float)FLAGS_num_executions) << std::endl; fout.close(); } diff --git a/examples/samsung/scripts/mobilebert_finetune.py b/examples/samsung/scripts/mobilebert_finetune.py index 76c1e9b03d3..bee2909ffb2 100644 --- a/examples/samsung/scripts/mobilebert_finetune.py +++ b/examples/samsung/scripts/mobilebert_finetune.py @@ -117,7 +117,7 @@ def build_loader_from_dataset(self, dataset, batch_size, usage="train"): return data_loader - def get_finetune_mobilebert(self, artifacts_dir): + def get_finetune_mobilebert(self, artifacts_dir, batch_size=64): # Pretrained bert's output ranges in a large scale. It is challenge for enn backend to support directly. # Please finetune mobilebert on specific tasks, make sure that bert's output and hidden states are friendly # to resource-constraint device. @@ -138,7 +138,7 @@ def get_finetune_mobilebert(self, artifacts_dir): labels_set = train_data.label.unique() train_data_loader = self.build_loader_from_dataset( - train_data, batch_size=64, usage="train" + train_data, batch_size=batch_size, usage="train" ) val_url = "https://raw.githubusercontent.com/clairett/pytorch-sentiment-classification/refs/heads/master/data/SST2/test.tsv" @@ -147,7 +147,7 @@ def get_finetune_mobilebert(self, artifacts_dir): BytesIO(content), delimiter="\t", header=None, names=["text", "label"] ) val_data_loader = self.build_loader_from_dataset( - val_data, batch_size=64, usage="val" + val_data, batch_size=batch_size, usage="val" ) artifacts_dir = artifacts_dir if artifacts_dir is not None else "./mobilebert" diff --git a/examples/selective_build/README.md b/examples/selective_build/README.md index c6c8dc1ba57..1e1e556f415 100644 --- a/examples/selective_build/README.md +++ b/examples/selective_build/README.md @@ -18,7 +18,7 @@ python -m examples.portable.scripts.export --model_name="mv2" # Create a PTE fil cd examples/selective_build/basic mkdir cmake-out && cd cmake-out cmake .. -DEXECUTORCH_SELECT_OPS_MODEL="../../mv2.pte" # Build with kernels needed for mv2.pte -cmake --build . -j8 +cmake --build . -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) ./selective_build_test --model_path="../../mv2.pte" # Run the model with the selective kernel library ``` @@ -78,7 +78,7 @@ python -m examples.portable.custom_ops.custom_ops_1 # Create a model PTE file cd examples/selective_build/basic mkdir cmake-out && cd cmake-out cmake .. -DEXECUTORCH_SELECT_OPS_MODEL="../../custom_ops_1.pte" -DEXECUTORCH_EXAMPLE_USE_CUSTOM_OPS=ON # Build with kernels needed for the model -cmake --build . -j8 +cmake --build . -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) ./selective_build_test --model_path="../../custom_ops_1.pte" # Run the model with the selective kernel library ``` diff --git a/examples/selective_build/test_selective_build.sh b/examples/selective_build/test_selective_build.sh index c1b5c627c42..bd286f8db19 100755 --- a/examples/selective_build/test_selective_build.sh +++ b/examples/selective_build/test_selective_build.sh @@ -105,7 +105,7 @@ aten,aten::clone.out" \ ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config Release + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release echo 'Running selective build test' ${build_dir}/selective_build_test --model_path="./mv2.pte" @@ -129,7 +129,7 @@ test_cmake_select_ops_in_yaml() { ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config Release + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release echo 'Running selective build test' ${build_dir}/selective_build_test --model_path="./custom_ops_1.pte" @@ -158,7 +158,7 @@ test_cmake_select_ops_in_model() { ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config $CMAKE_BUILD_TYPE + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config $CMAKE_BUILD_TYPE echo "Verifying auto-right-sized MAX_KERNEL_NUM header was generated" local generated_header diff --git a/examples/wasm/README.md b/examples/wasm/README.md index 15ce07493d1..716e1115846 100644 --- a/examples/wasm/README.md +++ b/examples/wasm/README.md @@ -51,13 +51,13 @@ Use -h (or --help) to see all the supported models. For the browser example, mak (mkdir cmake-out-wasm \ && cd cmake-out-wasm \ && emcmake cmake -DEXECUTORCH_PAL_DEFAULT=posix ..) \ - && cmake --build cmake-out-wasm -j32 --target executor_runner + && cmake --build cmake-out-wasm -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target executor_runner ``` If you need to rebuild `executor_runner` after modifying the contents of `./models/`, you can run the following command ```bash -cmake --build cmake-out-wasm -j32 --target executor_runner --clean-first +cmake --build cmake-out-wasm -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target executor_runner --clean-first ``` 4. Run the model with Node.js. Emscripten should come preinstalled with a compatible version of Node.js. If you have an incompatible version of Node.js installed, you can use the Emscripten-provided version by running `$EMSDK_NODE` instead of `node`. @@ -91,7 +91,7 @@ echo $EMSDK_NODE The file may not have been present while building the Wasm binary. You can rebuild with the following command ```bash -cmake --build cmake-out-wasm -j32 --target executor_runner --clean-first +cmake --build cmake-out-wasm -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target executor_runner --clean-first ``` The path may also be incorrect. The files in the `WASM_MODEL_DIR` are placed into the root directory of the virtual file system, so you would use `--model_path mv2.pte` instead of `--model_path models/mv2.pte`, for example. diff --git a/examples/wasm/test_build_wasm.sh b/examples/wasm/test_build_wasm.sh index f7144a209df..ef836b45901 100644 --- a/examples/wasm/test_build_wasm.sh +++ b/examples/wasm/test_build_wasm.sh @@ -23,7 +23,7 @@ test_build_wasm() { retry emcmake cmake -DWASM_MODEL_DIR="$(realpath "${model_dir_name}")" -B${build_dir} . echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --target executor_runner + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target executor_runner echo "Removing ${model_dir_name}" rm -rf "${model_dir_name}" diff --git a/examples/xnnpack/quantization/test_quantize.sh b/examples/xnnpack/quantization/test_quantize.sh index 1f50667c788..1211470b084 100644 --- a/examples/xnnpack/quantization/test_quantize.sh +++ b/examples/xnnpack/quantization/test_quantize.sh @@ -56,7 +56,7 @@ test_cmake_quantization() { -DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON \ -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" ..) - cmake --build cmake-out -j4 + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) EXT=$(get_shared_lib_ext) SO_LIB="cmake-out/kernels/quantized/libquantized_ops_aot_lib$EXT" diff --git a/exir/backend/test/demos/BUCK b/exir/backend/test/demos/BUCK index 8404c5982c6..b022a9f8fcc 100644 --- a/exir/backend/test/demos/BUCK +++ b/exir/backend/test/demos/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") +oncall("executorch") + fbcode_target(_kind = python_unittest, name = "test_delegate_aten_mode", srcs = [ diff --git a/exir/lowered_backend_module.py b/exir/lowered_backend_module.py index a4c2f2cfe79..358bcb25e8c 100644 --- a/exir/lowered_backend_module.py +++ b/exir/lowered_backend_module.py @@ -384,7 +384,7 @@ def arrange_graph_placeholders( ) -> torch.fx.GraphModule: """ Modifies the graph of the given graphmodule with one that contains the same nodes as the original, - but with placeholders in order of (Params + Buffers) (User Inputs) + but with placeholders in order of (Params + Buffers + Constants) (User Inputs) This is used by the delegate api which disturbs the placeholder ordering when creating a submodule from partitioned nodes @@ -403,32 +403,30 @@ def arrange_graph_placeholders( graph_sign = owning_program.graph_signature # Add all placeholders into the graph first: - # Cache these properties — each call rebuilds the dict from input_specs. + # Cache these properties to avoid rebuilding the dict on each access. params_map = graph_sign.inputs_to_parameters buffers_map = graph_sign.inputs_to_buffers + constants_map = graph_sign.inputs_to_lifted_tensor_constants param_nodes = [] buffer_nodes = [] + constant_nodes = [] input_nodes = [] for node in gm.graph.nodes: if node.op != "placeholder": continue - if node.name in params_map and node.meta.get("delegation_tag", None) == tag: + is_tagged = node.meta.get("delegation_tag", None) == tag + if node.name in params_map and is_tagged: param_nodes.append(node) - elif node.name in buffers_map and node.meta.get("delegation_tag", None) == tag: + elif node.name in buffers_map and is_tagged: buffer_nodes.append(node) + elif node.name in constants_map and is_tagged: + constant_nodes.append(node) else: input_nodes.append(node) - for param_node in param_nodes: - new_node = new_graph.node_copy(param_node, lambda x: node_map[x]) - node_map[param_node] = new_node - for buffer_node in buffer_nodes: - new_node = new_graph.node_copy(buffer_node, lambda x: node_map[x]) - node_map[buffer_node] = new_node - for input_node in input_nodes: - new_node = new_graph.node_copy(input_node, lambda x: node_map[x]) - node_map[input_node] = new_node + for node in param_nodes + buffer_nodes + constant_nodes + input_nodes: + node_map[node] = new_graph.node_copy(node, lambda x: node_map[x]) # Now add all the other nodes in order for node in gm.graph.nodes: diff --git a/exir/pass_manager.py b/exir/pass_manager.py index 829486fa0ce..97c27901d73 100644 --- a/exir/pass_manager.py +++ b/exir/pass_manager.py @@ -1,12 +1,12 @@ # 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. # pyre-strict import copy -import inspect import logging import operator from typing import Callable, List, Optional, Type, TypeAlias, Union @@ -36,7 +36,9 @@ def _get_pass_name(fn: PassType) -> str: """Returns a human-readable name for a pass.""" - return fn.__name__ if inspect.isfunction(fn) else type(fn).__name__ + if hasattr(fn, "__name__"): + return fn.__name__ + return type(fn).__name__ def _can_eliminate_common_getitems(gm: torch.fx.GraphModule) -> bool: diff --git a/exir/passes/constant_prop_pass.py b/exir/passes/constant_prop_pass.py index 11640d875c0..ea8ee1ad3a9 100644 --- a/exir/passes/constant_prop_pass.py +++ b/exir/passes/constant_prop_pass.py @@ -146,6 +146,10 @@ def get_propagated_const_tensor_dict( node.op != "call_function" or node.target is memory.alloc or node.target in all_skip_targets + # Ops with side effects (RNG draws, mutation) have to run at + # runtime. `aten.rand` has no tensor inputs, so without this check + # it would be folded into a single frozen draw. + or node.is_impure() ): continue diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 0b586bd44cd..54211a490a3 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -2475,6 +2475,31 @@ def forward(self, x): # 1 constant: a (= self.w @ self.cst) self.assertEqual(1, len(pass_result.constants)) + def test_constant_prop_pass_skips_nondeterministic_ops(self) -> None: + """ + Ops that draw from the RNG take no tensor inputs, so they look constant + to the pass. They have to stay in the graph: folding one would freeze a + single random draw into the program. + """ + + class RandomAdd(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + torch.rand(4) + + x = torch.zeros(4) + edge = to_edge(export(RandomAdd(), (x,), strict=True)) + new_ep = constant_prop_pass(edge.exported_program()) + + rand_nodes = [ + node + for node in new_ep.graph.nodes + if node.target == exir_ops.edge.aten.rand.default + ] + self.assertEqual(len(rand_nodes), 1) + self.assertEqual(len(new_ep.constants), 0) + module = new_ep.module() + self.assertFalse(torch.equal(module(x), module(x))) + def test_constant_prop_pass_zero_stride_tensors(self) -> None: """ Test that constant propagation correctly handles tensors with zero strides diff --git a/export/export.py b/export/export.py index f569a4196ee..fa8e534a430 100644 --- a/export/export.py +++ b/export/export.py @@ -283,13 +283,17 @@ def _get_default_pipeline(self) -> List[StageType]: if self._input_model_type != "ExportedProgram": stages.append(StageType.TORCH_EXPORT) - # Always include edge and executorch stages - stages.extend( - [ - StageType.TO_EDGE_TRANSFORM_AND_LOWER, - StageType.TO_EXECUTORCH, - ] - ) + stages.append(StageType.TO_EDGE_TRANSFORM_AND_LOWER) + + # This is the only stage that runs edge_manager_transform_passes, so a + # recipe declaring them would otherwise have them silently dropped. + if ( + self._lowering_recipe + and self._lowering_recipe.edge_manager_transform_passes + ): + stages.append(StageType.EDGE_PROGRAM_MANAGER_TRANSFORM) + + stages.append(StageType.TO_EXECUTORCH) return stages diff --git a/export/recipe.py b/export/recipe.py index 1609b989273..0d244b76556 100644 --- a/export/recipe.py +++ b/export/recipe.py @@ -1,15 +1,18 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # Copyright 2025 Arm Limited and/or its affiliates. +# Copyright 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 copy import dataclasses +import logging from abc import ABCMeta, abstractmethod -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum, EnumMeta -from typing import Callable, Dict, List, Optional, Union +from typing import Callable, Dict, Iterable, List, Optional, Union import torch from executorch.exir import EdgeProgramManager, ExportedProgram @@ -129,16 +132,54 @@ class QuantizationRecipe: """ Configuration recipe for quantization. - This class holds the configuration parameters for quantizing a model. + This class holds the configuration parameters for quantizing a model, supporting + both post-training quantization (PTQ) and quantization-aware training (QAT). Attributes: - quantizers: Optional list of quantizers for model quantization + quantizers: Optional list of quantizers for model quantization. ao_quantization_configs: Optional list of AOQuantizationConfig objects that pair - AOBaseConfig with optional filter functions + AOBaseConfig with optional filter functions. + is_qat: If True, use the QAT flow (prepare_qat_pt2e -> train_fn -> convert_pt2e). + If False (default), use the PTQ flow (prepare_pt2e -> calibrate -> convert_pt2e). + dynamic_batch_size: If True, dimension 0 (batch) of the calibration/QAT + training inputs may vary. Otherwise it is fixed to the + example inputs' batch size. + calibration_inputs_fn: Optional callable returning an iterable of input tuples used for + PTQ calibration. When None (default), the example inputs are used. + Ignored when is_qat=True. + train_fn: Callable that receives the prepared GraphModule and trains it. + Required when is_qat=True; ignored otherwise. + pre_prepare_passes: Optional list of callables applied to the captured GraphModule + before prepare_pt2e / prepare_qat_pt2e. + Each callable receives a GraphModule and must return a GraphModule. + post_prepare_passes: Optional list of callables applied to the prepared GraphModule + after prepare_pt2e / prepare_qat_pt2e and before calibration / training. + Each callable receives a GraphModule and must return a GraphModule. + pre_convert_passes: Optional list of callables applied to the GraphModule after + calibration (PTQ) or training (QAT) and before convert_pt2e. + Each callable receives a GraphModule and must return a GraphModule. + post_convert_passes: Optional list of callables applied to the GraphModule after convert_pt2e. + Each callable receives a GraphModule and must return a GraphModule. """ quantizers: Optional[List[Quantizer]] = None ao_quantization_configs: Optional[List[AOQuantizationConfig]] = None + is_qat: bool = False + dynamic_batch_size: bool = False + calibration_inputs_fn: Optional[Callable[[], Iterable[tuple]]] = None + train_fn: Optional[Callable[["torch.fx.GraphModule"], None]] = None + pre_prepare_passes: Optional[ + List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]] + ] = None + post_prepare_passes: Optional[ + List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]] + ] = None + pre_convert_passes: Optional[ + List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]] + ] = None + post_convert_passes: Optional[ + List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]] + ] = None def get_quantizers(self) -> Optional[List[Quantizer]]: """ @@ -168,6 +209,8 @@ class LoweringRecipe: edge_manager_transform_passes: Optional list of callables that take EdgeProgramManager as argument and return passes to be applied. Applied sequentially after TO_EDGE stage. edge_compile_config: Optional edge compilation configuration + pre_partitioning_callback: Optional callable invoked just before partitioning with + `(partitioners, programs)` arguments. """ partitioners: Optional[Union[List[Partitioner], Dict[str, List[Partitioner]]]] = ( @@ -182,6 +225,36 @@ class LoweringRecipe: ) = None # pyre-ignore[11]: Type not defined edge_compile_config: Optional[EdgeCompileConfig] = None + pre_partitioning_callback: Optional[ + Callable[[Optional[list[Partitioner]], dict[str, ExportedProgram]], None] + ] = None + + +@dataclass +class _CombineAccumulator: + """Private accumulator used by ExportRecipe._collect_recipe_fields.""" + + partitioners: list = field(default_factory=list) + partitioners_by_method: dict = field(default_factory=dict) + quantizers: list = field(default_factory=list) + ao_quantization_configs: list = field(default_factory=list) + pre_edge_passes: list = field(default_factory=list) + edge_transform_passes: list = field(default_factory=list) + edge_manager_transform_passes: list = field(default_factory=list) + pre_prepare_passes: list = field(default_factory=list) + post_prepare_passes: list = field(default_factory=list) + pre_convert_passes: list = field(default_factory=list) + post_convert_passes: list = field(default_factory=list) + is_qat_values: list = field(default_factory=list) + dynamic_batch_size_values: list = field(default_factory=list) + train_fn_values: list = field(default_factory=list) + calibration_inputs_fn_values: list = field(default_factory=list) + strict_values: list = field(default_factory=list) + mode_values: list = field(default_factory=list) + pipeline_stages_values: list = field(default_factory=list) + source_transform_in_place_values: list = field(default_factory=list) + backend_config: object = None + pre_partitioning_callbacks: list = field(default_factory=list) @experimental( @@ -281,89 +354,103 @@ def combine( return cls._combine_recipes(recipes, recipe_name) + @staticmethod + def _assert_scalar_fields_agree(field_name: str, values: list) -> None: + """Raise ValueError when a scalar field has conflicting values across recipes.""" + unique = set(values) + if len(unique) > 1: + raise ValueError( + f"Cannot combine recipes with conflicting '{field_name}' values: {unique}" + ) + @classmethod - def _combine_recipes( # noqa: C901 - cls, backend_recipes: List["ExportRecipe"], recipe_name: Optional[str] = None - ) -> "ExportRecipe": + def _combine_quantization_recipe( + cls, + is_qat_values: list, + dynamic_batch_size_values: list, + train_fn_values: list, + calibration_inputs_fn_values: list, + all_quantizers: list, + all_ao_quantization_configs: list, + all_pre_prepare_passes: list, + all_post_prepare_passes: list, + all_pre_convert_passes: list, + all_post_convert_passes: list, + ) -> "Optional[QuantizationRecipe]": """ - Util to combine multiple backend recipes into a single multi-backend recipe. - - Args: - backend_recipes: List of ExportRecipe objects to combine - recipe_name: Optional name for the combined recipe + Build the combined QuantizationRecipe from per-recipe collected lists. - Returns: - Combined ExportRecipe for multi-backend deployment + Returns None when no recipe contributed any quantization fields, and logs + an INFO message so callers know quantization is absent from the combination. """ - overriding = [ - r.name or f"recipes[{i}]" - for i, r in enumerate(backend_recipes) - if r.pipeline_stages - ] - if overriding: + # is_qat must agree: the two flows (QAT vs PTQ) are incompatible. + cls._assert_scalar_fields_agree("is_qat", is_qat_values) + + # At most one recipe may supply a train_fn. + non_none_train_fns = [f for f in train_fn_values if f is not None] + if len(non_none_train_fns) > 1: raise ValueError( - "Cannot combine recipes that override pipeline_stages, there is no " - f"correct way to merge the orderings: {overriding}" + "Cannot combine recipes: more than one recipe provides a train_fn." ) - # Extract components from individual recipes - all_partitioners = [] - all_partitioners_by_method = {} - all_quantizers = [] - all_ao_quantization_configs = [] - all_pre_edge_passes = [] - all_transform_passes = [] - combined_backend_config = None + # Multiple calibration_inputs_fn values are chained into a single factory. + non_none_calib_fns = [f for f in calibration_inputs_fn_values if f is not None] + if len(non_none_calib_fns) > 1: + _fns = non_none_calib_fns - for recipe in backend_recipes: - # Collect pre-edge transform passes - if recipe.aten_transform_passes: - all_pre_edge_passes.extend(recipe.aten_transform_passes) - - # Collect partitioners from lowering recipes - if recipe.lowering_recipe and recipe.lowering_recipe.partitioners: - partitioners = recipe.lowering_recipe.partitioners - if isinstance(partitioners, dict): - for method_name, method_partitioners in partitioners.items(): - all_partitioners_by_method.setdefault(method_name, []).extend( - method_partitioners - ) - else: - all_partitioners.extend(partitioners) - - # Collect transform passes from lowering recipes - if recipe.lowering_recipe and recipe.lowering_recipe.edge_transform_passes: - all_transform_passes.extend( - recipe.lowering_recipe.edge_transform_passes - ) - - # Collect for quantize stage - if quantization_recipe := recipe.quantization_recipe: - # Collect PT2E quantizers - if quantization_recipe.quantizers: - all_quantizers.extend(quantization_recipe.quantizers) - - # Collect source transform configs - if quantization_recipe.ao_quantization_configs: - all_ao_quantization_configs.extend( - quantization_recipe.ao_quantization_configs - ) + def _combined_calib_fn(): + for _fn in _fns: + yield from _fn() - # Use the first backend config as base - if combined_backend_config is None and recipe.executorch_backend_config: - combined_backend_config = copy.deepcopy( - recipe.executorch_backend_config - ) - - # Create combined quantization recipe - combined_quantization_recipe = None - if all_quantizers or all_ao_quantization_configs: - combined_quantization_recipe = QuantizationRecipe( - quantizers=all_quantizers if all_quantizers else None, - ao_quantization_configs=( - all_ao_quantization_configs if all_ao_quantization_configs else None - ), + combined_calib_fn: "Optional[Callable[[], Iterable[tuple]]]" = ( + _combined_calib_fn ) + else: + combined_calib_fn = non_none_calib_fns[0] if non_none_calib_fns else None + + if not ( + all_quantizers + or all_ao_quantization_configs + or all_pre_prepare_passes + or all_post_prepare_passes + or all_pre_convert_passes + or all_post_convert_passes + ): + logging.info( + "Combined recipe has no quantizers, quantization configs, or " + "quantization passes; quantization_recipe will be None." + ) + return None + + return QuantizationRecipe( + quantizers=all_quantizers or None, + ao_quantization_configs=all_ao_quantization_configs or None, + is_qat=is_qat_values[0] if is_qat_values else False, + dynamic_batch_size=any(dynamic_batch_size_values), + train_fn=non_none_train_fns[0] if non_none_train_fns else None, + calibration_inputs_fn=combined_calib_fn, + pre_prepare_passes=all_pre_prepare_passes or None, + post_prepare_passes=all_post_prepare_passes or None, + pre_convert_passes=all_pre_convert_passes or None, + post_convert_passes=all_post_convert_passes or None, + ) + + @classmethod + def _combine_lowering_recipe( + cls, + backend_recipes: "List[ExportRecipe]", + all_partitioners: list, + all_partitioners_by_method: dict, + all_edge_transform_passes: list, + all_edge_manager_transform_passes: list, + all_pre_partitioning_callbacks: list, + ) -> "Optional[LoweringRecipe]": + """ + Build the combined LoweringRecipe from per-recipe collected lists. + + Returns None when no recipe contributed any lowering fields, and logs + an INFO message so callers know lowering is absent from the combination. + """ if all_partitioners and all_partitioners_by_method: raise ValueError( @@ -372,7 +459,7 @@ def _combine_recipes( # noqa: C901 ) combined_partitioners = all_partitioners_by_method or all_partitioners - # By value, not identity: every provider builds a fresh config object, + # Compare edge_compile_confgs by value, not identity: every provider builds a fresh config object, # so asking for the same thing twice is not a conflict. distinct: List[tuple[str, EdgeCompileConfig]] = [] for i, recipe in enumerate(backend_recipes): @@ -394,23 +481,185 @@ def _combine_recipes( # noqa: C901 ) edge_compile_config = copy.deepcopy(distinct[0][1]) if distinct else None - combined_lowering_recipe = None - if combined_partitioners or all_transform_passes or edge_compile_config: - combined_lowering_recipe = LoweringRecipe( - partitioners=combined_partitioners if combined_partitioners else None, - edge_transform_passes=( - all_transform_passes if all_transform_passes else None - ), - edge_compile_config=edge_compile_config or EdgeCompileConfig(), + combined_pre_partitioning_callback = None + if all_pre_partitioning_callbacks: + _cbs = all_pre_partitioning_callbacks + + def _chained_pre_partitioning_callback(partitioners, programs): + for cb in _cbs: + try: + cb(partitioners, programs) + except Exception as e: + name = getattr(cb, "__qualname__", repr(cb)) + raise RuntimeError( + f"Pre-partitioning callback `{name}` failed: {e}" + ) from e + + combined_pre_partitioning_callback = _chained_pre_partitioning_callback + + if not ( + combined_partitioners + or all_edge_transform_passes + or all_edge_manager_transform_passes + or edge_compile_config + or combined_pre_partitioning_callback + ): + logging.info( + "Combined recipe has no lowering fields; lowering_recipe will be None." ) + return None + + return LoweringRecipe( + partitioners=combined_partitioners or None, + edge_transform_passes=all_edge_transform_passes or None, + edge_manager_transform_passes=all_edge_manager_transform_passes or None, + edge_compile_config=edge_compile_config or EdgeCompileConfig(), + pre_partitioning_callback=combined_pre_partitioning_callback, + ) + + @staticmethod + def _collect_lowering_fields( + acc: "_CombineAccumulator", lr: "LoweringRecipe" + ) -> None: + """Accumulate fields from a single LoweringRecipe into acc.""" + if lr.partitioners: + if isinstance(lr.partitioners, dict): + for method_name, method_partitioners in lr.partitioners.items(): + acc.partitioners_by_method.setdefault(method_name, []).extend( + method_partitioners + ) + else: + acc.partitioners.extend(lr.partitioners) + if lr.edge_transform_passes: + acc.edge_transform_passes.extend(lr.edge_transform_passes) + if lr.edge_manager_transform_passes: + acc.edge_manager_transform_passes.extend(lr.edge_manager_transform_passes) + if lr.pre_partitioning_callback: + acc.pre_partitioning_callbacks.append(lr.pre_partitioning_callback) + + @staticmethod + def _collect_quantization_fields( + acc: "_CombineAccumulator", qr: "QuantizationRecipe" + ) -> None: + """Accumulate fields from a single QuantizationRecipe into acc.""" + if qr.quantizers: + acc.quantizers.extend(qr.quantizers) + if qr.ao_quantization_configs: + acc.ao_quantization_configs.extend(qr.ao_quantization_configs) + acc.is_qat_values.append(qr.is_qat) + acc.dynamic_batch_size_values.append(qr.dynamic_batch_size) + acc.train_fn_values.append(qr.train_fn) + acc.calibration_inputs_fn_values.append(qr.calibration_inputs_fn) + if qr.pre_prepare_passes: + acc.pre_prepare_passes.extend(qr.pre_prepare_passes) + if qr.post_prepare_passes: + acc.post_prepare_passes.extend(qr.post_prepare_passes) + if qr.pre_convert_passes: + acc.pre_convert_passes.extend(qr.pre_convert_passes) + if qr.post_convert_passes: + acc.post_convert_passes.extend(qr.post_convert_passes) + + @classmethod + def _collect_recipe_fields( + cls, + backend_recipes: "List[ExportRecipe]", + ) -> "_CombineAccumulator": + """ + Iterate over all recipes and accumulate their fields into a single + _CombineAccumulator for later merging. + """ + acc = _CombineAccumulator() + + for recipe in backend_recipes: + if recipe.aten_transform_passes: + acc.pre_edge_passes.extend(recipe.aten_transform_passes) + + if lr := recipe.lowering_recipe: + cls._collect_lowering_fields(acc, lr) + + if qr := recipe.quantization_recipe: + cls._collect_quantization_fields(acc, qr) + + acc.strict_values.append(recipe.strict) + acc.mode_values.append(recipe.mode) + acc.pipeline_stages_values.append( + tuple(recipe.pipeline_stages) if recipe.pipeline_stages else None + ) + acc.source_transform_in_place_values.append( + recipe.source_transform_in_place + ) + + # Use the executorch_backend_config from the first recipe that supplies one. + if acc.backend_config is None and recipe.executorch_backend_config: + acc.backend_config = copy.deepcopy(recipe.executorch_backend_config) + + return acc + + @classmethod + def _combine_recipes( + cls, backend_recipes: "List[ExportRecipe]", recipe_name: "Optional[str]" = None + ) -> "ExportRecipe": + """ + Util to combine multiple backend recipes into a single multi-backend recipe. + + Args: + backend_recipes: List of ExportRecipe objects to combine + recipe_name: Optional name for the combined recipe + + Returns: + Combined ExportRecipe for multi-backend deployment + """ + acc = cls._collect_recipe_fields(backend_recipes) + + # Validate scalar fields that must agree across all recipes. + cls._assert_scalar_fields_agree("strict", acc.strict_values) + cls._assert_scalar_fields_agree("mode", acc.mode_values) + cls._assert_scalar_fields_agree("pipeline_stages", acc.pipeline_stages_values) + cls._assert_scalar_fields_agree( + "source_transform_in_place", acc.source_transform_in_place_values + ) + + combined_quantization_recipe = cls._combine_quantization_recipe( + is_qat_values=acc.is_qat_values, + dynamic_batch_size_values=acc.dynamic_batch_size_values, + train_fn_values=acc.train_fn_values, + calibration_inputs_fn_values=acc.calibration_inputs_fn_values, + all_quantizers=acc.quantizers, + all_ao_quantization_configs=acc.ao_quantization_configs, + all_pre_prepare_passes=acc.pre_prepare_passes, + all_post_prepare_passes=acc.post_prepare_passes, + all_pre_convert_passes=acc.pre_convert_passes, + all_post_convert_passes=acc.post_convert_passes, + ) + + combined_lowering_recipe = cls._combine_lowering_recipe( + backend_recipes=backend_recipes, + all_partitioners=acc.partitioners, + all_partitioners_by_method=acc.partitioners_by_method, + all_edge_transform_passes=acc.edge_transform_passes, + all_edge_manager_transform_passes=acc.edge_manager_transform_passes, + all_pre_partitioning_callbacks=acc.pre_partitioning_callbacks, + ) recipe_name = recipe_name or "_".join( [r.name for r in backend_recipes if r.name is not None] ) + # All pipeline_stages values are equal (enforced above); use the first non-None one. + shared_pipeline_stages = next( + (r.pipeline_stages for r in backend_recipes if r.pipeline_stages), None + ) return cls( name=recipe_name, quantization_recipe=combined_quantization_recipe, - aten_transform_passes=all_pre_edge_passes, + aten_transform_passes=acc.pre_edge_passes or None, lowering_recipe=combined_lowering_recipe, - executorch_backend_config=combined_backend_config, + executorch_backend_config=acc.backend_config, + pipeline_stages=shared_pipeline_stages, + strict=acc.strict_values[0] if acc.strict_values else True, + mode=acc.mode_values[0] if acc.mode_values else Mode.RELEASE, + source_transform_in_place=( + acc.source_transform_in_place_values[0] + if acc.source_transform_in_place_values + else False + ), ) diff --git a/export/stages.py b/export/stages.py index 96f46d750c3..06bcf8a0b74 100644 --- a/export/stages.py +++ b/export/stages.py @@ -1,6 +1,7 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # Copyright 2025 Arm Limited and/or its affiliates. +# Copyright 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. @@ -23,7 +24,16 @@ from torch._export.pass_base import PassType from torch.fx.passes.infra.pass_manager import PassManager as GraphModulePassManager from torchao.quantization import quantize_ -from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e +from torchao.quantization.pt2e import ( + allow_exported_model_train_eval, + move_exported_model_to_eval, + move_exported_model_to_train, +) +from torchao.quantization.pt2e.quantize_pt2e import ( + convert_pt2e, + prepare_pt2e, + prepare_qat_pt2e, +) from torchao.quantization.pt2e.quantizer import ( ComposableQuantizer, Quantizer as TorchAOPT2EQuantizer, @@ -283,6 +293,15 @@ def run(self, artifact: PipelineArtifact) -> None: # method the dict does not name, so it would copy to apply nothing. final_passes = pass_manager or _drop_empty(transform_passes) or None + export_recipe = artifact.context.get("export_recipe") + lowering_recipe = getattr(export_recipe, "lowering_recipe", None) + + if ( + lowering_recipe is not None + and lowering_recipe.pre_partitioning_callback is not None + ): + lowering_recipe.pre_partitioning_callback(self._partitioners, artifact.data) + with validation_disabled(): edge_program_manager = to_edge_transform_and_lower( exported_programs, @@ -465,16 +484,31 @@ def _get_quantizer_for_prepare_pt2e(self, quantizers: List[Any]): else: raise ValueError("No quantizers detected") + @staticmethod + def _apply_passes( + model: "torch.fx.GraphModule", + passes: Optional[List[Callable]], + ) -> "torch.fx.GraphModule": + for pass_fn in passes or []: + try: + model = pass_fn(model) + except Exception as exc: + raise RuntimeError( + f"QuantizeStage: Pass '{pass_fn!r}' raised an error: {exc}" + ) from exc + return model + def run(self, artifact: PipelineArtifact) -> None: if not self._quantization_recipe or not self._quantization_recipe.quantizers: logging.info( - "Quantization recipe is invalid to run QunatizeStage, returning original model" + "Quantization recipe is invalid to run QuantizeStage, returning original model" ) self._artifact = artifact return assert isinstance(artifact.data, dict) + recipe = self._quantization_recipe models = artifact.data example_inputs = artifact.get_context("example_inputs") @@ -487,17 +521,78 @@ def run(self, artifact: PipelineArtifact) -> None: ) inputs = example_inputs[method_name][0] - captured_graph = torch.export.export(model, inputs, strict=True).module() - quantizer = self._get_quantizer_for_prepare_pt2e( - self._quantization_recipe.quantizers # pyre-ignore + # When dynamic_batch_size is requested, mark dimension 0 of every + # tensor input as dynamic so that a QAT training loop can feed + # mini-batches of arbitrary size through the prepared graph. + export_dynamic_shapes = None + if recipe.dynamic_batch_size: + from torch.export import Dim + + batch = Dim("batch", min=1) + export_dynamic_shapes = tuple( + {0: batch} if isinstance(t, torch.Tensor) else None for t in inputs + ) + + # QAT requires the model to be in training mode at capture time so + # that batch_norm and dropout decompose with training-mode semantics. + if recipe.is_qat: + model.train() + + captured_graph = torch.export.export( + model, inputs, dynamic_shapes=export_dynamic_shapes, strict=True + ).module() + + # Pass 1: pre-prepare passes. + captured_graph = self._apply_passes( + captured_graph, recipe.pre_prepare_passes ) - prepared_model = prepare_pt2e(captured_graph, quantizer) - for calibration_input in example_inputs[method_name]: - prepared_model(*calibration_input) + quantizer = self._get_quantizer_for_prepare_pt2e(recipe.quantizers) + + if recipe.is_qat: + if recipe.train_fn is None: + raise ValueError("train_fn must be provided when is_qat=True") + prepared_model = prepare_qat_pt2e(captured_graph, quantizer) + + # Pass 2: post-prepare passes. + prepared_model = self._apply_passes( + prepared_model, recipe.post_prepare_passes + ) + + allow_exported_model_train_eval(prepared_model) + move_exported_model_to_train(prepared_model) + recipe.train_fn(prepared_model) + move_exported_model_to_eval(prepared_model) + else: + prepared_model = prepare_pt2e(captured_graph, quantizer) + + # Pass 2: post-prepare passes. + prepared_model = self._apply_passes( + prepared_model, recipe.post_prepare_passes + ) + + # Use custom calibration inputs when provided; fall back to example inputs. + if recipe.calibration_inputs_fn is not None: + calibration_inputs = recipe.calibration_inputs_fn() + else: + calibration_inputs = example_inputs[method_name] + + for calibration_input in calibration_inputs: + prepared_model(*calibration_input) + + # Pass 3: pre-convert passes. + prepared_model = self._apply_passes( + prepared_model, recipe.pre_convert_passes + ) quantized_model = convert_pt2e(prepared_model) + + # Pass 4: post-convert passes. + quantized_model = self._apply_passes( + quantized_model, recipe.post_convert_passes + ) + quantized_models[method_name] = quantized_model self._artifact = artifact.copy_with_new_data(quantized_models) @@ -616,7 +711,7 @@ def stage_type(self) -> str: def valid_predecessor_stages(self) -> List["StageType"]: return [ StageType.TO_EDGE, - # StageType.TO_EDGE_TRANSFORM_AND_LOWER, # TODO + StageType.TO_EDGE_TRANSFORM_AND_LOWER, ] @property diff --git a/export/tests/test_export_recipe.py b/export/tests/test_export_recipe.py index 5fc4701a5f6..e1488ffe58c 100644 --- a/export/tests/test_export_recipe.py +++ b/export/tests/test_export_recipe.py @@ -7,12 +7,18 @@ # pyre-strict import unittest -from typing import Any, Dict, Optional, Sequence +from typing import Any, Dict, List, Optional, Sequence from unittest.mock import Mock import torch -from executorch.export.recipe import ExportRecipe, RecipeType +from executorch.export.recipe import ( + ExportRecipe, + LoweringRecipe, + Mode, + QuantizationRecipe, + RecipeType, +) from executorch.export.recipe_provider import BackendRecipeProvider from executorch.export.recipe_registry import recipe_registry @@ -293,8 +299,8 @@ def test_combine_keeps_partitioners(self) -> None: def test_combine_rejects_pipeline_stages(self) -> None: from executorch.export.types import StageType - # Unnamed recipes fall back to their position; named ones are named. - with self.assertRaisesRegex(ValueError, r"pipeline_stages.*recipes\[0\]"): + # Recipes with different pipeline_stages (including None vs. a list) must be rejected. + with self.assertRaisesRegex(ValueError, r"pipeline_stages"): ExportRecipe.combine( [ ExportRecipe( @@ -303,7 +309,7 @@ def test_combine_rejects_pipeline_stages(self) -> None: ExportRecipe(name="b"), ] ) - with self.assertRaisesRegex(ValueError, r"pipeline_stages.*'stagey'"): + with self.assertRaisesRegex(ValueError, r"pipeline_stages"): ExportRecipe.combine( [ ExportRecipe( @@ -334,3 +340,440 @@ def test_combine_keeps_edge_transform_passes(self) -> None: self.assertEqual( combined.lowering_recipe.edge_transform_passes, [first, second] ) + + +# --------------------------------------------------------------------------- +# Helpers shared by combine-recipe tests +# --------------------------------------------------------------------------- + + +def _make_pass(name: str, call_log: List[str]): + """Return a graph-module pass that appends *name* to *call_log*.""" + + def pass_fn(m): + call_log.append(name) + return m + + return pass_fn + + +class TestCombineRecipesEmpty(unittest.TestCase): + def test_empty_recipes_raises(self) -> None: + with self.assertRaises(ValueError): + ExportRecipe.combine([]) + + +class TestCombineRecipesSingleRecipe(unittest.TestCase): + def test_single_recipe_returned_unchanged(self) -> None: + recipe = ExportRecipe(name="solo") + result = ExportRecipe.combine([recipe]) + self.assertIs(result, recipe) + + +class TestCombineRecipesScalarFields(unittest.TestCase): + """Fields that must be identical across all combined recipes.""" + + def test_conflicting_strict_raises(self) -> None: + r1 = ExportRecipe(name="a", strict=True) + r2 = ExportRecipe(name="b", strict=False) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("strict", str(cm.exception)) + + def test_conflicting_mode_raises(self) -> None: + r1 = ExportRecipe(name="a", mode=Mode.DEBUG) + r2 = ExportRecipe(name="b", mode=Mode.RELEASE) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("mode", str(cm.exception)) + + def test_conflicting_source_transform_in_place_raises(self) -> None: + r1 = ExportRecipe(name="a", source_transform_in_place=True) + r2 = ExportRecipe(name="b", source_transform_in_place=False) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("source_transform_in_place", str(cm.exception)) + + def test_agreeing_scalar_fields_are_preserved(self) -> None: + r1 = ExportRecipe( + name="a", strict=False, mode=Mode.DEBUG, source_transform_in_place=True + ) + r2 = ExportRecipe( + name="b", strict=False, mode=Mode.DEBUG, source_transform_in_place=True + ) + result = ExportRecipe.combine([r1, r2]) + self.assertFalse(result.strict) + self.assertEqual(result.mode, Mode.DEBUG) + self.assertTrue(result.source_transform_in_place) + + def test_name_is_joined_from_input_recipe_names(self) -> None: + r1 = ExportRecipe(name="backend_a") + r2 = ExportRecipe(name="backend_b") + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.name, "backend_a_backend_b") + + def test_custom_recipe_name_is_used(self) -> None: + r1 = ExportRecipe(name="a") + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2], recipe_name="custom_name") + self.assertEqual(result.name, "custom_name") + + +class TestCombineRecipesAtenTransformPasses(unittest.TestCase): + def test_aten_transform_passes_merged(self) -> None: + pass1 = Mock() + pass2 = Mock() + r1 = ExportRecipe(name="a", aten_transform_passes=[pass1]) + r2 = ExportRecipe(name="b", aten_transform_passes=[pass2]) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.aten_transform_passes, [pass1, pass2]) + + def test_aten_transform_passes_none_when_both_empty(self) -> None: + r1 = ExportRecipe(name="a") + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2]) + self.assertIsNone(result.aten_transform_passes) + + def test_aten_transform_passes_one_side_none(self) -> None: + pass1 = Mock() + r1 = ExportRecipe(name="a", aten_transform_passes=[pass1]) + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.aten_transform_passes, [pass1]) + + +class TestCombineRecipesLowering(unittest.TestCase): + def test_partitioners_merged(self) -> None: + p1 = Mock() + p2 = Mock() + r1 = ExportRecipe(name="a", lowering_recipe=LoweringRecipe(partitioners=[p1])) + r2 = ExportRecipe(name="b", lowering_recipe=LoweringRecipe(partitioners=[p2])) + result = ExportRecipe.combine([r1, r2]) + self.assertIsNotNone(result.lowering_recipe) + self.assertEqual(result.lowering_recipe.partitioners, [p1, p2]) + + def test_edge_transform_passes_merged(self) -> None: + pass1 = Mock() + pass2 = Mock() + r1 = ExportRecipe( + name="a", lowering_recipe=LoweringRecipe(edge_transform_passes=[pass1]) + ) + r2 = ExportRecipe( + name="b", lowering_recipe=LoweringRecipe(edge_transform_passes=[pass2]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.lowering_recipe.edge_transform_passes, [pass1, pass2]) + + def test_edge_manager_transform_passes_merged(self) -> None: + pass1 = Mock() + pass2 = Mock() + r1 = ExportRecipe( + name="a", + lowering_recipe=LoweringRecipe(edge_manager_transform_passes=[pass1]), + ) + r2 = ExportRecipe( + name="b", + lowering_recipe=LoweringRecipe(edge_manager_transform_passes=[pass2]), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual( + result.lowering_recipe.edge_manager_transform_passes, [pass1, pass2] + ) + + def test_lowering_recipe_none_when_nothing_contributed(self) -> None: + r1 = ExportRecipe(name="a") + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2]) + self.assertIsNone(result.lowering_recipe) + + def test_edge_compile_config_taken_from_first_recipe_with_one(self) -> None: + from executorch.exir.capture import EdgeCompileConfig + + config = EdgeCompileConfig() + r1 = ExportRecipe(name="a") + r2 = ExportRecipe( + name="b", + lowering_recipe=LoweringRecipe( + partitioners=[Mock()], edge_compile_config=config + ), + ) + result = ExportRecipe.combine([r1, r2]) + # combine() deepcopies the config so the combined recipe cannot mutate + # the provider's shared object; assert value-equality, not identity. + self.assertIsNotNone(result.lowering_recipe) + self.assertEqual(result.lowering_recipe.edge_compile_config, config) + self.assertIsNot(result.lowering_recipe.edge_compile_config, config) + + +class TestCombineRecipesQuantization(unittest.TestCase): + def test_quantizers_merged(self) -> None: + q1 = Mock() + q2 = Mock() + r1 = ExportRecipe( + name="a", quantization_recipe=QuantizationRecipe(quantizers=[q1]) + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[q2]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertIsNotNone(result.quantization_recipe) + self.assertEqual(result.quantization_recipe.quantizers, [q1, q2]) + + def test_ao_quantization_configs_merged(self) -> None: + from executorch.export.recipe import AOQuantizationConfig + from torchao.core.config import AOBaseConfig + + cfg1 = AOQuantizationConfig(ao_base_config=Mock(spec=AOBaseConfig)) + cfg2 = AOQuantizationConfig(ao_base_config=Mock(spec=AOBaseConfig)) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe(ao_quantization_configs=[cfg1]), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe(ao_quantization_configs=[cfg2]), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual( + result.quantization_recipe.ao_quantization_configs, [cfg1, cfg2] + ) + + def test_quantization_recipe_none_when_nothing_contributed(self) -> None: + r1 = ExportRecipe(name="a") + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2]) + self.assertIsNone(result.quantization_recipe) + + def test_conflicting_is_qat_raises(self) -> None: + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=True), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=False), + ) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("is_qat", str(cm.exception)) + + def test_agreeing_is_qat_preserved(self) -> None: + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=True), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=True), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertTrue(result.quantization_recipe.is_qat) + + def test_two_train_fns_raises(self) -> None: + fn1 = Mock() + fn2 = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], is_qat=True, train_fn=fn1 + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], is_qat=True, train_fn=fn2 + ), + ) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("train_fn", str(cm.exception)) + + def test_single_train_fn_preserved(self) -> None: + fn = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], is_qat=True, train_fn=fn + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=True), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertIs(result.quantization_recipe.train_fn, fn) + + def test_single_calibration_inputs_fn_preserved(self) -> None: + fn = Mock(return_value=[(1,), (2,)]) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn + ), + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertIs(result.quantization_recipe.calibration_inputs_fn, fn) + + def test_two_calibration_inputs_fns_chained(self) -> None: + fn1 = Mock(return_value=[(1,), (2,)]) + fn2 = Mock(return_value=[(3,), (4,)]) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn1 + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn2 + ), + ) + result = ExportRecipe.combine([r1, r2]) + combined_fn = result.quantization_recipe.calibration_inputs_fn + self.assertIsNotNone(combined_fn) + # Each factory must be called exactly once when the combined factory is consumed. + all_inputs = list(combined_fn()) + fn1.assert_called_once_with() + fn2.assert_called_once_with() + self.assertEqual(all_inputs, [(1,), (2,), (3,), (4,)]) + + def test_three_calibration_inputs_fns_chained_in_order(self) -> None: + fn1 = Mock(return_value=[(1,)]) + fn2 = Mock(return_value=[(2,)]) + fn3 = Mock(return_value=[(3,)]) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn1 + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn2 + ), + ) + r3 = ExportRecipe( + name="c", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn3 + ), + ) + result = ExportRecipe.combine([r1, r2, r3]) + combined_fn = result.quantization_recipe.calibration_inputs_fn + self.assertEqual(list(combined_fn()), [(1,), (2,), (3,)]) + + def test_no_calibration_inputs_fn_stays_none(self) -> None: + r1 = ExportRecipe( + name="a", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertIsNone(result.quantization_recipe.calibration_inputs_fn) + + def test_pre_prepare_passes_merged(self) -> None: + log: List[str] = [] + p1 = _make_pass("pre_a", log) + p2 = _make_pass("pre_b", log) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_prepare_passes=[p1] + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_prepare_passes=[p2] + ), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.pre_prepare_passes, [p1, p2]) + + def test_post_prepare_passes_merged(self) -> None: + p1 = Mock() + p2 = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], post_prepare_passes=[p1] + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], post_prepare_passes=[p2] + ), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.post_prepare_passes, [p1, p2]) + + def test_pre_convert_passes_merged(self) -> None: + p1 = Mock() + p2 = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_convert_passes=[p1] + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_convert_passes=[p2] + ), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.pre_convert_passes, [p1, p2]) + + def test_post_convert_passes_merged(self) -> None: + p1 = Mock() + p2 = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], post_convert_passes=[p1] + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], post_convert_passes=[p2] + ), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.post_convert_passes, [p1, p2]) + + def test_all_pass_lists_none_when_nothing_contributed(self) -> None: + r1 = ExportRecipe( + name="a", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + result = ExportRecipe.combine([r1, r2]) + qr = result.quantization_recipe + self.assertIsNone(qr.pre_prepare_passes) + self.assertIsNone(qr.post_prepare_passes) + self.assertIsNone(qr.pre_convert_passes) + self.assertIsNone(qr.post_convert_passes) + + def test_pass_lists_preserved_when_only_one_recipe_contributes(self) -> None: + p = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_prepare_passes=[p] + ), + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.pre_prepare_passes, [p]) diff --git a/export/tests/test_export_session.py b/export/tests/test_export_session.py index 2e9112f7995..93902fd86d7 100644 --- a/export/tests/test_export_session.py +++ b/export/tests/test_export_session.py @@ -642,6 +642,36 @@ def test_dict_exported_program_input_type_detection(self) -> None: pipeline = session._get_default_pipeline() self.assertNotIn(StageType.TORCH_EXPORT, pipeline) + def test_edge_manager_transform_passes_get_their_stage(self) -> None: + # Nothing else in the default pipeline runs them, so without this the + # recipe's passes are accepted and then silently never applied. + session = ExportSession( + model=self.model, + example_inputs=[self.example_inputs], + export_recipe=ExportRecipe( + name="t", + lowering_recipe=LoweringRecipe( + edge_manager_transform_passes=[lambda epm: []] + ), + ), + ) + pipeline = session._get_default_pipeline() + self.assertIn(StageType.EDGE_PROGRAM_MANAGER_TRANSFORM, pipeline) + self.assertGreater( + pipeline.index(StageType.EDGE_PROGRAM_MANAGER_TRANSFORM), + pipeline.index(StageType.TO_EDGE_TRANSFORM_AND_LOWER), + ) + + def test_no_transform_passes_means_no_stage(self) -> None: + session = ExportSession( + model=self.model, + example_inputs=[self.example_inputs], + export_recipe=self.recipe, + ) + self.assertNotIn( + StageType.EDGE_PROGRAM_MANAGER_TRANSFORM, session._get_default_pipeline() + ) + def test_example_inputs_required_for_nn_module(self) -> None: """Test that example_inputs are required for nn.Module.""" with self.assertRaises(ValueError) as cm: diff --git a/export/tests/test_export_stages.py b/export/tests/test_export_stages.py index 935a796591a..351ae748b86 100644 --- a/export/tests/test_export_stages.py +++ b/export/tests/test_export_stages.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 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. @@ -409,6 +410,7 @@ def test_run_no_quantizers(self) -> None: result_artifact = stage.get_artifacts() self.assertEqual(result_artifact, artifact) + @patch("executorch.export.stages.move_exported_model_to_eval") @patch("executorch.export.stages.convert_pt2e") @patch("executorch.export.stages.prepare_pt2e") @patch("executorch.export.stages.ComposableQuantizer") @@ -419,11 +421,19 @@ def test_run_with_quantizers( mock_composable_quantizer: Mock, mock_prepare_pt2e: Mock, mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, ) -> None: """Test execution with quantizers""" mock_quantizer = self.create_dummy_quantizer() mock_recipe = Mock(spec=QuantizationRecipe) mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.dynamic_batch_size = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None stage = QuantizeStage(mock_recipe) # Mock the torch.export.export chain @@ -443,9 +453,9 @@ def test_run_with_quantizers( artifact = PipelineArtifact(data=self.models_dict, context=self.context) stage.run(artifact) - # Verify torch.export.export was called + # Verify torch.export.export was called with dynamic_shapes=None (no dynamic batch) mock_torch_export.assert_called_once_with( - self.model, self.example_inputs[0], strict=True + self.model, self.example_inputs[0], dynamic_shapes=None, strict=True ) # Verify ComposableQuantizer was created with the quantizers @@ -471,11 +481,401 @@ def test_run_with_quantizers( self.assertEqual(artifact.data["forward"], self.model) self.assertIsNot(result_artifact.data["forward"], self.model) + @patch("executorch.export.stages.allow_exported_model_train_eval") + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.move_exported_model_to_train") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_qat_calls_prepare_qat_pt2e( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_train: Mock, + mock_move_to_eval: Mock, + mock_allow_train_eval: Mock, + ) -> None: + """QAT flow: prepare_qat_pt2e is called and train_fn is invoked with the prepared model. + allow_exported_model_train_eval must be called after preparation. + move_exported_model_to_train must be called before train_fn, and + move_exported_model_to_eval must be called after train_fn.""" + mock_quantizer = self.create_dummy_quantizer() + call_order = [] + mock_allow_train_eval.side_effect = lambda m: call_order.append( + "allow_train_eval" + ) + mock_move_to_train.side_effect = lambda m: call_order.append("to_train") + mock_move_to_eval.side_effect = lambda m: call_order.append("to_eval") + + train_fn = Mock(side_effect=lambda m: call_order.append("train_fn")) + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = True + mock_recipe.dynamic_batch_size = False + mock_recipe.train_fn = train_fn + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_captured_graph = Mock() + mock_exported_program.module.return_value = mock_captured_graph + mock_torch_export.return_value = mock_exported_program + + mock_composed_quantizer = Mock() + mock_composable_quantizer.return_value = mock_composed_quantizer + mock_prepared_model = Mock() + mock_prepare_qat_pt2e.return_value = mock_prepared_model + mock_quantized_model = Mock() + mock_convert_pt2e.return_value = mock_quantized_model + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + # prepare_qat_pt2e must be called, not prepare_pt2e + mock_prepare_qat_pt2e.assert_called_once_with( + mock_captured_graph, mock_composed_quantizer + ) + # allow_exported_model_train_eval before move_to_train, then train_fn, then to_eval + self.assertEqual( + call_order, ["allow_train_eval", "to_train", "train_fn", "to_eval"] + ) + mock_allow_train_eval.assert_called_once_with(mock_prepared_model) + mock_move_to_train.assert_called_once_with(mock_prepared_model) + mock_move_to_eval.assert_called_once_with(mock_prepared_model) + # train_fn must be called with the prepared model + train_fn.assert_called_once_with(mock_prepared_model) + # convert_pt2e must still be called after training + mock_convert_pt2e.assert_called_once_with(mock_prepared_model) + + result_artifact = stage.get_artifacts() + self.assertEqual(result_artifact.data["forward"], mock_quantized_model) + + @patch("torch.export.export") + def test_run_qat_missing_train_fn_raises(self, mock_torch_export: Mock) -> None: + """QAT flow with train_fn=None must raise ValueError.""" + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = True + mock_recipe.dynamic_batch_size = False + mock_recipe.train_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + + with self.assertRaises(ValueError) as cm: + stage.run(artifact) + self.assertIn("train_fn must be provided when is_qat=True", str(cm.exception)) + + @patch("executorch.export.stages.allow_exported_model_train_eval") + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.move_exported_model_to_train") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_qat_model_put_in_train_mode_before_export( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_train: Mock, + mock_move_to_eval: Mock, + mock_allow_train_eval: Mock, + ) -> None: + """QAT: model.train() must be called before torch.export.export so that + batch_norm and dropout decompose with training-mode semantics.""" + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = True + mock_recipe.dynamic_batch_size = False + mock_recipe.train_fn = Mock() + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + # Start model in eval mode; the stage must switch it to train. + self.model.eval() + self.assertFalse(self.model.training) + + training_at_export_time = [] + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + + def capture_training_flag(model, *args, **kwargs): + training_at_export_time.append(model.training) + return mock_exported_program + + mock_torch_export.side_effect = capture_training_flag + mock_prepare_qat_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + stage.run(PipelineArtifact(data=self.models_dict, context=self.context)) + + # The model must have been in training mode when export was called. + self.assertEqual(training_at_export_time, [True]) + + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_ptq_does_not_call_prepare_qat_pt2e( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, + ) -> None: + """PTQ flow must not call prepare_qat_pt2e (regression guard).""" + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.dynamic_batch_size = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepare_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + mock_prepare_pt2e.assert_called_once() + mock_prepare_qat_pt2e.assert_not_called() + + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_ptq_four_passes_called_in_order( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, + ) -> None: + """All four pass hooks are called at the correct points in the PTQ flow.""" + call_order = [] + + def make_pass(name): + def pass_fn(m): + call_order.append(name) + return m + + return pass_fn + + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.dynamic_batch_size = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = [make_pass("pre_prepare")] + mock_recipe.post_prepare_passes = [make_pass("post_prepare")] + mock_recipe.pre_convert_passes = [make_pass("pre_convert")] + mock_recipe.post_convert_passes = [make_pass("post_convert")] + + mock_exported_program = Mock(spec=ExportedProgram) + mock_graph = Mock() + mock_exported_program.module.return_value = mock_graph + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepare_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + self.assertEqual( + call_order, + ["pre_prepare", "post_prepare", "pre_convert", "post_convert"], + ) + + @patch("executorch.export.stages.allow_exported_model_train_eval") + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.move_exported_model_to_train") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_qat_four_passes_called_in_order( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_train: Mock, + mock_move_to_eval: Mock, + mock_allow_train_eval: Mock, + ) -> None: + """All four pass hooks are called at the correct points in the QAT flow.""" + call_order = [] + + def make_pass(name): + def pass_fn(m): + call_order.append(name) + return m + + return pass_fn + + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = True + mock_recipe.dynamic_batch_size = False + mock_recipe.train_fn = Mock() + mock_recipe.pre_prepare_passes = [make_pass("pre_prepare")] + mock_recipe.post_prepare_passes = [make_pass("post_prepare")] + mock_recipe.pre_convert_passes = [make_pass("pre_convert")] + mock_recipe.post_convert_passes = [make_pass("post_convert")] + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepare_qat_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + self.assertEqual( + call_order, + ["pre_prepare", "post_prepare", "pre_convert", "post_convert"], + ) + + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_ptq_uses_calibration_inputs_fn_when_provided( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, + ) -> None: + """When calibration_inputs_fn is set, it is called and its output is used for calibration.""" + custom_input = (torch.randn(2, 10),) + calibration_inputs_fn = Mock(return_value=[custom_input]) + + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.dynamic_batch_size = False + mock_recipe.calibration_inputs_fn = calibration_inputs_fn + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepared_model = Mock() + mock_prepare_pt2e.return_value = mock_prepared_model + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + # calibration_inputs_fn must be called with no arguments + calibration_inputs_fn.assert_called_once_with() + # prepared model must be called with the custom calibration input + mock_prepared_model.assert_called_once_with(*custom_input) + + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_ptq_falls_back_to_example_inputs_when_no_calibration_fn( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, + ) -> None: + """When calibration_inputs_fn is None, example inputs are used for calibration.""" + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.dynamic_batch_size = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepared_model = Mock() + mock_prepare_pt2e.return_value = mock_prepared_model + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + # The prepared model must be called with the example inputs (one tuple) + mock_prepared_model.assert_called_once_with(*self.example_inputs[0]) + def test_run_empty_example_inputs(self) -> None: """Test error when example inputs list is empty.""" mock_quantizer = Mock() mock_recipe = Mock(spec=QuantizationRecipe) mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None stage = QuantizeStage(mock_recipe) context = {"example_inputs": {"forward": []}} artifact = PipelineArtifact(data=self.models_dict, context=context) @@ -650,6 +1050,175 @@ def test_run_edge_manager_none(self) -> None: self.assertIn("Edge program manager is not set", str(cm.exception)) +class TestQuantizeStageExportDynamicShapes(unittest.TestCase): + """Tests for the dynamic_batch_size export behavior in QuantizeStage.""" + + def setUp(self) -> None: + self.model = torch.nn.Linear(10, 5) + self.models_dict = {"forward": self.model} + self.example_inputs = [(torch.randn(1, 10),)] + self.context = {"example_inputs": {"forward": self.example_inputs}} + + @staticmethod + def _make_recipe(is_qat: bool, dynamic_batch_size: bool) -> Mock: + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [Mock(spec=TorchAOPT2EQuantizer)] + mock_recipe.is_qat = is_qat + mock_recipe.dynamic_batch_size = dynamic_batch_size + mock_recipe.calibration_inputs_fn = None + mock_recipe.train_fn = Mock() if is_qat else None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + return mock_recipe + + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_dynamic_batch_size_false_exports_without_dynamic_shapes( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, + ) -> None: + """When dynamic_batch_size=False, torch.export.export is called with dynamic_shapes=None.""" + mock_ep = Mock(spec=ExportedProgram) + mock_ep.module.return_value = Mock() + mock_torch_export.return_value = mock_ep + mock_composable_quantizer.return_value = Mock() + mock_prepare_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + recipe = self._make_recipe(is_qat=False, dynamic_batch_size=False) + stage = QuantizeStage(recipe) + stage.run(PipelineArtifact(data=self.models_dict, context=self.context)) + + mock_torch_export.assert_called_once_with( + self.model, + self.example_inputs[0], + dynamic_shapes=None, + strict=True, + ) + + @patch("executorch.export.stages.allow_exported_model_train_eval") + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.move_exported_model_to_train") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_dynamic_batch_size_true_exports_with_dynamic_batch_dim( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_train: Mock, + mock_move_to_eval: Mock, + mock_allow_train_eval: Mock, + ) -> None: + """When dynamic_batch_size=True, torch.export.export is called with a + dynamic_shapes tuple where dimension 0 of every tensor is dynamic.""" + mock_ep = Mock(spec=ExportedProgram) + mock_ep.module.return_value = Mock() + mock_torch_export.return_value = mock_ep + mock_composable_quantizer.return_value = Mock() + mock_prepare_qat_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + recipe = self._make_recipe(is_qat=True, dynamic_batch_size=True) + stage = QuantizeStage(recipe) + stage.run(PipelineArtifact(data=self.models_dict, context=self.context)) + + call_kwargs = mock_torch_export.call_args + dynamic_shapes_arg = call_kwargs.kwargs.get( + "dynamic_shapes", call_kwargs.args[2] if len(call_kwargs.args) > 2 else None + ) + # dynamic_shapes must be a non-None tuple with one entry per input tensor. + self.assertIsNotNone(dynamic_shapes_arg) + self.assertIsInstance(dynamic_shapes_arg, tuple) + self.assertEqual(len(dynamic_shapes_arg), len(self.example_inputs[0])) + # The entry for the single tensor input must map dim 0 to a Dim. + first_entry = dynamic_shapes_arg[0] + self.assertIsInstance(first_entry, dict) + self.assertIn(0, first_entry) + + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_dynamic_batch_size_true_ptq_exports_with_dynamic_shapes( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, + ) -> None: + """When dynamic_batch_size=True and is_qat=False, export is called with dynamic shapes.""" + mock_ep = Mock(spec=ExportedProgram) + mock_ep.module.return_value = Mock() + mock_torch_export.return_value = mock_ep + mock_composable_quantizer.return_value = Mock() + mock_prepare_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + recipe = self._make_recipe(is_qat=False, dynamic_batch_size=True) + stage = QuantizeStage(recipe) + stage.run(PipelineArtifact(data=self.models_dict, context=self.context)) + + call_kwargs = mock_torch_export.call_args + dynamic_shapes_arg = call_kwargs.kwargs.get( + "dynamic_shapes", call_kwargs.args[2] if len(call_kwargs.args) > 2 else None + ) + self.assertIsNotNone(dynamic_shapes_arg) + self.assertIsInstance(dynamic_shapes_arg, tuple) + self.assertEqual(len(dynamic_shapes_arg), len(self.example_inputs[0])) + first_entry = dynamic_shapes_arg[0] + self.assertIsInstance(first_entry, dict) + self.assertIn(0, first_entry) + + def test_dynamic_batch_size_true_ptq_calibration_with_variable_batch_sizes( + self, + ) -> None: + """PTQ calibration runs without error when batch sizes vary across calibration inputs.""" + from executorch.export.recipe import QuantizationRecipe + + class PassthroughQuantizer(TorchAOPT2EQuantizer): + def annotate(self, model): + return model + + def validate(self, model): + pass + + def calibration_inputs_fn(): + for batch_size in (2, 4, 8): + yield (torch.randn(batch_size, 10),) + + recipe = QuantizationRecipe( + quantizers=[PassthroughQuantizer()], + is_qat=False, + dynamic_batch_size=True, + calibration_inputs_fn=calibration_inputs_fn, + ) + stage = QuantizeStage(recipe) + # Use batch size 2 for the example input so torch.export does not + # specialize dim 0 as the constant 1. + context = {"example_inputs": {"forward": [(torch.randn(2, 10),)]}} + artifact = PipelineArtifact( + data={"forward": torch.nn.Linear(10, 5)}, + context=context, + ) + stage.run(artifact) + self.assertIn("forward", stage.get_artifacts().data) + + class TestEmptyPassDictIsNotApplied(unittest.TestCase): """`EdgeProgramManager.transform` deep-copies the graph and weights of every method the pass dict does not name, so handing it an empty dict copies @@ -662,6 +1231,15 @@ def _manager(self) -> Mock: manager.exported_program.return_value = Mock() return manager + def test_edge_program_manager_stage_may_follow_partitioning(self) -> None: + # The point of the stage for a delegate recipe: its passes act on what + # the partitioner left outside the delegates, so it has to be able to + # run after TO_EDGE_TRANSFORM_AND_LOWER and not only after TO_EDGE. + self.assertEqual( + set(EdgeProgramManagerTransformStage().valid_predecessor_stages), + {StageType.TO_EDGE, StageType.TO_EDGE_TRANSFORM_AND_LOWER}, + ) + def test_edge_program_manager_stage_skips_empty_transform(self) -> None: manager = self._manager() stage = EdgeProgramManagerTransformStage( diff --git a/extension/llm/batching/CMakeLists.txt b/extension/llm/batching/CMakeLists.txt index 6d00f601b27..fac3cfff5f0 100644 --- a/extension/llm/batching/CMakeLists.txt +++ b/extension/llm/batching/CMakeLists.txt @@ -9,12 +9,16 @@ # scheduler and the executor seam are header-only and free of ExecuTorch runtime # types; the runner owns a thread, so this is a static library rather than an # INTERFACE target. +# +# extension_llm_batching_module is a separate target because it implements that +# seam against a program and a KV cache, and so carries the runtime types the +# seam itself is kept clear of. if(NOT EXECUTORCH_ROOT) set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) endif() -add_library(extension_llm_batching runner.cpp) +add_library(extension_llm_batching runner.cpp metrics.cpp) # std::optional, std::function, and std::future in the public headers. target_compile_features(extension_llm_batching PUBLIC cxx_std_17) target_include_directories( @@ -26,8 +30,22 @@ target_compile_options(extension_llm_batching PUBLIC ${_common_compile_options}) find_package(Threads REQUIRED) target_link_libraries(extension_llm_batching PUBLIC Threads::Threads) +add_library(extension_llm_batching_module module_executor.cpp) +target_link_libraries( + extension_llm_batching_module + PUBLIC extension_llm_batching extension_llm_cache extension_module + extension_tensor + PRIVATE extension_llm_sampler +) +target_include_directories( + extension_llm_batching_module PUBLIC ${_common_include_directories} +) +target_compile_options( + extension_llm_batching_module PUBLIC ${_common_compile_options} +) + install( - TARGETS extension_llm_batching + TARGETS extension_llm_batching extension_llm_batching_module EXPORT ExecuTorchTargets DESTINATION ${CMAKE_INSTALL_LIBDIR} INCLUDES diff --git a/extension/llm/batching/metrics.cpp b/extension/llm/batching/metrics.cpp new file mode 100644 index 00000000000..365e2e1bb79 --- /dev/null +++ b/extension/llm/batching/metrics.cpp @@ -0,0 +1,140 @@ +/* + * 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. + */ + +// The report formatters, kept out of metrics.h so the stream headers they need +// stay out of every translation unit that merely reads a counter. + +#include + +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +namespace detail { + +inline std::string fixed(double v, int places) { + std::ostringstream os; + os << std::fixed << std::setprecision(places) << v; + return os.str(); +} + +inline std::string ms(std::int64_t microseconds) { + return fixed(static_cast(microseconds) / 1000.0, 2); +} + +} // namespace detail + +// Returned rather than printed: these headers stay free of ExecuTorch's +// logging, so the caller decides where it goes. +std::string format_report(const GenerationMetrics& m) { + std::ostringstream os; + os << "session " << m.sid << ": " << m.n_prompt_tokens << " prompt + " + << m.n_generated_tokens << " generated tokens\n" + << " ttft " << detail::ms(m.ttft_us()) << " ms = queue " + << detail::ms(m.queue_wait_us()) << " + prefill " + << detail::ms(m.prefill_span_us()) << " (" + << detail::fixed(m.prefill_tokens_per_sec(), 1) << " tok/s, " + << m.n_prefill_steps + << (m.n_prefill_steps == 1 ? " step)\n" : " steps)\n"); + if (m.itl_count > 0) { + // The decode span and its rate are omitted: both follow from the mean + // below, which is decode time over decode tokens by construction. + os << " decode " << detail::ms(m.itl_sum_us / m.itl_count) + << " ms/token, max " << detail::ms(m.itl_max_us) << " -> " + << detail::fixed(m.decode_tokens_per_sec(), 1) << " tok/s over " + << m.n_decode_steps << " steps\n"; + } + os << " total " << detail::ms(m.e2e_us()) << " ms\n"; + return os.str(); +} + +std::string format_report(const EngineMetrics& m) { + std::ostringstream os; + const double wall = m.wall_us(); + const auto share = [wall](std::int64_t part) { + return wall > 0.0 ? detail::fixed(100.0 * part / wall, 0) + : std::string("0"); + }; + // An empty bucket has no mean. Printing 0.00 ms would read as a measurement + // rather than an absence, and the count of zero is the informative part. + const auto bucket = [](double mean_us, std::uint64_t n) { + return n > 0 ? detail::ms(static_cast(mean_us)) + " ms" + : std::string("--"); + }; + + os << "engine\n" + << " run " << detail::fixed(wall / 1e6, 2) << " s wall, " + << m.steps << " steps"; + if (m.steps_failed > 0) { + os << " (" << m.steps_failed << " failed)"; + } + os << ", " << detail::fixed(100.0 * m.idle_fraction(), 1) << "% idle\n"; + if (m.init_us > 0) { + os << " init " << detail::ms(m.init_us) + << " ms executor setup, before the wall above\n"; + } + os << " generations " << m.generations_completed << " of " + << m.generations_started << ": " << m.finished_stop_token << " stop, " + << m.finished_token_limit << " limit, " << m.finished_cancelled + << " cancelled, " << m.finished_failed << " failed\n" + << " sessions " << m.peak_concurrent_generations + << " peak concurrent, " << m.sessions_refused << " refused at capacity\n" + << "\n" + // A partition, so the three shares add up. The gap between the first two + // is what packing prefill alongside decode cost the decodes. + << " step time decode-only " + << bucket(m.mean_decode_only_step_us(), m.decode_only_steps()) << " x" + << m.decode_only_steps() << " (" << share(m.decode_only_latency_sum_us) + << "%)\n" + << " mixed " + << bucket(m.mean_mixed_step_us(), m.mixed_steps()) << " x" + << m.mixed_steps() << " (" << share(m.mixed_latency_sum_us) << "%)\n" + << " prefill-only " + << bucket(m.mean_prefill_only_step_us(), m.prefill_only_steps()) << " x" + << m.prefill_only_steps() << " (" << share(m.prefill_only_latency_sum_us) + << "%)\n" + << "\n" + << " decode " << detail::fixed(m.decode_only_tokens_per_sec(), 1) + << " tok/s on decode-only steps at " + << detail::fixed(m.mean_decode_step_sessions(), 2) << " sessions/step\n" + << " " + << detail::fixed(m.mean_decode_us_per_token() / 1000.0, 2) + << " ms/token per session across all " << m.steps_with_decode + << " decode steps\n" + << " context " << detail::fixed(m.mean_context_per_step(), 0) + << " tokens resident/step, max " << m.context_max << " per session\n" + << "\n" + << " scheduler admitted " + << detail::fixed(m.mean_admitted_decode_sessions(), 2) << " of " + << detail::fixed(m.mean_ready_decode_sessions(), 2) + << " ready decode sessions (" + << detail::fixed(100.0 * m.admitted_ratio(), 1) << "%)\n" + << " ttft mean " << detail::ms(m.mean_ttft_us()) << " ms, min " + << detail::ms(m.min_ttft_us()) << " ms, max " << detail::ms(m.ttft_max_us) + << " ms over " << m.ttft_count + << " generations\n" + // Wall-clock rates: these move with the prompt-to-generation mix, so they + // describe the workload as much as the engine. + << " delivered " << m.total_generated_tokens << " generated -> " + << detail::fixed(m.generated_tokens_per_sec(), 1) << " tok/s wall\n" + << " processed " << m.model_input_tokens() << " tokens -> " + << detail::fixed(m.model_input_tokens_per_sec(), 1) << " tok/s (" + << m.decode_tokens_total << " decode + " << m.prefill_tokens_total + << " prompt)\n"; + return os.str(); +} + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/batching/metrics.h b/extension/llm/batching/metrics.h new file mode 100644 index 00000000000..312c62d5e1c --- /dev/null +++ b/extension/llm/batching/metrics.h @@ -0,0 +1,416 @@ +/* + * 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 + +// What a batched run measured, in two tiers. +// +// GenerationMetrics covers one generation and is published on its handle, so a +// caller reads it beside finish_reason(). EngineMetrics covers the engine and +// is owned by the runner. +// +// There is deliberately no session tier: Session::position() already reports a +// session's context, and a session's generations are recovered by grouping on +// GenerationMetrics::sid. Keeping generations separate is the point -- a second +// generation on a warm session has a different profile from a cold one, and an +// average over the two hides exactly that. +// +// Free of ExecuTorch runtime types, like the rest of these headers: no ET_LOG, +// no exceptions, and nothing here allocates outside the format_report calls. + +#include +#include +#include +#include + +#include +#include // ET_EXPERIMENTAL + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +// Monotonic: these are durations, and a wall-clock adjustment mid-run would +// otherwise produce negative ones. +using MetricsClock = std::chrono::steady_clock; +using MetricsTime = MetricsClock::time_point; + +ET_EXPERIMENTAL inline std::int64_t us_between( + MetricsTime from, + MetricsTime to) { + return std::chrono::duration_cast(to - from) + .count(); +} + +// A default-constructed time_point means the event never happened, which is +// distinct from it happening at time zero: the clock's epoch is the process's, +// so no real event lands there. +ET_EXPERIMENTAL inline bool stamped(MetricsTime t) { + return t.time_since_epoch().count() != 0; +} + +// One generation's timeline and counts. The finish reason is not duplicated +// here; GenerationHandle::finish_reason() already carries it, and two copies +// would be free to disagree. +struct ET_EXPERIMENTAL GenerationMetrics { + SessionId sid = 0; + + // generate_async() was called. Taken on the caller's thread, before the + // request is queued. + MetricsTime t_submit{}; + // The first batch this generation appeared in. Everything between here and + // t_submit is time the scheduler did not pick it. + MetricsTime t_first_step{}; + MetricsTime t_first_token{}; + MetricsTime t_end{}; + + std::int64_t n_prompt_tokens = 0; + std::int64_t n_generated_tokens = 0; + std::int32_t n_prefill_steps = 0; + std::int32_t n_decode_steps = 0; + + // Caller-visible inter-token latency, excluding the gap to the first token, + // which is TTFT. The first token in a later callback gets the elapsed gap; + // further tokens in that callback get zero-time samples. Summary rather than + // samples: exact mean at four scalars, where a per-token vector would cost + // 8 KB on a long generation. + std::int64_t itl_count = 0; + std::int64_t itl_sum_us = 0; + std::int64_t itl_min_us = std::numeric_limits::max(); + std::int64_t itl_max_us = 0; + + // Time to first token: what a caller waits before anything appears. + std::int64_t ttft_us() const { + return stamped(t_submit) && stamped(t_first_token) + ? us_between(t_submit, t_first_token) + : 0; + } + + // The queueing share of ttft_us(). Large means the scheduler was busy, not + // that prefill was slow. + std::int64_t queue_wait_us() const { + return stamped(t_submit) && stamped(t_first_step) + ? us_between(t_submit, t_first_step) + : 0; + } + + // The compute share of ttft_us(). Includes other sessions' work in the steps + // this one's prefill was spread across, so read it with n_prefill_steps. + std::int64_t prefill_span_us() const { + return stamped(t_first_step) && stamped(t_first_token) + ? us_between(t_first_step, t_first_token) + : 0; + } + + // Generation proper, after the first token. + std::int64_t decode_span_us() const { + return stamped(t_first_token) && stamped(t_end) + ? us_between(t_first_token, t_end) + : 0; + } + + std::int64_t e2e_us() const { + return stamped(t_submit) && stamped(t_end) ? us_between(t_submit, t_end) + : 0; + } + + double itl_mean_us() const { + return itl_count > 0 ? static_cast(itl_sum_us) / itl_count : 0.0; + } + + // Zero rather than the sentinel when no gap was ever sampled, as happens to + // any generation that produced a single token. Mirrors min_ttft_us(). + std::int64_t min_itl_us() const { + return itl_count > 0 ? itl_min_us : 0; + } + + // What this one caller saw, which is not the engine's aggregate rate. + double decode_tokens_per_sec() const { + if (itl_count == 0) { + return 0.0; + } + return itl_sum_us > 0 ? 1e6 * static_cast(itl_count) / itl_sum_us + : std::numeric_limits::infinity(); + } + + // Prompt tokens over the wall time this generation's prefill took. Like the + // decode rate above it is what this caller experienced, so it includes any + // work sharing those steps -- a prompt that waits behind another session's + // chunks reports a lower rate, which is what that caller actually got. + double prefill_tokens_per_sec() const { + const std::int64_t span = prefill_span_us(); + return span > 0 ? 1e6 * static_cast(n_prompt_tokens) / span : 0.0; + } +}; + +// The engine's own view, accumulated on the engine thread and read once it has +// stopped. +struct ET_EXPERIMENTAL EngineMetrics { + std::uint64_t steps = 0; + std::uint64_t steps_failed = 0; + + // Summed over steps, so dividing by `steps` gives the mean. Sessions, not + // tokens: a prefill chunk is one session and many tokens. + std::uint64_t decode_sessions_total = 0; + std::uint64_t prefill_sessions_total = 0; + // Generations that could have decoded in this step: alive and past their + // first token. Prefilling generations are excluded because they are not + // waiting on a decode slot, they are doing their own work. Against + // decode_sessions_total this says how much of the eligible work the scheduler + // ran, without naming any scheduler's limits. + std::uint64_t ready_total = 0; + + std::int64_t step_latency_sum_us = 0; + std::int64_t step_latency_max_us = 0; + + // Every token the engine processed. Task::is_decode classifies each one + // exactly, so these are complete. + // + // There is deliberately no rate over every step that held them. A step + // mixing both kinds runs them in one forward pass over one weight read, so + // no share of its latency belongs to either, and dividing by the time of + // those steps would make packing prefill alongside decode -- which raises + // total throughput -- look like a decode regression. The rate that is safe + // to publish is decode_only_tokens_per_sec(), taken over steps that held no + // prefill at all, where the attribution is not in question. + std::uint64_t decode_tokens_total = 0; + std::uint64_t prefill_tokens_total = 0; + // The part of decode_tokens_total that ran on decode-only steps, so it can + // be divided by their time. The comparable figure to a fixed-batch decode + // benchmark, which measures exactly this shape of step. + std::uint64_t decode_only_tokens = 0; + + // Steps holding at least one task of each kind. A step can hold both, so + // these overlap by mixed_steps(). + std::uint64_t steps_with_decode = 0; + std::uint64_t steps_with_prefill = 0; + + // step_latency_sum_us split by what the step held. A partition, unlike the + // two counts above: every step is exactly one of these three, so they sum to + // step_latency_sum_us. Comparing the first two prices the scheduler's choice + // to pack prefill alongside decode, which no single blended mean can show. + std::int64_t decode_only_latency_sum_us = 0; + std::int64_t mixed_latency_sum_us = 0; + std::int64_t prefill_only_latency_sum_us = 0; + + // Committed session length summed across the batch, once per step: what + // attention had to carry. Generic -- executor.h defines a session's length, + // saying nothing about how the state is stored -- and the largest reason one + // decode step costs more than another. + std::int64_t context_sum = 0; + std::int64_t context_max = 0; // longest single session seen in any step + + // Step latency charged once to every decode session in the step. A latency + // may be counted for several sessions because each of them really did wait + // it -- unlike time as a cost, waiting is not divided up. Includes mixed + // steps, where a decode stuck behind a prefill chunk waited the whole thing. + std::int64_t decode_session_time_sum_us = 0; + + // Opens the executor refused for being at capacity, and the most generations + // seen installed at once. The peak is sampled once per executed step, so a + // generation that began and ended between two steps is not counted. + std::uint64_t sessions_refused = 0; + std::uint64_t peak_concurrent_generations = 0; + + std::uint64_t generations_started = 0; + std::uint64_t generations_completed = 0; + std::uint64_t finished_stop_token = 0; + std::uint64_t finished_token_limit = 0; + std::uint64_t finished_cancelled = 0; + std::uint64_t finished_failed = 0; + + // Over generations that reached a first token, which is not every + // completion: one cancelled or failed during prefill has no TTFT to report. + std::uint64_t ttft_count = 0; + std::int64_t ttft_sum_us = 0; + std::int64_t ttft_min_us = std::numeric_limits::max(); + std::int64_t ttft_max_us = 0; + + std::int64_t total_prompt_tokens = 0; + std::int64_t total_generated_tokens = 0; + + MetricsTime t_first_step{}; + MetricsTime t_last_step{}; + + // Executor::initialize(), timed here because nothing else can: it runs on the + // engine thread after the constructor has already returned, so a caller has + // no two points to measure between. Outside wall_us(), which starts at the + // first step -- folding it in would hide one-time setup inside the run. + std::int64_t init_us = 0; + + double wall_us() const { + return stamped(t_first_step) && stamped(t_last_step) + ? static_cast(us_between(t_first_step, t_last_step)) + : 0.0; + } + + // Mean decode sessions per step, and the mean that were eligible. Raw + // counts: normalising against a scheduler's decode limit would tie these to + // one scheduler, and against a workload smaller than that limit it would + // report idle capacity that no prompt existed to fill. + // + // Both divide by every step, so the pair is comparable. For the width of an + // actual decode forward, which is what a batched matmul sees, use + // mean_decode_step_sessions(). + double mean_admitted_decode_sessions() const { + return steps > 0 ? static_cast(decode_sessions_total) / steps : 0.0; + } + + double mean_ready_decode_sessions() const { + return steps > 0 ? static_cast(ready_total) / steps : 0.0; + } + + // Decode sessions per step that actually held a decode. Each contributes one + // token, so this is also the decode token width of the forward. + double mean_decode_step_sessions() const { + return steps_with_decode > 0 + ? static_cast(decode_sessions_total) / steps_with_decode + : 0.0; + } + + // Every session in the step, both kinds. How full the forward pass was, + // which is the batching question; the decode figures above are the + // scheduling one. + double mean_step_sessions() const { + return steps > 0 + ? static_cast(decode_sessions_total + prefill_sessions_total) / + steps + : 0.0; + } + + // Below 1 the scheduler is holding eligible work back rather than running + // out of it. Slightly under 1 is normal: a generation is briefly eligible + // but unqueued between its output being handled and its continuation being + // submitted. + double admitted_ratio() const { + return ready_total > 0 + ? static_cast(decode_sessions_total) / ready_total + : 0.0; + } + + // What the engine processed. Distinct from the generation tier's totals, + // which count what callers were given: a generation's first token comes out + // of a prefill step, and prompt tokens are processed but never emitted. + std::uint64_t model_input_tokens() const { + return decode_tokens_total + prefill_tokens_total; + } + + std::int64_t total_tokens() const { + return total_prompt_tokens + total_generated_tokens; + } + + double mean_step_tokens() const { + return steps > 0 ? static_cast(model_input_tokens()) / steps : 0.0; + } + + // Wall time the engine spent outside execute(): waiting for work, draining + // commands, running callbacks. + double idle_fraction() const { + const double wall = wall_us(); + return wall > 0.0 ? 1.0 - static_cast(step_latency_sum_us) / wall + : 0.0; + } + + // What callers were given, per second. The comparable figure when swapping + // executors: a speculative one feeds one token per decode step and returns + // several, so processed tokens would stay flat while this rises. + double generated_tokens_per_sec() const { + const double wall = wall_us(); + return wall > 0.0 ? 1e6 * static_cast(total_generated_tokens) / wall + : 0.0; + } + + // What the engine put through the model, per second. Against the rate above + // it shows how much output each processed token bought. + double model_input_tokens_per_sec() const { + const double wall = wall_us(); + return wall > 0.0 ? 1e6 * static_cast(model_input_tokens()) / wall + : 0.0; + } + + // Steps that held both kinds at once. Implied by the overlap rather than + // counted: every step holds decode, prefill, or both. + std::uint64_t mixed_steps() const { + const std::uint64_t overlap = steps_with_decode + steps_with_prefill; + return overlap > steps ? overlap - steps : 0; + } + + double mean_ttft_us() const { + return ttft_count > 0 ? static_cast(ttft_sum_us) / ttft_count : 0.0; + } + + // The sentinel is never shown: with no samples there is no minimum. + // Mean wall time a decode session waited per step it took part in, mixed + // steps included. The engine-side counterpart to per-generation inter-token + // latency, and a cross-check on it. Weighted by sessions, so it is not + // directly comparable to the unweighted bucket means above -- read it as + // what sessions experienced, not as what a step cost. + double mean_decode_us_per_token() const { + return decode_sessions_total > 0 + ? static_cast(decode_session_time_sum_us) / + decode_sessions_total + : 0.0; + } + + // Steps holding only one kind. Derived, because every step is decode-only, + // prefill-only, or mixed. + std::uint64_t decode_only_steps() const { + const std::uint64_t mixed = mixed_steps(); + return steps_with_decode > mixed ? steps_with_decode - mixed : 0; + } + + std::uint64_t prefill_only_steps() const { + const std::uint64_t mixed = mixed_steps(); + return steps_with_prefill > mixed ? steps_with_prefill - mixed : 0; + } + + double mean_decode_only_step_us() const { + const std::uint64_t n = decode_only_steps(); + return n > 0 ? static_cast(decode_only_latency_sum_us) / n : 0.0; + } + + double mean_mixed_step_us() const { + const std::uint64_t n = mixed_steps(); + return n > 0 ? static_cast(mixed_latency_sum_us) / n : 0.0; + } + + double mean_prefill_only_step_us() const { + const std::uint64_t n = prefill_only_steps(); + return n > 0 ? static_cast(prefill_only_latency_sum_us) / n : 0.0; + } + + // Decode throughput measured only where no prefill shared the forward, so + // every microsecond in the denominator was spent on these tokens. The one + // rate here that compares across engines, since a fixed-batch decode + // benchmark runs exactly this shape of step. + double decode_only_tokens_per_sec() const { + return decode_only_latency_sum_us > 0 ? 1e6 * + static_cast(decode_only_tokens) / decode_only_latency_sum_us + : 0.0; + } + + // Committed tokens the batch carried, averaged over steps. + double mean_context_per_step() const { + return steps > 0 ? static_cast(context_sum) / steps : 0.0; + } + + std::int64_t min_ttft_us() const { + return ttft_count > 0 ? ttft_min_us : 0; + } +}; +// Human-readable reports. Defined in metrics.cpp so a runner that never formats +// one does not pull and into its translation units -- +// everything above is trivially inline and allocation-free, these are not. +ET_EXPERIMENTAL std::string format_report(const GenerationMetrics& m); +ET_EXPERIMENTAL std::string format_report(const EngineMetrics& m); + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/batching/module_executor.cpp b/extension/llm/batching/module_executor.cpp new file mode 100644 index 00000000000..a2e920445af --- /dev/null +++ b/extension/llm/batching/module_executor.cpp @@ -0,0 +1,639 @@ +/* + * 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. + */ + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +using ::executorch::extension::make_tensor_ptr; +using ::executorch::runtime::Error; +using ::executorch::runtime::Result; + +namespace { + +bool is_supported_logits_type(::executorch::aten::ScalarType type) { + using ScalarType = ::executorch::aten::ScalarType; + return type == ScalarType::Float || type == ScalarType::Half || + type == ScalarType::BFloat16 || type == ScalarType::UInt16; +} + +// Constant methods carry no delegate, so the layout reads with the program +// loaded and the method not. Sizing is the caller's and is left unset. +Result config_from_program(Module& module) { + const auto read_int = [&module](const char* name) -> std::optional { + const auto r = module.execute(name); + if (!r.ok() || r->empty() || !r->at(0).isInt()) { + return std::nullopt; + } + return r->at(0).toInt(); + }; + const auto read_ints = + [&module](const char* name) -> std::optional> { + const auto r = module.execute(name); + if (!r.ok() || r->empty() || !r->at(0).isTensor()) { + return std::nullopt; + } + const auto t = r->at(0).toTensor(); + if (t.scalar_type() != ::executorch::aten::ScalarType::Int) { + return std::nullopt; + } + const int32_t* p = t.const_data_ptr(); + return std::vector(p, p + t.numel()); + }; + + const auto n_caches = read_int("get_n_caches"); + const auto kv_heads = read_ints("get_kv_heads"); + const auto head_dims = read_ints("get_head_dims"); + const auto windows = read_ints("get_windows"); + ET_CHECK_OR_RETURN_ERROR( + n_caches && kv_heads && head_dims && windows, + InvalidArgument, + "ModuleExecutor: the program publishes no KV layout"); + const auto n = static_cast(*n_caches); + ET_CHECK_OR_RETURN_ERROR( + kv_heads->size() == n && head_dims->size() == n && windows->size() == n, + InvalidArgument, + "ModuleExecutor: the published KV layout names %zu caches inconsistently", + n); + + cache::CacheConfig cfg{}; + cfg.n_layers = static_cast(n); + cfg.layers.reserve(n); + for (size_t l = 0; l < n; ++l) { + cache::LayerConfig lc{}; + lc.n_kv_heads = (*kv_heads)[l]; + lc.head_dim = (*head_dims)[l]; + lc.policy = (*windows)[l] > 0 + ? cache::LayerPolicy{cache::LayerPolicy::Kind::Ring, (*windows)[l]} + : cache::LayerPolicy{cache::LayerPolicy::Kind::Flat, 0}; + cfg.layers.push_back(lc); + } + return cfg; +} + +std::uint64_t nondeterministic_seed() { + std::random_device device; + return device(); +} + +} // namespace + +Result ModuleExecutor::build_step( + const BatchInput& batch) { + // Flatten the batch and truncate whatever it reopens; execute() declares each + // slice to the cache as it runs it. A per-sequence cursor carries the batch's + // own writes, so consecutive chunks of one prompt abut and only the first can + // reopen committed ground. Every input is checked before any is truncated, so + // a refusal leaves the cache untouched. + Step step; + const std::size_t total = batch.size(); + step.tokens.reserve(total); + step.positions.reserve(total); + step.logit_indices.reserve(batch.inputs.size()); + + step.seq_ids.reserve(total); + // Truncations the batch asks for, held until every input has been checked. + std::vector> rewinds; + // Where each sequence stands mid-batch: the cache still reports what it held + // before the step, so the batch's own writes live here. + std::unordered_map cursor; + + for (const Input& input : batch.inputs) { + const auto seq_it = sessions_.find(input.sid); + if (seq_it == sessions_.end()) { + ET_LOG(Error, "build_step: session %" PRId64 " is not open", input.sid); + return Error::InvalidArgument; + } + const std::int32_t seq_id = seq_it->second.seq_id; + if (input.size == 0 || !input.tokens || + input.offset > input.tokens->size() || + input.size > input.tokens->size() - input.offset) { + ET_LOG( + Error, + "build_step: session %" PRId64 " gave a slice its tokens do not hold", + input.sid); + return Error::InvalidArgument; + } + + const std::int64_t start = static_cast(input.position) + + static_cast(input.offset); + const auto [cursor_it, first_for_seq] = + cursor.try_emplace(seq_id, ctl_->next_pos(seq_id)); + int& at = cursor_it->second; + if (start > at) { + // Positions nothing attended, and nothing later reaches back to fill. + ET_LOG( + Error, + "build_step: session %" PRId64 " starts at %" PRId64 + " over a sequence holding %d", + input.sid, + start, + at); + return Error::InvalidArgument; + } + if (start < at) { + if (!first_for_seq) { + // Its predecessor in this batch has already been laid down, so a + // rewind now would truncate committed cells for a step whose + // positions repeat and cannot be placed. + ET_LOG( + Error, + "build_step: session %" PRId64 " overlaps its earlier input", + input.sid); + return Error::InvalidArgument; + } + if (start == 0) { + // Emptying a sequence hands its id back, and the step names it. + ET_LOG( + Error, + "build_step: session %" PRId64 " reopens from the start", + input.sid); + return Error::InvalidArgument; + } + rewinds.emplace_back(seq_id, static_cast(start)); + at = static_cast(start); + } + + const std::int64_t end = start + static_cast(input.size); + if (end > max_session_tokens_) { + ET_LOG( + Error, + "build_step: session %" PRId64 " reaches %" PRId64 " of %d cells", + input.sid, + end, + max_session_tokens_); + return Error::OutOfResources; + } + + const Token* slice = input.tokens->data() + input.offset; + for (std::size_t k = 0; k < input.size; ++k) { + step.tokens.push_back(static_cast(slice[k])); + } + for (std::size_t k = 0; k < input.size; ++k) { + step.positions.push_back(start + static_cast(k)); + } + step.seq_ids.insert(step.seq_ids.end(), input.size, seq_id); + at = static_cast(end); + step.logit_indices.push_back( + input.produce_output ? static_cast(step.tokens.size()) - 1 : -1); + } + + for (const auto& [seq_id, from] : rewinds) { + if (!ctl_->seq_rm(seq_id, from, std::nullopt)) { + ET_LOG(Error, "build_step: sequence %d would not truncate", seq_id); + return Error::Internal; + } + } + return step; +} + +ModuleExecutor::ModuleExecutor( + std::unique_ptr module, + std::shared_ptr cache, + int max_sessions, + int max_session_tokens, + std::string backend_id, + std::string method, + std::int32_t vocab_size, + int max_step_tokens, + LogitsToKeepMode logits_to_keep_mode) + : install_guard_(cache), + module_(std::move(module)), + ctl_(cache->as()), + max_sessions_(max_sessions), + max_session_tokens_(max_session_tokens), + backend_id_(std::move(backend_id)), + method_(std::move(method)), + vocab_size_(vocab_size), + max_step_tokens_(max_step_tokens), + logits_to_keep_mode_(logits_to_keep_mode) {} + +ModuleExecutor::~ModuleExecutor() = default; + +Result> ModuleExecutor::create( + std::unique_ptr module, + int max_sessions, + int max_session_tokens, + int kv_dtype, + int initial_capacity, + std::string cache_kind, + std::string method) { + if (module == nullptr) { + ET_LOG(Error, "ModuleExecutor: no program"); + return Error::InvalidArgument; + } + if (max_sessions <= 0 || max_session_tokens <= 0) { + ET_LOG(Error, "ModuleExecutor: session limits must be positive"); + return Error::InvalidArgument; + } + const Error load_error = + module->load(); // no-op once the caller has loaded it + if (load_error != Error::Ok) { + ET_LOG(Error, "ModuleExecutor: the program did not load"); + return load_error; + } + + const auto max_context_length = read_max_context_length(*module); + if (!max_context_length.ok()) { + ET_LOG(Error, "ModuleExecutor: the program's metadata is malformed"); + return max_context_length.error(); + } + if (max_session_tokens > *max_context_length) { + ET_LOG( + Error, + "ModuleExecutor: max session tokens %d exceeds model context length %" PRId64, + max_session_tokens, + *max_context_length); + return Error::InvalidArgument; + } + const auto logits_mode_result = read_logits_to_keep_mode(*module); + if (!logits_mode_result.ok()) { + ET_LOG(Error, "ModuleExecutor: the program's metadata is malformed"); + return logits_mode_result.error(); + } + const LogitsToKeepMode logits_mode = *logits_mode_result; + if (logits_mode == LogitsToKeepMode::Last) { + ET_LOG( + Error, + "ModuleExecutor: logits-to-keep mode last is incompatible with batched execution"); + return Error::NotSupported; + } + + auto cfg = config_from_program(*module); + if (!cfg.ok()) { + return cfg.error(); + } + if (max_sessions > std::numeric_limits::max() / max_session_tokens) { + ET_LOG(Error, "ModuleExecutor: total cache capacity exceeds int range"); + return Error::InvalidArgument; + } + cfg->capacity = max_sessions * max_session_tokens; + cfg->kv_dtype = kv_dtype; + if (initial_capacity >= 0) { + cfg->initial_capacity = initial_capacity; + } + if (!cache::valid(*cfg)) { + ET_LOG(Error, "ModuleExecutor: the program's layout is unusable"); + return Error::InvalidProgram; + } + + const auto meta = module->method_meta(method); + if (!meta.ok()) { + ET_LOG(Error, "ModuleExecutor: %s has no metadata", method.c_str()); + return meta.error(); + } + + const std::size_t expected_inputs = + logits_mode == LogitsToKeepMode::Selected ? 3 : 2; + if (meta->num_inputs() != expected_inputs) { + ET_LOG( + Error, + "ModuleExecutor: %s expects %zu inputs for its logits mode, got %zu", + method.c_str(), + expected_inputs, + meta->num_inputs()); + return Error::InvalidProgram; + } + + const auto tokens_info = meta->input_tensor_meta(0); + const auto positions_info = meta->input_tensor_meta(1); + if (!tokens_info.ok() || !positions_info.ok()) { + ET_LOG(Error, "ModuleExecutor: %s inputs must be tensors", method.c_str()); + return Error::InvalidProgram; + } + const auto token_sizes = tokens_info->sizes(); + const auto position_sizes = positions_info->sizes(); + if (tokens_info->scalar_type() != ::executorch::aten::ScalarType::Long || + token_sizes.size() != 2 || token_sizes[0] != 1 || token_sizes[1] <= 0 || + positions_info->scalar_type() != ::executorch::aten::ScalarType::Long || + position_sizes.size() != 1 || position_sizes[0] != token_sizes[1]) { + ET_LOG( + Error, + "ModuleExecutor: %s must take Long[1, T] tokens and Long[T] positions", + method.c_str()); + return Error::InvalidProgram; + } + if (logits_mode == LogitsToKeepMode::Selected) { + const auto selector_info = meta->input_tensor_meta(2); + if (!selector_info.ok() || + selector_info->scalar_type() != ::executorch::aten::ScalarType::Long || + selector_info->sizes().size() != 1) { + ET_LOG( + Error, + "ModuleExecutor: %s selected logits selector must be rank-one Long", + method.c_str()); + return Error::InvalidProgram; + } + } + + if (meta->num_outputs() == 0) { + ET_LOG(Error, "ModuleExecutor: %s publishes no outputs", method.c_str()); + return Error::InvalidProgram; + } + const auto logits_info = meta->output_tensor_meta(0); + if (!logits_info.ok()) { + ET_LOG(Error, "ModuleExecutor: %s has no logits metadata", method.c_str()); + return logits_info.error(); + } + const auto logits_sizes = logits_info->sizes(); + if (logits_sizes.size() < 2 || logits_sizes[logits_sizes.size() - 1] <= 0 || + !is_supported_logits_type(logits_info->scalar_type())) { + ET_LOG( + Error, + "ModuleExecutor: %s logits must have supported dtype and shape [..., vocab]", + method.c_str()); + return Error::InvalidProgram; + } + const auto published_vocab_size = read_vocab_size(*module); + if (!published_vocab_size.ok()) { + ET_LOG( + Error, "ModuleExecutor: invalid get_vocab_size for %s", method.c_str()); + return published_vocab_size.error(); + } + const auto vocab_size = check_vocab_size( + *published_vocab_size, logits_sizes[logits_sizes.size() - 1]); + if (!vocab_size.ok()) { + ET_LOG( + Error, + "ModuleExecutor: invalid get_vocab_size for %s output width", + method.c_str()); + return vocab_size.error(); + } + + std::string backend_id; + for (std::size_t i = 0; i < meta->num_backends(); ++i) { + const auto name = meta->get_backend_name(i); + if (!name.ok()) { + ET_LOG( + Error, "ModuleExecutor: %s has an unnamed delegate", method.c_str()); + return name.error(); + } + if (backend_id.empty()) { + backend_id = name.get(); + } else if (backend_id != name.get()) { + ET_LOG( + Error, + "ModuleExecutor: %s spans more than one backend, so which holds the " + "cache is ambiguous", + method.c_str()); + return Error::InvalidProgram; + } + } + if (backend_id.empty()) { + ET_LOG(Error, "ModuleExecutor: %s delegates to nothing", method.c_str()); + return Error::InvalidProgram; + } + + auto built = + cache::CacheFactory::global().build(backend_id, cache_kind, *cfg); + if (!built.ok()) { + ET_LOG( + Error, + "ModuleExecutor: backend %s registers no %s cache", + backend_id.c_str(), + cache_kind.c_str()); + return built.error(); + } + std::shared_ptr cache = built.get(); + if (cache->as() == nullptr) { + ET_LOG(Error, "ModuleExecutor: the cache carries no sequence identity"); + return Error::InvalidType; + } + + return std::unique_ptr(new ModuleExecutor( + std::move(module), + std::move(cache), + max_sessions, + max_session_tokens, + std::move(backend_id), + std::move(method), + *vocab_size, + token_sizes[1], + logits_mode)); +} + +bool ModuleExecutor::initialize() { + // The delegate resolves the cache from this key while the method loads. + ::executorch::runtime::BackendOptions<1> options; + ::executorch::runtime::LoadBackendOptionsMap options_map; + if (install_guard_.set_option(options) != Error::Ok || + options_map.set_options(backend_id_.c_str(), options.view()) != + Error::Ok) { + ET_LOG(Error, "ModuleExecutor: could not name the cache to the backend"); + return false; + } + if (module_->load_method( + method_, + /*planned_memory=*/nullptr, + /*event_tracer=*/nullptr, + &options_map) != Error::Ok) { + ET_LOG(Error, "ModuleExecutor: could not load %s", method_.c_str()); + return false; + } + return true; +} + +std::optional ModuleExecutor::open_session() { + if (static_cast(sessions_.size()) >= max_sessions_) { + return std::nullopt; + } + const std::optional seq_id = ctl_->seq_new(); + if (!seq_id) { + return std::nullopt; + } + const SessionId session = next_session_++; + sessions_.emplace(session, SessionState{*seq_id, nullptr}); + return session; +} + +void ModuleExecutor::close_session(SessionId session) { + const auto it = sessions_.find(session); + if (it == sessions_.end()) { + return; + } + // Frees the cells and hands the sequence id back. The session id is not. + ctl_->seq_rm(it->second.seq_id, 0, std::nullopt); + sessions_.erase(it); +} + +void ModuleExecutor::set_sampling( + SessionId session, + const SamplingParams& params, + std::optional seed) { + const auto it = sessions_.find(session); + if (it == sessions_.end()) { + return; + } + // One sampler per generation, carrying its own generator state from here on. + it->second.sampler = std::make_unique( + vocab_size_, + params.temperature, + params.top_p, + seed.value_or(nondeterministic_seed())); + it->second.sampler->set_topk(params.top_k); +} + +bool ModuleExecutor::execute(const BatchInput& batch, BatchOutput& out) { + out.outputs.clear(); + out.outputs.resize(batch.inputs.size()); + + const Result step = build_step(batch); + if (!step.ok()) { + return false; + } + + // A batch wider than the method was traced at runs as several forwards. They + // go in order, so a slice attends the cells its predecessors wrote, and each + // input's logits row falls in exactly one of them. + const int total = static_cast(step->tokens.size()); + for (int off = 0; off < total; off += max_step_tokens_) { + const int n = std::min(max_step_tokens_, total - off); + // Placement checks the forward's token count against the declaration, so + // each slice declares its own. + if (!ctl_->declare_step(std::vector( + step->seq_ids.begin() + off, step->seq_ids.begin() + off + n))) { + ET_LOG(Error, "ModuleExecutor: the cache refused a slice of %d", n); + return false; + } + auto tokens = make_tensor_ptr( + {1, n}, + std::vector( + step->tokens.begin() + off, step->tokens.begin() + off + n)); + auto positions = make_tensor_ptr( + {n}, + std::vector( + step->positions.begin() + off, step->positions.begin() + off + n)); + std::vector selector_values; + std::vector selected_inputs; + if (logits_to_keep_mode_ == LogitsToKeepMode::Selected) { + for (std::size_t i = 0; i < step->logit_indices.size(); ++i) { + const int row = step->logit_indices[i]; + if (row >= off && row < off + n) { + selector_values.push_back(row - off); + selected_inputs.push_back(i); + } + } + if (selector_values.empty()) { + selector_values.push_back(n - 1); + } + } + + const int expected_rows = logits_to_keep_mode_ == LogitsToKeepMode::Selected + ? static_cast(selector_values.size()) + : n; + auto result = [&]() -> Result> { + if (logits_to_keep_mode_ == LogitsToKeepMode::Selected) { + auto selector = + make_tensor_ptr({expected_rows}, std::move(selector_values)); + return module_->execute(method_, {tokens, positions, selector}); + } + return module_->execute(method_, {tokens, positions}); + }(); + if (!result.ok()) { + ET_LOG( + Error, + "ModuleExecutor: %s failed with 0x%x", + method_.c_str(), + static_cast(result.error())); + return false; + } + if (result->empty() || !result->at(0).isTensor()) { + ET_LOG(Error, "ModuleExecutor: %s returned no logits", method_.c_str()); + return false; + } + // Non-const: the sampler reduces each row in place. Each is read once. + auto logits = result->at(0).toTensor(); + if (logits.dim() < 2 || logits.size(logits.dim() - 1) != vocab_size_ || + logits.numel() != + static_cast(expected_rows) * vocab_size_) { + ET_LOG( + Error, + "ModuleExecutor: %s returned invalid logits shape for %d rows and vocab %d", + method_.c_str(), + expected_rows, + vocab_size_); + return false; + } + + if (logits_to_keep_mode_ == LogitsToKeepMode::Selected) { + for (std::size_t row = 0; row < selected_inputs.size(); ++row) { + const std::size_t input_index = selected_inputs[row]; + const SessionId session = batch.inputs[input_index].sid; + const std::optional token = + sample_row(logits, static_cast(row), session); + if (!token) { + return false; + } + out.outputs[input_index] = Output{session, {*token}}; + } + } else { + for (std::size_t i = 0; i < batch.inputs.size(); ++i) { + const int row = step->logit_indices[i]; + if (row < off || row >= off + n) { + continue; // another slice's row, or a dropped chunk prediction + } + const SessionId session = batch.inputs[i].sid; + const std::optional token = + sample_row(logits, row - off, session); + if (!token) { + return false; + } + out.outputs[i] = Output{session, {*token}}; + } + } + } + return true; +} + +std::optional ModuleExecutor::sample_row( + ::executorch::aten::Tensor& logits, + int row, + SessionId session) { + const auto it = sessions_.find(session); + if (it == sessions_.end() || it->second.sampler == nullptr) { + ET_LOG( + Error, + "ModuleExecutor: session %" PRId64 " has no sampling policy", + session); + return std::nullopt; + } + if (row >= logits.numel() / vocab_size_) { + ET_LOG(Error, "ModuleExecutor: logits hold no row %d", row); + return std::nullopt; + } + // A one-row view over the model's own output: sample_from_logits reduces in + // place and reads the last dimension. + auto one_row = make_tensor_ptr( + {vocab_size_}, + static_cast(logits.mutable_data_ptr()) + + static_cast(row) * vocab_size_ * + ::executorch::runtime::elementSize(logits.scalar_type()), + logits.scalar_type()); + return static_cast(sample_from_logits(*one_row, *it->second.sampler)); +} + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/batching/module_executor.h b/extension/llm/batching/module_executor.h new file mode 100644 index 00000000000..e0832f5fbc6 --- /dev/null +++ b/extension/llm/batching/module_executor.h @@ -0,0 +1,141 @@ +/* + * 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 + +// An Executor that runs an ExecuTorch Module over a registered KV cache, built +// from the layout the program publishes. A session is one cache sequence, a +// batch is one forward carrying every input's tokens on a single axis, and the +// cache's mask keeps the sequences apart. + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include // ET_EXPERIMENTAL + +namespace executorch { +namespace extension { +namespace llm { + +class Sampler; + +namespace batching { + +namespace cache = ::executorch::extension::llm::cache; + +class ET_EXPERIMENTAL ModuleExecutor : public Executor { + public: + ~ModuleExecutor() override; + + // Builds the cache from the layout `module` publishes and pairs it with the + // backend, which is read from the program -- so the method's attention must + // be delegated to just one. The method itself loads in initialize(); the + // program must be loaded and its method must not be, since the delegate + // resolves the cache while that load runs. + // + // Capacity is `max_sessions` x `max_session_tokens` cells exactly, and + // open_session() holds the count, so exhaustion is unreachable rather than + // handled. `kv_dtype` is the ET ScalarType K/V is stored in; a negative + // `initial_capacity` leaves the pools to grow from their own default. + // `cache_kind` must name a builder that carries batch control -- a cache + // serving one sequence cannot back a batch of them. + // + // Returns an error for unusable limits, no published KV layout, a method + // spanning several backends, or no such cache for the backend it names. A + // method that will not load is reported by initialize(). + static ::executorch::runtime::Result> create( + std::unique_ptr module, + int max_sessions, + int max_session_tokens, + int kv_dtype, + int initial_capacity = -1, + std::string cache_kind = cache::kind::kBatchedCell, + std::string method = "forward"); + + // The widest step this method takes, from the shape its token input was + // traced at. A wider batch is sliced; a narrower one leaves the forward + // partly unused. + std::size_t preferred_batch_tokens() const override { + return static_cast(max_step_tokens_); + } + + // Loads the method here so the delegate that resolves the cache binds on the + // thread that runs it. + bool initialize() override; + + std::optional open_session() override; + void close_session(SessionId session) override; + void set_sampling( + SessionId session, + const SamplingParams& params, + std::optional seed) override; + bool execute(const BatchInput& batch, BatchOutput& out) override; + + private: + struct SessionState { + std::int32_t seq_id; + std::unique_ptr sampler; + }; + + struct Step { + std::vector tokens; + std::vector positions; + std::vector seq_ids; + std::vector logit_indices; + }; + + ::executorch::runtime::Result build_step(const BatchInput& batch); + + ModuleExecutor( + std::unique_ptr module, + std::shared_ptr cache, + int max_sessions, + int max_session_tokens, + std::string backend_id, + std::string method, + std::int32_t vocab_size, + int max_step_tokens, + LogitsToKeepMode logits_to_keep_mode); + + // Draw the token an input produced from its row of `logits`, which the + // session's sampler consumes in place. + std::optional + sample_row(::executorch::aten::Tensor& logits, int row, SessionId session); + + // Ordered so the module dies first, releasing the delegate that resolved the + // cache before the registry entry naming it goes. + cache::InstallGuard install_guard_; + std::unique_ptr module_; + cache::BatchControl* const ctl_; + int max_sessions_; + int max_session_tokens_; + std::string backend_id_; + std::string method_; + // The method's logits width, so a sampler can be built by its policy. + std::int32_t vocab_size_; + int max_step_tokens_; + LogitsToKeepMode logits_to_keep_mode_; + + SessionId next_session_ = 1; // never reused, unlike the cache's sequence ids + std::unordered_map sessions_; +}; + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/batching/runner.cpp b/extension/llm/batching/runner.cpp index 0df1a8f6ae7..2ea0e69ba92 100644 --- a/extension/llm/batching/runner.cpp +++ b/extension/llm/batching/runner.cpp @@ -66,6 +66,7 @@ struct GenerationHandleState { CompletionPhase phase = CompletionPhase::Pending; std::optional reason; std::string error_message; + GenerationMetrics metrics; std::atomic cancelled{false}; }; @@ -120,11 +121,14 @@ class TerminalCompletion { ~TerminalCompletion() { if (state_) { - finish(TerminalOutcome::failed({})); + finish(TerminalOutcome::failed({}), GenerationMetrics{}); } } - void finish(TerminalOutcome outcome) { + // Metrics ride alongside the outcome rather than inside it: they describe + // the generation, not the reason it ended. Published here, with the reason + // and under the same lock, so a handle never shows one without the other. + void finish(TerminalOutcome outcome, const GenerationMetrics& metrics) { assert(state_); auto state = std::move(state_); { @@ -132,6 +136,7 @@ class TerminalCompletion { assert(state->phase == CompletionPhase::DeliveringCallback); state->reason = outcome.reason; state->error_message = std::move(outcome.error_message); + state->metrics = metrics; state->phase = CompletionPhase::Done; } state->cv.notify_all(); @@ -184,7 +189,8 @@ void finalize_terminal( if (!callback_result.succeeded) { outcome = TerminalOutcome::failed(std::move(callback_result.error_message)); } - completion->finish(std::move(outcome)); + // Rejected before admission, so there is no timeline to report. + completion->finish(std::move(outcome), GenerationMetrics{}); } // --- GenerationHandle ------------------------------------------------------ @@ -236,6 +242,17 @@ std::string GenerationHandle::error_message() const { return state_->error_message; } +// Published with the reason, which lands only after the terminal callback +// returns. Reading this from inside that callback yields an empty snapshot, +// exactly as finish_reason() yields nullopt there. +GenerationMetrics GenerationHandle::metrics() const { + if (!state_) { + return {}; + } + std::lock_guard lock(state_->mutex); + return state_->metrics; +} + // Everything the runner owns. Held by shared_ptr from both Runner and every // Session, so a Session outliving its Runner finds a stopped object rather // than a dangling one. @@ -270,6 +287,11 @@ class RunnerImpl : public std::enable_shared_from_this { GenConfig config, GenerationCallback on_update); + // Engine-thread data, so only stable once that thread is joined. + EngineMetrics metrics() const { + return metrics_; + } + private: enum class Lifecycle { Running, Stopping, Stopped }; @@ -281,6 +303,10 @@ class RunnerImpl : public std::enable_shared_from_this { // Shared with every handle, so cancelling needs no route back to the // runner and works after it is gone. std::shared_ptr state; + GenerationMetrics m; + // Engine-side only. An inter-token gap needs the previous delivery, and + // the published metrics keep only the summary, not the last timestamp. + MetricsTime last_token_at{}; }; // Start-only data. The sampling policy is installed on the executor at @@ -433,14 +459,27 @@ class RunnerImpl : public std::enable_shared_from_this { CallbackResult dispatch_update_( const Generation& generation, GenerationUpdate update); - void deliver_claimed_terminal_( + FinishReason deliver_claimed_terminal_( const Generation& generation, TerminalCompletion completion, TerminalOutcome outcome); // Invoke the terminal callback and publish state for a detached generation. - void complete_generation_(Generation generation, TerminalOutcome outcome); - void complete_request_(GenerationRequest request, TerminalOutcome outcome); + // + // `on_engine_thread` is false only on the post-shutdown path out of + // generate_async(), which runs on the caller's thread. The engine may still + // be draining there, so that path publishes to the handle but must leave the + // engine's own counters alone. + void complete_generation_( + Generation generation, + TerminalOutcome outcome, + bool on_engine_thread); + void complete_request_( + GenerationRequest request, + TerminalOutcome outcome, + bool on_engine_thread); + // Engine thread only: rolls one finished generation into metrics_. + void record_completion_(const GenerationMetrics& m, FinishReason reason); std::optional detach_active_generation_(SessionId session); void complete_active_generation_(SessionId session, TerminalOutcome outcome); void fail_active_generation_after_callback_( @@ -470,6 +509,7 @@ class RunnerImpl : public std::enable_shared_from_this { // Kept after records close to enforce executor IDs are lifetime-unique. std::unordered_set issued_session_ids_; TaskId next_tid_ = 1; + EngineMetrics metrics_; std::thread engine_; }; @@ -557,6 +597,10 @@ void Runner::shutdown() { impl_->shutdown(); } +EngineMetrics Runner::metrics() const { + return impl_->metrics(); +} + std::future> Runner::open_session_async() { return impl_->open_session_async(); } @@ -627,7 +671,10 @@ void RunnerImpl::run_() { // Before the loop and before any command is answered, so a caller is never // handed a session for an executor that did not come up, and so one-time // setup is not charged to whichever generation happened to go first. - if (!executor_.initialize()) { + const MetricsTime init_start = MetricsClock::now(); + const bool ready = executor_.initialize(); + metrics_.init_us = us_between(init_start, MetricsClock::now()); + if (!ready) { // Stop without running work. The drain below still answers whatever was // queued while this was starting, so no caller is left waiting. std::lock_guard lock(control_mutex_); @@ -673,7 +720,9 @@ void RunnerImpl::run_() { for (auto& entry : open_sessions) { if (entry.second) { complete_generation_( - std::move(*entry.second), TerminalOutcome::cancelled()); + std::move(*entry.second), + TerminalOutcome::cancelled(), + /*on_engine_thread=*/true); } executor_.close_session(entry.first); } @@ -714,7 +763,13 @@ void RunnerImpl::process_pending_commands_() { } void RunnerImpl::process_command_(OpenCommand command) { - auto sid = is_running_() ? executor_.open_session() : std::nullopt; + const bool running = is_running_(); + auto sid = running ? executor_.open_session() : std::nullopt; + if (running && !sid) { + // At capacity. Counted apart from a refusal caused by the runner stopping, + // which says nothing about how loaded the executor was. + ++metrics_.sessions_refused; + } bool newly_issued = false; bool published = false; if (sid) { @@ -750,7 +805,9 @@ void RunnerImpl::process_command_(CloseCommand command) { } if (retired->active_generation) { complete_generation_( - std::move(*retired->active_generation), TerminalOutcome::cancelled()); + std::move(*retired->active_generation), + TerminalOutcome::cancelled(), + /*on_engine_thread=*/true); } executor_.close_session(command.session); } @@ -789,9 +846,166 @@ bool RunnerImpl::execute_one_batch_() { return false; } + // Composition is read here, before to_batch_input moves the Inputs out and + // drops is_decode with the rest of the scheduling fields. + // + // Decode tasks are one sequence each. A session's prefill can arrive as + // several chunks which are not necessarily adjacent -- DecodeFirstScheduler + // rotates, taking one chunk per session per pass, so two sessions prefilling + // together interleave as A, B, A, B. Track every session seen rather than + // comparing with the previous one, which would count each chunk as a new + // sequence and charge the step to the generation repeatedly. + std::uint64_t decode_sessions = 0; + std::uint64_t prefill_sessions = 0; + std::uint64_t decode_tokens = 0; + std::uint64_t prefill_tokens = 0; + std::vector prefilling; // small: bounded by the batch width + const MetricsTime step_start = MetricsClock::now(); + for (const Task& task : tasks) { + bool first_chunk = false; + if (task.is_decode) { + ++decode_sessions; + decode_tokens += task.input.size; + } else { + prefill_tokens += task.input.size; + first_chunk = + std::find(prefilling.begin(), prefilling.end(), task.input.sid) == + prefilling.end(); + if (first_chunk) { + ++prefill_sessions; + prefilling.push_back(task.input.sid); + } + } + // Charge the step to the generation once, however many chunks it brought. + if (!task.is_decode && !first_chunk) { + continue; + } + auto session = sessions_.find(task.input.sid); + if (session == sessions_.end() || !session->second.active_generation) { + continue; + } + GenerationMetrics& m = session->second.active_generation->m; + if (!stamped(m.t_first_step)) { + m.t_first_step = step_start; + } + if (task.is_decode) { + ++m.n_decode_steps; + } else { + ++m.n_prefill_steps; + } + } + // Generations eligible for a decode slot, admitted or not. Past their first + // token, so a generation still prefilling is not counted as held back when + // it is simply busy elsewhere. Against decode_sessions this is what + // separates a scheduler holding work back from there being no work. + // + // The same pass records how many generations are installed at once, so the + // batch widths above can be read against the concurrency that was available. + std::uint64_t live_generations = 0; + for (const auto& entry : sessions_) { + const auto& generation = entry.second.active_generation; + if (!generation) { + continue; + } + ++live_generations; + if (stamped(generation->m.t_first_token)) { + ++metrics_.ready_total; + } + } + metrics_.peak_concurrent_generations = + std::max(metrics_.peak_concurrent_generations, live_generations); + + // Committed context the batch carries into the forward: what attention has + // to cover. Read from the input positions, which say nothing about how the + // executor stores the state. + // + // Counted once per session, at the end of its furthest chunk. A prompt split + // across several chunks of one step is one context, not one per chunk, and + // scheduler.h lets a session appear more than once for exactly that reason. + std::vector> context_ends; + for (const Task& task : tasks) { + const auto end = static_cast(task.input.position) + + static_cast(task.input.offset) + + static_cast(task.input.size); + metrics_.context_max = std::max(metrics_.context_max, end); + auto entry = std::find_if( + context_ends.begin(), context_ends.end(), [&](const auto& seen) { + return seen.first == task.input.sid; + }); + if (entry == context_ends.end()) { + context_ends.emplace_back(task.input.sid, end); + } else { + entry->second = std::max(entry->second, end); + } + } + for (const auto& entry : context_ends) { + metrics_.context_sum += entry.second; + } + + // The forward and the batch handed to it, timed apart from step_start. The + // scans above are measurement, and folding them into the buckets would bias + // exactly the numbers metrics.h offers for comparison against other engines; + // they are also the parts that grow with concurrency, so the bias would not + // be constant. step_start still marks the step for the generation timeline. + const MetricsTime exec_start = MetricsClock::now(); BatchInput batch = to_batch_input(tasks); BatchOutput out; const bool ok = executor_.execute(batch, out); + const MetricsTime step_end = MetricsClock::now(); + + const std::int64_t latency = us_between(exec_start, step_end); + ++metrics_.steps; + metrics_.decode_sessions_total += decode_sessions; + metrics_.prefill_sessions_total += prefill_sessions; + // Only what the model is known to have taken in. A failed execute leaves + // what it processed unknown -- that is why the batch is condemned and its + // sessions poisoned -- so counting the attempt as throughput would credit + // work that may never have happened. The time is still counted below, + // because it was really spent, and steps_failed records the attempt. + if (ok) { + metrics_.decode_tokens_total += decode_tokens; + metrics_.prefill_tokens_total += prefill_tokens; + } + metrics_.step_latency_sum_us += latency; + metrics_.step_latency_max_us = + std::max(metrics_.step_latency_max_us, latency); + if (!stamped(metrics_.t_first_step)) { + metrics_.t_first_step = step_start; + } + metrics_.t_last_step = step_end; + if (decode_tokens > 0) { + ++metrics_.steps_with_decode; + // Charged once per session: each of them waited this whole step. + metrics_.decode_session_time_sum_us += + latency * static_cast(decode_sessions); + } + if (prefill_tokens > 0) { + ++metrics_.steps_with_prefill; + } + // Exactly one of the three, so the sums partition step_latency_sum_us. The + // three step-count conditions below use the same two predicates, so the + // counts partition `steps` too. + assert( + (decode_tokens > 0 || prefill_tokens > 0) && + "a non-empty batch carries tokens of at least one kind"); + if (decode_tokens > 0 && prefill_tokens > 0) { + metrics_.mixed_latency_sum_us += latency; + } else if (decode_tokens > 0) { + metrics_.decode_only_latency_sum_us += latency; + // Attributable, because no prefill shared this forward. Gated on `ok` for + // the same reason as decode_tokens_total: a failed execute leaves what the + // model consumed unknown. The latency above is still counted, since it was + // really spent and the three buckets have to partition the total. + if (ok) { + metrics_.decode_only_tokens += decode_tokens; + } + } else { + metrics_.prefill_only_latency_sum_us += latency; + } + if (!ok) { + ++metrics_.steps_failed; + } + if (!is_running_()) { return true; // discard an in-flight result after the stop boundary } @@ -937,6 +1151,10 @@ GenerationHandle RunnerImpl::generate_async( request.generation.stop_tokens = std::move(config.stop_tokens); request.generation.on_update = std::move(on_update); request.generation.state = state; + // The caller's thread, before the request is queued: the wait a caller sees + // starts here, not when the engine gets round to it. + request.generation.m.sid = session; + request.generation.m.t_submit = MetricsClock::now(); auto handle = GenerationHandle(state); bool admitted = false; @@ -954,7 +1172,10 @@ GenerationHandle RunnerImpl::generate_async( // After shutdown nothing drains the inbox, so complete synchronously instead // of admitting a start that can never report completion. - complete_request_(std::move(request), TerminalOutcome::cancelled()); + complete_request_( + std::move(request), + TerminalOutcome::cancelled(), + /*on_engine_thread=*/false); return handle; } @@ -1005,24 +1226,40 @@ std::shared_ptr> RunnerImpl::build_initial_delta_( } void RunnerImpl::start_generation_(GenerationRequest request) { + // Counted on arrival at the engine, not on successful install: every path + // below ends in a completion, so deferring this would let completions + // exceed starts. + ++metrics_.generations_started; if (!is_running_() || request.generation.state->cancelled.load()) { - complete_request_(std::move(request), TerminalOutcome::cancelled()); + complete_request_( + std::move(request), + TerminalOutcome::cancelled(), + /*on_engine_thread=*/true); return; } auto session = sessions_.find(request.session); if (session == sessions_.end()) { complete_request_( - std::move(request), TerminalOutcome::failed("session is not open")); + std::move(request), + TerminalOutcome::failed("session is not open"), + /*on_engine_thread=*/true); return; } auto& record = session->second; if (auto rejection = validate_generation_start_(request, record)) { - complete_request_(std::move(request), std::move(*rejection)); + complete_request_( + std::move(request), std::move(*rejection), /*on_engine_thread=*/true); return; } const auto start_position = record.position(); + // The caller's own delta, captured before build_initial_delta_ moves it and + // before any carried token is prepended. The generation tier counts what + // callers gave; tokens actually fed to the model are EngineMetrics' + // model_input_tokens(). + request.generation.m.n_prompt_tokens = + static_cast(request.delta->size()); auto delta = build_initial_delta_(request, record); executor_.set_sampling(request.session, request.sampling, request.seed); @@ -1036,7 +1273,10 @@ void RunnerImpl::start_generation_(GenerationRequest request) { } if (!installed) { // Sampling began before the stop transition, but no task was submitted. - complete_request_(std::move(request), TerminalOutcome::cancelled()); + complete_request_( + std::move(request), + TerminalOutcome::cancelled(), + /*on_engine_thread=*/true); return; } @@ -1178,6 +1418,34 @@ std::optional RunnerImpl::prepare_output_( auto interpreted = interpret_output_(generation, *output); generation.remaining_tokens = interpreted.remaining_tokens; + // Counted here rather than at delivery: this is the one place that sees the + // emitted run with the generation in hand, and it covers both the + // completing and the continuing branch below, which move the tokens away. + if (!interpreted.emitted_tokens.empty()) { + const auto emitted = + static_cast(interpreted.emitted_tokens.size()); + const MetricsTime now = MetricsClock::now(); + generation.m.n_generated_tokens += emitted; + if (!stamped(generation.m.t_first_token)) { + generation.m.t_first_token = now; + // The first token's wait is TTFT. Further tokens in the same run have + // zero caller-visible latency between them. + const std::int64_t intra_burst = emitted - 1; + if (intra_burst > 0) { + generation.m.itl_count += intra_burst; + generation.m.itl_min_us = 0; + } + } else { + const std::int64_t gap = us_between(generation.last_token_at, now); + generation.m.itl_count += emitted; + generation.m.itl_sum_us += gap; + generation.m.itl_min_us = std::min( + generation.m.itl_min_us, emitted > 1 ? std::int64_t{0} : gap); + generation.m.itl_max_us = std::max(generation.m.itl_max_us, gap); + } + generation.last_token_at = now; + } + // The transcript grows by what the caller keeps, capped by what the executor // committed. The last emitted token lands only when it is fed back. record.advance(interpreted.committed_tokens); @@ -1282,7 +1550,10 @@ CallbackResult RunnerImpl::dispatch_update_( return invoke_callback(generation.on_update, std::move(update)); } -void RunnerImpl::deliver_claimed_terminal_( +// Returns the reason actually published, which is not the one passed in when +// the terminal callback throws. The engine counts that final reason, so the +// tallies agree with what the handle reports. +FinishReason RunnerImpl::deliver_claimed_terminal_( const Generation& generation, TerminalCompletion completion, TerminalOutcome outcome) { @@ -1293,12 +1564,16 @@ void RunnerImpl::deliver_claimed_terminal_( if (!callback_result.succeeded) { outcome = TerminalOutcome::failed(std::move(callback_result.error_message)); } - completion.finish(std::move(outcome)); + const FinishReason reason = outcome.reason; + completion.finish(std::move(outcome), generation.m); + return reason; } void RunnerImpl::complete_generation_( Generation generation, - TerminalOutcome outcome) { + TerminalOutcome outcome, + bool on_engine_thread) { + generation.m.t_end = MetricsClock::now(); std::optional completion; { // The claim is taken under the runner lock. User code still runs only after @@ -1319,14 +1594,54 @@ void RunnerImpl::complete_generation_( } completion.emplace(std::move(*claimed)); } - deliver_claimed_terminal_( + const FinishReason reason = deliver_claimed_terminal_( generation, std::move(*completion), std::move(outcome)); + if (on_engine_thread) { + record_completion_(generation.m, reason); + } +} + +void RunnerImpl::record_completion_( + const GenerationMetrics& m, + FinishReason reason) { + ++metrics_.generations_completed; + switch (reason) { + case FinishReason::StopToken: + ++metrics_.finished_stop_token; + break; + case FinishReason::NewTokenLimit: + ++metrics_.finished_token_limit; + break; + case FinishReason::Cancelled: + ++metrics_.finished_cancelled; + break; + case FinishReason::Failed: + ++metrics_.finished_failed; + break; + } + metrics_.total_prompt_tokens += m.n_prompt_tokens; + metrics_.total_generated_tokens += m.n_generated_tokens; + // Zero for a generation that never reached a first token. Counted + // separately from completions so the mean divides by the samples it has, + // and so the minimum stays untouched when there are none. + // Gated on the event, not on it having taken measurable time: a first token + // in the same microsecond as the submit is still a first token, and stamped() + // is what "happened" means everywhere else here. + if (stamped(m.t_first_token)) { + const std::int64_t ttft = m.ttft_us(); + ++metrics_.ttft_count; + metrics_.ttft_sum_us += ttft; + metrics_.ttft_min_us = std::min(metrics_.ttft_min_us, ttft); + metrics_.ttft_max_us = std::max(metrics_.ttft_max_us, ttft); + } } void RunnerImpl::complete_request_( GenerationRequest request, - TerminalOutcome outcome) { - complete_generation_(std::move(request.generation), std::move(outcome)); + TerminalOutcome outcome, + bool on_engine_thread) { + complete_generation_( + std::move(request.generation), std::move(outcome), on_engine_thread); } std::optional RunnerImpl::detach_active_generation_( @@ -1353,7 +1668,8 @@ void RunnerImpl::complete_active_generation_( if (!active) { return; } - complete_generation_(std::move(*active), std::move(outcome)); + complete_generation_( + std::move(*active), std::move(outcome), /*on_engine_thread=*/true); } void RunnerImpl::fail_active_generation_after_callback_( @@ -1366,11 +1682,18 @@ void RunnerImpl::fail_active_generation_after_callback_( if (!active) { return; } + active->m.t_end = MetricsClock::now(); auto completion = TerminalCompletion::try_claim(active->state); if (!completion) { return; } - completion->finish(TerminalOutcome::failed(std::move(error_message))); + completion->finish( + TerminalOutcome::failed(std::move(error_message)), active->m); + // This path claims the terminal itself instead of going through + // complete_generation_, so it has to count its own completion. Without this + // the engine would report fewer completions than starts, and precisely for + // the generations that failed most interestingly. + record_completion_(active->m, FinishReason::Failed); } } // namespace batching diff --git a/extension/llm/batching/runner.h b/extension/llm/batching/runner.h index b26d943fbe7..a3af2fde5c7 100644 --- a/extension/llm/batching/runner.h +++ b/extension/llm/batching/runner.h @@ -40,6 +40,7 @@ #include #include +#include #include #include #include // ET_EXPERIMENTAL @@ -147,6 +148,10 @@ class ET_EXPERIMENTAL GenerationHandle { // when valid() && done(); empty when no diagnostic is available. std::string error_message() const; + // This generation's timeline and counts, complete once done(). Empty on a + // default-constructed handle. + GenerationMetrics metrics() const; + private: friend class RunnerImpl; friend class Session; @@ -246,6 +251,11 @@ class ET_EXPERIMENTAL Runner { // from its callback. void shutdown(); + // What the engine measured. Read it after shutdown(): the counters are the + // engine thread's, so joining it is what makes them stable and visible. A + // call before then returns a torn snapshot. + EngineMetrics metrics() const; + private: std::shared_ptr impl_; }; diff --git a/extension/llm/batching/test/fake_executor.h b/extension/llm/batching/test/fake_executor.h index 1f51f416fed..a4d90cdb5ac 100644 --- a/extension/llm/batching/test/fake_executor.h +++ b/extension/llm/batching/test/fake_executor.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include #include #include @@ -17,6 +18,7 @@ #include #include #include +#include #include #include @@ -51,6 +53,9 @@ class FakeExecutor : public Executor { }; bool initialize() override { + if (initialize_delay.count() > 0) { + std::this_thread::sleep_for(initialize_delay); + } std::lock_guard lock(mutex_); ++initialize_calls_; return !fail_initialize; @@ -132,17 +137,19 @@ class FakeExecutor : public Executor { int capacity = 8; // Refuse to come up, so the runner should admit no work at all. bool fail_initialize = false; + // Stand in for real setup, so its cost is large enough to assert on. + std::chrono::milliseconds initialize_delay{0}; // Batch index from which execute() starts failing. Negative never fails. int fail_batches_from = -1; // Once a session has produced emit_before_stop tokens, every later one is // stop_token. Counted per session across the whole run, so a stop can be - // placed part way into a multi-token decode. A negative stop_token disables - // this. - Token stop_token = -1; + // placed part way into a multi-token decode. Unset disables this; a sentinel + // cannot, since Token is unsigned here and every value is a valid token. + std::optional stop_token; int emit_before_stop = 0; - // Tokens a decode step produces. 1 is a plain executor; more simulates a - // speculative one answering with the run it accepted plus the model's own - // next token. Prefill always produces one whatever this is. + // Tokens an output-producing step returns. Values above 1 simulate a + // speculative executor answering with an accepted run plus the next token. + std::size_t tokens_per_prefill = 1; std::size_t tokens_per_decode = 1; // Malformed answers. An Output carries only the tokens an input produced, so // the only ways to break the contract are to produce none, or to answer for @@ -240,7 +247,8 @@ class FakeExecutor : public Executor { // Task::is_decode. Good enough for a fake: the runner only ever feeds one // token to continue. std::vector produce(const Input& input) { - const std::size_t n = input.size == 1 ? tokens_per_decode : 1; + const std::size_t n = + input.size == 1 ? tokens_per_decode : tokens_per_prefill; std::vector produced; produced.reserve(n); for (std::size_t i = 0; i < n; ++i) { @@ -251,8 +259,8 @@ class FakeExecutor : public Executor { Token next_token(SessionId session) { const int n = ++produced_[session]; - if (stop_token >= 0 && n > emit_before_stop) { - return stop_token; + if (stop_token && n > emit_before_stop) { + return *stop_token; } auto it = sampling_.find(session); if (it == sampling_.end()) { diff --git a/extension/llm/batching/test/runner_test.cpp b/extension/llm/batching/test/runner_test.cpp index f021f2c429c..7e258802400 100644 --- a/extension/llm/batching/test/runner_test.cpp +++ b/extension/llm/batching/test/runner_test.cpp @@ -29,6 +29,7 @@ #include using executorch::extension::llm::batching::DecodeFirstScheduler; +using executorch::extension::llm::batching::EngineMetrics; using executorch::extension::llm::batching::Executor; using executorch::extension::llm::batching::FinishReason; using executorch::extension::llm::batching::GenConfig; @@ -1017,6 +1018,27 @@ TEST(GenerationTest, ImmediateStopTokenIsDeliveredAndWinsOverBudget) { << "the prompt produces STOP, so no continuation is scheduled"; } +TEST(SpeculativeTest, FirstBurstRecordsIntraBurstTokenLatencies) { + FakeExecutor executor; + executor.tokens_per_prefill = 3; + Fixture fixture(executor); + Session session = open(fixture.runner); + auto updates = std::make_shared(); + + GenerationHandle handle = generate(session, tokens(2), config(3), updates); + ASSERT_TRUE(updates->wait()); + handle.wait(); + + const auto metrics = handle.metrics(); + EXPECT_EQ(metrics.n_generated_tokens, 3); + EXPECT_EQ(metrics.itl_count, 2); + EXPECT_EQ(metrics.itl_sum_us, 0); + EXPECT_EQ(metrics.itl_min_us, 0); + EXPECT_EQ(metrics.itl_max_us, 0); + EXPECT_EQ( + metrics.decode_tokens_per_sec(), std::numeric_limits::infinity()); +} + // Regression: the runner counts a step's produced tokens as far as the // executor committed them, so the next turn resumes past them rather than on // top of them. The last token of a run is not committed until it is fed back, @@ -1663,6 +1685,316 @@ TEST(ShutdownTest, ConcurrentAdmissionDoesNotStrandCallers) { EXPECT_EQ(stranded.load(), 0); } +// --- engine-tier metrics --------------------------------------------------- +// +// EngineMetrics is only stable once the engine thread is joined, so every test +// here shuts the runner down before reading it. + +namespace { + +// The four terminal reasons partition the completions, so this is the tally +// that reveals a terminal path which forgot to account for itself. +void expect_balanced(const EngineMetrics& m, std::uint64_t expected_starts) { + EXPECT_EQ(m.generations_started, expected_starts); + EXPECT_EQ(m.generations_completed, m.generations_started) + << "every generation the engine started must also be counted as done"; + EXPECT_EQ( + m.finished_stop_token + m.finished_token_limit + m.finished_cancelled + + m.finished_failed, + m.generations_completed) + << "the per-reason tallies must partition the completions"; +} + +} // namespace + +TEST(EngineMetricsTest, CountsStepsSequencesAndTokens) { + FakeExecutor executor; + // Chunk 8 with a 20-token budget, so a 10-token prompt arrives as 8 + 2 and + // both chunks still fit in one batch. + Fixture fixture(executor, 4, 8); + Session session = open(fixture.runner); + auto updates = std::make_shared(); + + GenerationHandle handle = generate(session, tokens(10), config(3), updates); + ASSERT_TRUE(updates->wait()); + handle.wait(); + ASSERT_EQ(updates->finish(), FinishReason::NewTokenLimit); + + session = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + // One prefill step carrying both chunks, then one decode step per token + // after the first, which the prefill produced. + EXPECT_EQ(m.steps, 3u); + EXPECT_EQ(m.steps_failed, 0u); + EXPECT_EQ(m.steps_with_prefill, 1u); + EXPECT_EQ(m.steps_with_decode, 2u); + // Charged once per session per step, not once per chunk. + EXPECT_EQ(m.prefill_sessions_total, 1u); + EXPECT_EQ(m.decode_sessions_total, 2u); + EXPECT_EQ(m.prefill_tokens_total, 10u); + EXPECT_EQ(m.decode_tokens_total, 2u); + // The prompt plus the tokens fed back. The last token delivered is still + // pending, so it was never an input. + EXPECT_EQ(m.model_input_tokens(), 12u); + EXPECT_EQ(m.total_prompt_tokens, 10); + EXPECT_EQ(m.total_generated_tokens, 3); + EXPECT_EQ(m.ttft_count, 1u); + // The one sample is the minimum and the whole sum. Stated as a relation + // rather than "> 0" so it holds however fast the machine is. + EXPECT_EQ(m.min_ttft_us(), m.ttft_sum_us); + EXPECT_EQ(m.ttft_max_us, m.ttft_sum_us); + expect_balanced(m, 1); + EXPECT_EQ(m.finished_token_limit, 1u); +} + +TEST(EngineMetricsTest, StartsAndCompletionsBalanceAcrossMixedOutcomes) { + FakeExecutor executor; + executor.stop_token = 999; + executor.emit_before_stop = 1; + Fixture fixture(executor); + + // Runs to a stop token. + Session stopping = open(fixture.runner); + GenConfig stop_config = config(50); + stop_config.stop_tokens = {999}; + auto stopped = std::make_shared(); + generate(stopping, tokens(2), stop_config, stopped); + ASSERT_TRUE(stopped->wait()); + + // Rejected by validation once it reaches the engine. + Session invalid = open(fixture.runner); + auto rejected = std::make_shared(); + generate(invalid, tokens(2), config(0), rejected); + ASSERT_TRUE(rejected->wait()); + + // Cancelled explicitly. + Session cancelling = open(fixture.runner); + auto cancelled = std::make_shared(); + GenerationHandle doomed = + generate(cancelling, tokens(2), config(100000), cancelled); + doomed.cancel(); + ASSERT_TRUE(cancelled->wait()); + + stopping = Session{}; + invalid = Session{}; + cancelling = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + expect_balanced(m, 3); + EXPECT_EQ(m.finished_stop_token, 1u); + EXPECT_EQ(m.finished_failed, 1u) << "the rejected budget is a failure"; + EXPECT_EQ(m.finished_cancelled, 1u); + EXPECT_EQ(m.finished_token_limit, 0u); +} + +TEST(EngineMetricsTest, FailedBatchIsCountedWithoutCreditingItsTokens) { + FakeExecutor executor; + executor.fail_batches_from = 0; + Fixture fixture(executor); + Session session = open(fixture.runner); + auto updates = std::make_shared(); + + generate(session, tokens(4), config(2), updates); + ASSERT_TRUE(updates->wait()); + ASSERT_EQ(updates->finish(), FinishReason::Failed); + + session = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + EXPECT_EQ(m.steps, 1u); + EXPECT_EQ(m.steps_failed, 1u); + // The sequence was in the batch, but a failed execute leaves what the model + // consumed unknown, so its tokens are not credited as throughput. + EXPECT_EQ(m.prefill_sessions_total, 1u); + EXPECT_EQ(m.prefill_tokens_total, 0u); + EXPECT_EQ(m.model_input_tokens(), 0u); + EXPECT_EQ(m.total_generated_tokens, 0); + // Never reached a first token, so it contributes no TTFT sample and leaves + // the minimum untouched. + EXPECT_EQ(m.ttft_count, 0u); + EXPECT_EQ(m.min_ttft_us(), 0); + expect_balanced(m, 1); + EXPECT_EQ(m.finished_failed, 1u); +} + +TEST(EngineMetricsTest, SingleTokenGenerationReportsZeroMinimumItl) { + FakeExecutor executor; + Fixture fixture(executor); + Session session = open(fixture.runner); + auto updates = std::make_shared(); + + GenerationHandle handle = generate(session, tokens(2), config(1), updates); + ASSERT_TRUE(updates->wait()); + handle.wait(); + + // One token means no gap was ever sampled. The raw field still holds its + // sentinel; the accessor is what callers should read. + const auto m = handle.metrics(); + EXPECT_EQ(m.n_generated_tokens, 1); + EXPECT_EQ(m.itl_count, 0); + EXPECT_EQ(m.min_itl_us(), 0); + EXPECT_EQ(m.itl_mean_us(), 0.0); + EXPECT_EQ(m.itl_min_us, std::numeric_limits::max()); +} + +#if ET_HAS_EXCEPTIONS +// Regression: this path claims the terminal itself rather than going through +// complete_generation_, so it once published nothing and counted nothing, +// leaving completions permanently behind starts. +TEST(EngineMetricsTest, CallbackExceptionIsCountedAsACompletedFailure) { + FakeExecutor executor; + Fixture fixture(executor); + Session session = open(fixture.runner); + + GenerationHandle failed = session.generate_async( + tokens(2), config(4), [](const GenerationUpdate& update) { + if (!update.finish_reason) { + throw std::runtime_error("callback failed"); + } + }); + failed.wait(); + ASSERT_EQ(failed.finish_reason(), FinishReason::Failed); + + session = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + expect_balanced(m, 1); + EXPECT_EQ(m.finished_failed, 1u); + // The tokens delivered before the callback threw still count as generated. + EXPECT_GT(m.total_generated_tokens, 0); +} +#endif + +// The three latency buckets are a partition of step_latency_sum_us: every step +// holds decode, prefill, or both, and its whole latency lands in exactly one. +// Without this, a step kind added later could quietly go unaccounted. +TEST(EngineMetricsTest, LatencyBucketsPartitionTotalStepTime) { + FakeExecutor executor; + // Chunk 2 with a 10-token budget, so a 6-token prompt needs several chunks + // and the run produces decode-only, prefill-only, and mixed steps. + Fixture fixture(executor, 2, 2); + Session first = open(fixture.runner); + Session second = open(fixture.runner); + + auto first_updates = std::make_shared(); + generate(first, tokens(6), config(6), first_updates); + auto second_updates = std::make_shared(); + generate(second, tokens(6), config(6), second_updates); + ASSERT_TRUE(first_updates->wait()); + ASSERT_TRUE(second_updates->wait()); + + first = Session{}; + second = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + EXPECT_EQ( + m.decode_only_latency_sum_us + m.mixed_latency_sum_us + + m.prefill_only_latency_sum_us, + m.step_latency_sum_us); + EXPECT_EQ( + m.decode_only_steps() + m.mixed_steps() + m.prefill_only_steps(), + m.steps); + // Only the attributable subset feeds the decode-only rate. + EXPECT_LE(m.decode_only_tokens, m.decode_tokens_total); + expect_balanced(m, 2); +} + +TEST(EngineMetricsTest, ReportsContextConcurrencyAndRefusals) { + FakeExecutor executor; + executor.capacity = 2; + Fixture fixture(executor); + Session first = open(fixture.runner); + Session second = open(fixture.runner); + + // A third open has nowhere to go, which is invisible without the counter. + auto refused = fixture.runner.open_session_async(); + ASSERT_EQ(refused.wait_for(kTimeout), std::future_status::ready); + EXPECT_FALSE(refused.get().has_value()); + + auto first_updates = std::make_shared(); + generate(first, tokens(4), config(3), first_updates); + auto second_updates = std::make_shared(); + generate(second, tokens(4), config(3), second_updates); + ASSERT_TRUE(first_updates->wait()); + ASSERT_TRUE(second_updates->wait()); + + first = Session{}; + second = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + EXPECT_EQ(m.sessions_refused, 1u); + EXPECT_GE(m.peak_concurrent_generations, 1u); + EXPECT_LE(m.peak_concurrent_generations, 2u); + // Each session ends holding its prompt plus the tokens it was fed back, and + // context is summed across the batch, so the mean is at least one prompt. + EXPECT_GT(m.mean_context_per_step(), 0.0); + EXPECT_GE(m.context_max, 4); + // Counts, not the rate derived from them: a fake executor's whole run is a + // few microseconds, so us_between can floor it to 0 and the rate accessor + // then reports exactly 0.0. That says nothing about the metric. + EXPECT_GT(m.decode_only_steps(), 0u); + EXPECT_GT(m.decode_only_tokens, 0u); + EXPECT_GT(m.mean_decode_step_sessions(), 0.0); +} + +// Regression: context is a property of a session, not of a task. A prompt that +// arrives as several chunks of one step is one context; summing per chunk +// counted it once per chunk and inflated the total. +TEST(EngineMetricsTest, MultiChunkPromptCountsContextOncePerSession) { + FakeExecutor executor; + // Chunk 2 with a 6-token budget, so a 6-token prompt fits as 3 chunks in a + // single step. max_new_tokens 1 ends the generation on that step's output, + // leaving exactly one step to account for. + Fixture fixture(executor, 2, 2); + Session session = open(fixture.runner); + auto updates = std::make_shared(); + + generate(session, tokens(6), config(1), updates); + ASSERT_TRUE(updates->wait()); + ASSERT_EQ(updates->finish(), FinishReason::NewTokenLimit); + + session = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + ASSERT_EQ(m.steps, 1u) << "the whole prompt should fit in one step"; + // The session ends that step holding 6 tokens. Per-chunk summing would give + // 2 + 4 + 6 = 12. + EXPECT_EQ(m.context_sum, 6); + EXPECT_EQ(m.context_max, 6); + EXPECT_DOUBLE_EQ(m.mean_context_per_step(), 6.0); +} + +TEST(EngineMetricsTest, InitializationIsTimedAndKeptOutOfTheWall) { + FakeExecutor executor; + // Large enough to dwarf the fake's steps, so the comparison below is not a + // race between two similar durations. + executor.initialize_delay = std::chrono::milliseconds(50); + Fixture fixture(executor); + Session session = open(fixture.runner); + auto updates = std::make_shared(); + + generate(session, tokens(2), config(2), updates); + ASSERT_TRUE(updates->wait()); + + session = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + EXPECT_GE(m.init_us, 40000) << "setup should be measured, not dropped"; + // wall_us() starts at the first step. Were setup folded into it, the wall + // could not be shorter than setup. + EXPECT_GT(static_cast(m.init_us), m.wall_us()) + << "one-time setup must stay outside the run window"; +} + // --- executor initialization ----------------------------------------------- TEST(InitializeTest, RunsOnceBeforeAnythingElseIsAskedOfTheExecutor) { @@ -1695,6 +2027,7 @@ TEST(InitializeTest, FailureStopsTheRunnerAndRefusesSessions) { fixture.runner.shutdown(); EXPECT_EQ(executor.initialize_calls(), 1); + EXPECT_EQ(fixture.runner.metrics().steps, 0u); EXPECT_TRUE(executor.seen().empty()) << "no batch should have run"; EXPECT_TRUE(executor.opened().empty()) << "a session must not be opened on an executor that failed to start"; diff --git a/extension/llm/batching/types.h b/extension/llm/batching/types.h index 3bfa39b74c4..211b8a4be0f 100644 --- a/extension/llm/batching/types.h +++ b/extension/llm/batching/types.h @@ -26,7 +26,7 @@ namespace extension { namespace llm { namespace batching { -using Token = std::int64_t; +using Token = std::uint64_t; using SessionId = std::int64_t; using Position = std::int32_t; // Wide enough that a monotonically issued id cannot wrap in any realistic diff --git a/extension/llm/cache/README.md b/extension/llm/cache/README.md new file mode 100644 index 00000000000..0673956a077 --- /dev/null +++ b/extension/llm/cache/README.md @@ -0,0 +1,201 @@ +# Off-graph KV cache + +A KV cache that lives outside the exported graph. The runner creates it, the +backend writes into it during a forward, and neither holds a pointer to the +other. They meet through a string key. + +Keeping the cache out of the graph lets the runner do things the graph cannot +express: rewind a turn, clear between prompts, or hand one pool of memory to +several concurrent sequences. + +## Who uses what + +Three audiences touch this directory, and they need almost disjoint parts of +it. + +### The control plane + +The runner, or a batch executor. It decides *what* to cache and *when* to +discard it, and it runs between forwards, never during one. + +It uses `CacheFactory` to build a cache, `InstallGuard` to publish it, and one +runner-facing face for the rest of the session: + +```cpp +auto built = CacheFactory::global().build(kMLXBackendId, kind::kSingle, cfg); +if (!built.ok()) { return built.error(); } + +const std::shared_ptr kv = built.get(); +const InstallGuard guard{kv}; // published while in scope +guard.set_option(mlx_opts); // hand key to backend + +auto* ctl = kv->as(); +ctl->can_extend(n); ctl->rewind(len); ctl->clear(); +``` + +It includes `cache.h` and `cache_registry.h`. It never includes +`sequence_cache.h` or `cell_cache.h`, and never calls a planner face. It does +not know how bytes are arranged, only how much room is left and how to give +some back. + +### The backend + +The byte layer inside the delegate. It owns the actual tensors and runs during +a forward. + +At init it resolves the key and asks for its own face: + +```cpp +handle->cache_shared = CacheRegistry::global().get(cache_key); +handle->state.cache = handle->cache_shared->as(); +``` + +During a forward it asks a planner face where the bytes go: + +```cpp +auto plan = planner->plan(layer, position, T); // integers, no tensors +// ... write K/V into those rows, attend over those runs ... +planner->commit(*plan); +``` + +It includes the layout headers, because it subclasses them to attach its own +tensor storage. It is the only caller of `plan()`, `commit()`, and `place_step()`. + +### The cache implementer + +Someone adding a layout or a backend face. Subclass a neutral layout, add +whatever face your backend needs, and list them: + +```cpp +class MLXCellCache : public cache::CellCache, public MLXCache { + void* face(cache::FaceId id) override { + if (void* p = cache::CellCache::face(id)) { return p; } + return cache::expose(this, id); + } +}; +``` + +Then register a builder so the control plane can ask for it by name. Registration +is insertion-only: an empty builder or a duplicate `(backend_id, kind)` returns +`Error::InvalidArgument`, leaving any existing builder unchanged. + +## Faces + +A cache is owned as a `Cache*` and asked for the interface you want: + +```cpp +auto* ctl = cache->as(); // null if this cache does not offer it +``` + +`Cache` has one virtual method. Each face declares its own name, and an +implementation lists the faces it offers through `expose`. + +| | single sequence | pooled cells | +| ---------------- | ----------------- | -------------- | +| control plane | `SequenceControl` | `BatchControl` | +| backend | `SequencePlanner` | `CellStepper` | +| backend-specific | each backend names its own, such as `MLXCache` || + +Control-plane faces live in `cache.h`. A runner calls `as()` +on something it got from the registry, so it must see the face without choosing +a layout. Backend faces live with their layout in `sequence_cache.h` or +`cell_cache.h`, because only a byte layer calls them and it already includes +that header to construct the cache. + +**No RTTI.** The core avoids `dynamic_cast` so it can build with `-fno-rtti` +under `EXECUTORCH_OPTIMIZE_SIZE`. The `static_cast` inside `expose` also +applies the pointer adjustment a face at a non-zero offset needs, refuses to +compile if the type is not really a base, and is bound to its own name, so the +two cannot be mismatched. Because `as()` names `T::kFaceName`, asking for a +type that is not a face fails to compile instead of returning null. + +**Names, not an enum.** The set of faces is open. A backend adds one without +this directory learning about it: `MLXCache` declares its name in `MLXCache.h` +and `cache.h` never sees it. Names compare by pointer first and fall back to +`strcmp`, which covers a cache built in one shared object and queried from +another. + +A face name is a global ABI identifier. It must be non-null, remain stable, and +identify exactly one C++ interface across the core, every backend, and every +shared object. Reusing a name for an unrelated or incompatible interface makes +the erased pointer cast invalid. The raw lookup hook is protected; consumers use +`as()`, and each concrete cache must explicitly implement the faces it offers. + +## Layouts + +**`SequenceCache`** holds one sequence with a single logical length for the +whole model. Each layer is flat, keeping all history, or ring, sliding a +window, so a model that mixes both stays coherent. Offers `SequenceControl` and +`SequencePlanner`. + +**`CellCache`** holds many sequences over a shared pool of per-token cells. A +cell is freed once no sequence owns it. Offers `BatchControl` and `CellStepper`. + +## Cache kinds + +| constant | value | layout | +| -------------------- | --------------- | --------------------------------- | +| `kind::kSingle` | `single` | one sequence, per-layer runs | +| `kind::kBatchedCell` | `batched-cell` | many sequences over a shared pool | + +Kinds are strings so a backend can register a layout this directory has never +heard of. The constants name the kinds it does know about. Use them: a typo in +a literal is a runtime `NotFound`, while a typo in a constant does not compile. + +## Lifetimes + +`InstallGuard` is the only way to publish. `CacheRegistry::install` is private, +so an entry cannot outlive its owner and a second caller cannot clobber it. + +Three lifetimes overlap: + +- The **registry entry** must exist across every `load_method()` that resolves + the key. +- The **guard** controls that discoverability and may be destroyed after the + final such initialization. +- The **cache** may outlive the entry and guard. Each backend that resolved the + key holds its own `shared_ptr`. + +Destroying the guard unpublishes the key without invalidating an already +resolved cache. A later `load_method()` using that key fails, so the guard must +remain alive for as long as new delegates may still need to resolve it. + +## Two layers + +`cache.h`, `sequence_cache.h`, and `cell_cache.{h,cpp}` include nothing but the +C++ standard library. No tensors and no ExecuTorch. They describe where bytes +go using integers: which physical rows a step writes, which it reads, what the +mask should be. Failures come back as `bool` and `std::optional`. + +`cache_registry.{h,cpp}` is ExecuTorch-specific. It uses `Result`, `Error`, and +`ET_LOG`, but the stronger tie is its reason for existing. `DelegateHandle` is +opaque and backend options carry only strings, so a runner cannot pass the +backend a pointer. Publishing under a generated key works around that. Give a +framework where the cache can be handed to the op directly, and this layer +disappears. + +> The build does not yet honour this split. One `extension_llm_cache` target +> compiles both halves and links `executorch_core`, so the neutral core cannot +> currently be built without ExecuTorch. + +## Files + +``` +cache.h faces, the face mechanism, config neutral +sequence_cache.h SequenceCache, SequencePlanner, flat/ring neutral +cell_cache.{h,cpp} CellCache, CellStepper, the cell pool neutral +cache_registry.{h,cpp} CacheRegistry, CacheFactory, InstallGuard ExecuTorch +``` + +## Known gaps + +**`CacheConfig` fields do not mean the same thing to every layout.** `capacity` +is a position ceiling for `kSingle` and a count of pool slots for +`kBatchedCell`. `max_write` is read only by ring layers, which in turn ignore +`initial_capacity`. Splitting it into model-dictated shape and per-kind options +is the intended fix. + +**Registration relies on static-initializer side effects.** A builder is +registered only if the linker pulls its object file in. Whole-archive linking +of backends covers this today. If that changes, the symptom is a runtime +"no cache builder registered". diff --git a/extension/llm/cache/cache.h b/extension/llm/cache/cache.h index fb25e43f4a6..986c7a1b6a4 100644 --- a/extension/llm/cache/cache.h +++ b/extension/llm/cache/cache.h @@ -10,45 +10,72 @@ // Neutral, tensor-free, ET-independent KV-cache core shared across backends. A // cache exposes a runner-facing control face and a backend-facing planner face, -// recovered from the owning CacheBase*. Which pair it implements depends on the -// layout: one sequence over per-layer runs, or many sequences over a pool of -// per-token cells. +// recovered from the owning Cache* with as(). Which pair it implements +// depends on the layout: one sequence over per-layer runs, or many sequences +// over a pool of per-token cells. +// +// Here: the face machinery, the runner-facing faces, and the config a caller +// fills in, all usable without picking a layout. Each backend-facing planner +// face lives with its layout, in sequence_cache.h or cell_cache.h. #include +#include #include #include +#include // ET_EXPERIMENTAL + namespace executorch { namespace extension { namespace llm { namespace cache { -class SequenceControl; -class SequencePlanner; -class BatchControl; -class CellStepper; +// A face is named by a non-null string it declares itself, so a backend can +// add one without this header learning about it. Names are stable ABI +// identifiers: each must identify exactly one interface across all binaries. +using FaceId = const char*; + +// Pointer equality covers the common case. The strcmp catches a cache built in +// one shared object and queried from another, where the literals may differ. +ET_EXPERIMENTAL inline bool same_face(FaceId a, FaceId b) { + return a != nullptr && b != nullptr && (a == b || std::strcmp(a, b) == 0); +} + +// Hands back `self` as each face it names, or nullptr for one it does not. +// static_cast applies the pointer adjustment a face at a non-zero offset needs +// and refuses to compile if Self does not derive from it. Each cast is bound to +// its own name in the pack, so a name cannot be paired with the wrong face. +template +ET_EXPERIMENTAL void* expose(Self* self, FaceId id) { + void* out = nullptr; + const bool matched[] = { + (same_face(id, Fs::kFaceName) ? (out = static_cast(self), true) + : false)...}; + (void)matched; + return out; +} -// Registry ownership anchor. A cache returns `this` from the faces it -// implements and leaves the rest null. -class CacheBase { +// Registry ownership anchor. A cache names the faces it implements from +// face(); everything else it is asked for comes back null. +class ET_EXPERIMENTAL Cache { public: - virtual ~CacheBase() = default; - virtual SequenceControl* as_control() { - return nullptr; - } - virtual SequencePlanner* as_planner() { - return nullptr; - } - virtual BatchControl* as_batch_control() { - return nullptr; - } - virtual CellStepper* as_cell_stepper() { - return nullptr; + virtual ~Cache() = default; + + // Naming T::kFaceName means a type that is not a face fails to compile, + // rather than quietly returning null at run time. + template + T* as() { + return static_cast(face(T::kFaceName)); } + + protected: + // Implemented with expose<...>(this, id). Kept behind as() so callers do + // not handle erased pointers or face names directly. + virtual void* face(FaceId id) = 0; }; // Lifecycle and admission, tensor-free. -class CacheControl { +class ET_EXPERIMENTAL CacheControl { public: virtual ~CacheControl() = default; virtual bool can_extend(int n = 1) const = 0; // admission / hard-stop @@ -57,59 +84,21 @@ class CacheControl { }; // Application face of a single-sequence cache: one length to rewind. -class SequenceControl : public CacheControl { +class ET_EXPERIMENTAL SequenceControl : public CacheControl { public: + static constexpr const char* kFaceName = "et.cache.SequenceControl"; + // Truncate to new_len; false = cannot grow, or the target is older than an // evicting layer still retains. virtual bool rewind(int new_len) = 0; }; -// A contiguous span of physical rows in a layer's pool. -struct Run { - int start; - int len; -}; - -// Integer-only handoff to the backend byte layer. Runs are in logical order -// (oldest -> newest); a flat layer uses one, a ring layer two when it wraps. -// read_base_pos is the logical position of read[0].start. -struct SeqStepPlan { - Run write[2]; - int n_write; - Run read[2]; - int n_read; - int read_base_pos; -}; - -// Backend face. plan() is const: it computes a layer's layout without changing -// state, and commit() advances the shared logical length. nullopt = the step -// exceeds capacity, or `layer` is out of range. -class SequencePlanner { - public: - virtual ~SequencePlanner() = default; - virtual std::optional plan(int layer, int position, int T) - const = 0; - // Advance the logical length past this step. Idempotent, so once per step - // suffices. - virtual void commit(const SeqStepPlan& plan) = 0; -}; - -// Per-layer layout: flat keeps all history, ring slides a window. Stateless. -class LayoutPolicy { - public: - virtual ~LayoutPolicy() = default; - // Write/read runs for T cells at logical `position`. Precondition: T fits the - // policy's window. - virtual SeqStepPlan plan(int position, int T) const = 0; - // Oldest logical position still retained at this length: 0 for flat, - // length - window for ring. - virtual int retained_from(int length) const = 0; -}; - // Application face of any multi-sequence cache: the sequence verbs. They run // between forwards, never during one. -class BatchControl : public CacheControl { +class ET_EXPERIMENTAL BatchControl : public CacheControl { public: + static constexpr const char* kFaceName = "et.cache.BatchControl"; + // Which sequence each of the next forward's tokens belongs to, one entry per // token; every id must be one seq_new handed out. Also the admission gate: // false = rejected and nothing changed, and a step that passes has room for @@ -137,7 +126,7 @@ class BatchControl : public CacheControl { }; // Per-layer cache kind and its parameters. -struct LayerPolicy { +struct ET_EXPERIMENTAL LayerPolicy { enum class Kind : int { Flat = 0, Ring = 1 @@ -147,7 +136,7 @@ struct LayerPolicy { }; // Per-layer architecture facts + cache policy. -struct LayerConfig { +struct ET_EXPERIMENTAL LayerConfig { LayerPolicy policy; // default Flat int n_kv_heads; int head_dim; @@ -155,7 +144,7 @@ struct LayerConfig { // Model facts and the policy the byte layer sizes its pools from. `layers` is // per-layer: size 1 applies to every layer, else one entry each. -struct CacheConfig { +struct ET_EXPERIMENTAL CacheConfig { int capacity; // logical cap in cells int n_layers; std::vector layers; @@ -167,7 +156,7 @@ struct CacheConfig { }; // Whether `cfg` satisfies the contract above. -inline bool valid(const CacheConfig& cfg) { +ET_EXPERIMENTAL inline bool valid(const CacheConfig& cfg) { // initial_capacity may be 0 but not negative, and may exceed capacity -- the // byte layer clamps it. return cfg.capacity > 0 && cfg.n_layers > 0 && cfg.initial_capacity >= 0 && diff --git a/extension/llm/cache/cache_et.h b/extension/llm/cache/cache_et.h deleted file mode 100644 index 1c157f8ed60..00000000000 --- a/extension/llm/cache/cache_et.h +++ /dev/null @@ -1,64 +0,0 @@ -/* - * 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 - -// ExecuTorch adapter for the neutral cache core. The core (cache.h / -// sequence_cache.h) is ET-independent and reports failures as -// bool/std::optional so it is usable outside an ET runner. These thin inline -// adapters map those results to ExecuTorch Error/Result (logging on failure) -// for ET consumers -- the runner and the delegate byte layer. (The registry is -// delegate-specific and already returns Result directly, so it needs no -// adapter.) - -#include - -#include -#include -#include - -namespace executorch { -namespace extension { -namespace llm { -namespace cache { -namespace et { - -using ::executorch::runtime::Error; -using ::executorch::runtime::Result; - -// Plan a layer's step, or OutOfResources if it would exceed capacity (or the -// layer is out of range). -inline Result -plan(const SequencePlanner& planner, int layer, int position, int T) { - std::optional p = planner.plan(layer, position, T); - ET_CHECK_OR_RETURN_ERROR( - p.has_value(), - OutOfResources, - "cache: plan(layer=%d, position=%d, T=%d) exceeds capacity or bad layer", - layer, - position, - T); - return *p; -} - -// Truncate the history, or InvalidArgument if new_len would grow it (or is -// older than an evicting layer retains). -inline Error rewind(SequenceControl& control, int new_len) { - ET_CHECK_OR_RETURN_ERROR( - control.rewind(new_len), - InvalidArgument, - "rewind: cannot grow to %d", - new_len); - return Error::Ok; -} - -} // namespace et -} // namespace cache -} // namespace llm -} // namespace extension -} // namespace executorch diff --git a/extension/llm/cache/cache_registry.cpp b/extension/llm/cache/cache_registry.cpp index d54d05ddc73..51cb9cd711b 100644 --- a/extension/llm/cache/cache_registry.cpp +++ b/extension/llm/cache/cache_registry.cpp @@ -11,6 +11,7 @@ #include #include +#include namespace executorch { namespace extension { @@ -24,12 +25,12 @@ CacheRegistry& CacheRegistry::global() { void CacheRegistry::install( const std::string& key, - std::shared_ptr cache) { + std::shared_ptr cache) { std::lock_guard lock(mu_); caches_[key] = std::move(cache); } -std::shared_ptr CacheRegistry::get(const std::string& key) const { +std::shared_ptr CacheRegistry::get(const std::string& key) const { std::lock_guard lock(mu_); const auto it = caches_.find(key); return it == caches_.end() ? nullptr : it->second; @@ -40,20 +41,25 @@ void CacheRegistry::erase(const std::string& key) { caches_.erase(key); } -CacheBuilderRegistry& CacheBuilderRegistry::global() { - static CacheBuilderRegistry registry; +CacheFactory& CacheFactory::global() { + static CacheFactory registry; return registry; } -void CacheBuilderRegistry::register_builder( +Error CacheFactory::register_builder( const std::string& backend_id, const std::string& kind, CacheBuilder builder) { + if (!builder) { + return Error::InvalidArgument; + } std::lock_guard lock(mu_); - builders_[{backend_id, kind}] = std::move(builder); + const auto inserted = + builders_.emplace(std::make_pair(backend_id, kind), std::move(builder)); + return inserted.second ? Error::Ok : Error::InvalidArgument; } -Result> CacheBuilderRegistry::build( +Result> CacheFactory::build( const std::string& backend_id, const std::string& kind, const CacheConfig& cfg) const { @@ -61,12 +67,27 @@ Result> CacheBuilderRegistry::build( { std::lock_guard lock(mu_); const auto it = builders_.find({backend_id, kind}); - ET_CHECK_OR_RETURN_ERROR( - it != builders_.end(), - NotFound, - "no cache builder registered for %s:%s", - backend_id.c_str(), - kind.c_str()); + if (it == builders_.end()) { + // Name what is registered. A kind is a string, so a typo is otherwise a + // dead end. builders_ is ordered, so these come out sorted. + std::string known; + for (const auto& entry : builders_) { + if (entry.first.first != backend_id) { + continue; + } + if (!known.empty()) { + known += ", "; + } + known += entry.first.second; + } + ET_LOG( + Error, + "no '%s' cache registered for '%s'; registered: %s", + kind.c_str(), + backend_id.c_str(), + known.empty() ? "(none)" : known.c_str()); + return Error::NotFound; + } builder = it->second; } // Checked here rather than in each cache: `layers` is indexed directly, so a @@ -77,13 +98,36 @@ Result> CacheBuilderRegistry::build( "cache: invalid CacheConfig for %s:%s", backend_id.c_str(), kind.c_str()); - return builder(cfg); + // A builder that hands back null would otherwise travel as an ok() Result + // and be dereferenced by the caller. + auto cache = builder(cfg); + ET_CHECK_OR_RETURN_ERROR( + cache != nullptr, + Internal, + "cache: builder for %s:%s returned null", + backend_id.c_str(), + kind.c_str()); + return cache; } -std::string make_unique_key() { +namespace { +// Process-global atomic counter -> "cache-N". Internal: InstallGuard is the +// only thing that publishes, so nothing outside needs to mint a key. +std::string new_cache_key() { static std::atomic counter{0}; return "cache-" + std::to_string(counter.fetch_add(1)); } +} // namespace + +InstallGuard::InstallGuard(std::shared_ptr cache) + : key_(new_cache_key()), cache_(std::move(cache)) { + ET_CHECK_MSG(cache_ != nullptr, "Cannot install a null cache"); + CacheRegistry::global().install(key_, cache_); +} + +InstallGuard::~InstallGuard() { + CacheRegistry::global().erase(key_); +} } // namespace cache } // namespace llm diff --git a/extension/llm/cache/cache_registry.h b/extension/llm/cache/cache_registry.h index ae730344c8a..020a6fc2124 100644 --- a/extension/llm/cache/cache_registry.h +++ b/extension/llm/cache/cache_registry.h @@ -12,8 +12,8 @@ // is opaque to the host, so the runner (which knows the cache kind) creates the // cache and binds it to the delegate through a process-global registry; the two // sides rendezvous on a cache_key passed as a runtime backend-load option. -// Caches are owned as CacheBase* and the faces are recovered through its as_* -// accessors (no RTTI), each null for a face the cache does not implement. +// Caches are owned as Cache*; a face comes from as(), null when the cache +// does not implement it. #include #include @@ -24,8 +24,10 @@ #include #include +#include #include #include +#include namespace executorch { namespace extension { @@ -35,82 +37,100 @@ namespace cache { using ::executorch::runtime::Error; using ::executorch::runtime::Result; -// Process-global map>. Ownership is shared: -// the registry entry, the runner's session guard, and the delegate handle all -// hold the cache, so erasing the entry mid-method is safe. -class CacheRegistry { +// Backend-load option carrying the key of an installed cache. This name is the +// rendezvous contract shared by cache-owning runners and cache-aware backends. +inline constexpr char kCacheKeyOption[] = "llm_cache_registry_key"; + +// Process-global map>. Ownership is shared: +// the registry entry, the runner's guard, and the delegate handle all hold +// the cache, so erasing the entry mid-method is safe. +class ET_EXPERIMENTAL CacheRegistry { public: static CacheRegistry& global(); - void install(const std::string& key, std::shared_ptr cache); - std::shared_ptr get(const std::string& key) const; - void erase(const std::string& key); + // The delegate's half of the rendezvous: resolve a key it was handed as a + // backend option. Null if no cache is published under it. + std::shared_ptr get(const std::string& key) const; private: CacheRegistry() = default; + // Only InstallGuard may publish, so an entry cannot outlive its owner, two + // callers cannot collide on a key, and no erase can go unpaired. + friend class InstallGuard; + void install(const std::string& key, std::shared_ptr cache); + void erase(const std::string& key); + mutable std::mutex mu_; - std::unordered_map> caches_; + std::unordered_map> caches_; }; +// The registered cache kinds. Spelling one inline is a runtime NotFound rather +// than a compile error, so go through these. +namespace kind { +// One sequence over per-layer runs. +inline constexpr const char* kSingle = "single"; +// Many sequences sharing one pool of per-token cells. +inline constexpr const char* kBatchedCell = "batched-cell"; +} // namespace kind + // Cache kind is expressed by which factory you call: backends register a // builder per (backend_id, kind) and the kind survives only as an internal // lookup tag. -using CacheBuilder = - std::function(const CacheConfig&)>; +using CacheBuilder = std::function(const CacheConfig&)>; -class CacheBuilderRegistry { +class ET_EXPERIMENTAL CacheFactory { public: - static CacheBuilderRegistry& global(); + static CacheFactory& global(); + + // Public so a test can hold its own rather than registering builders into + // the process-global one, where they outlive it. + CacheFactory() = default; - void register_builder( + // Registers one builder without replacing an existing entry. Returns + // InvalidArgument if builder is empty or the pair is already registered. + ET_NODISCARD Error register_builder( const std::string& backend_id, const std::string& kind, CacheBuilder builder); - // Returns Error::NotFound if no builder is registered for (backend_id, kind). - Result> build( + // Returns NotFound if no builder is registered for (backend_id, kind), and + // Internal if the registered builder returns null. + Result> build( const std::string& backend_id, const std::string& kind, const CacheConfig& cfg) const; private: - CacheBuilderRegistry() = default; - mutable std::mutex mu_; - std::map, CacheBuilder> - builders_; // keyed by (backend_id, kind) + std::map, CacheBuilder> builders_; }; -// Process-global atomic counter -> "cache-N"; centralizes key generation so -// keys never collide. -std::string make_unique_key(); - -// RAII: installs the cache into the global registry under a unique key on -// construction and erases it on destruction (no leak on any exit path). Holds -// the runner's shared_ptr and exposes the control face for the generation loop. -class CacheSession { +// RAII over one registry entry: installs the cache under a key of its own +// making on construction and erases it on destruction (no leak on any exit +// path). Minting the key here rather than taking one means two live guards +// cannot collide on it. Must outlive the load_method() whose backend init +// resolves the key. +// +// Destruction removes discoverability only. A shared_ptr already returned by +// CacheRegistry::get() remains valid independently. +class ET_EXPERIMENTAL InstallGuard { public: - CacheSession(std::string key, std::shared_ptr cache) - : key_(std::move(key)), cache_(std::move(cache)) { - CacheRegistry::global().install(key_, cache_); - } - ~CacheSession() { - CacheRegistry::global().erase(key_); - } + explicit InstallGuard(std::shared_ptr cache); + ~InstallGuard(); - CacheSession(const CacheSession&) = delete; - CacheSession& operator=(const CacheSession&) = delete; + InstallGuard(const InstallGuard&) = delete; + InstallGuard& operator=(const InstallGuard&) = delete; - SequenceControl* control() const { - return cache_->as_control(); - } - const std::string& key() const { - return key_; + // Adds the complete cache rendezvous option. BackendOptions copies the key + // and value, so the resulting option remains valid independently. + template + Error set_option(::executorch::runtime::BackendOptions& options) const { + return options.set_option(kCacheKeyOption, key_.c_str()); } private: std::string key_; - std::shared_ptr cache_; + std::shared_ptr cache_; }; } // namespace cache diff --git a/extension/llm/cache/cell_cache.cpp b/extension/llm/cache/cell_cache.cpp index ceaa31a6686..43866357baf 100644 --- a/extension/llm/cache/cell_cache.cpp +++ b/extension/llm/cache/cell_cache.cpp @@ -19,8 +19,7 @@ namespace cache { CellCache::CellCache(const CacheConfig& cfg) : capacity_(cfg.capacity), pos_(cfg.capacity, -1), - owners_(cfg.capacity, 0), - served_(cfg.n_layers, false) { + owners_(cfg.capacity, 0) { assert(valid(cfg)); // One window per layer, from the same per-layer config the sequence cache // reads. Layers agreeing on a window share a step. @@ -53,7 +52,6 @@ void CellCache::clear() { declared_ = false; step_seq_ids_.clear(); step_pos_.clear(); - std::fill(served_.begin(), served_.end(), false); invalidate_steps(); } @@ -71,7 +69,6 @@ bool CellCache::declare_step(const std::vector& seq_ids) { step_seq_ids_ = seq_ids; declared_ = true; invalidate_steps(); - std::fill(served_.begin(), served_.end(), false); return true; } @@ -155,25 +152,32 @@ int CellCache::used_end() const { const CellStep* CellCache::place_step(int layer, const int32_t* positions, int length) { - if (layer < 0 || layer >= static_cast(windows_.size()) || - served_[layer]) { - return nullptr; // out of range, or a forward that skipped declare_step + if (layer < 0 || layer >= static_cast(windows_.size())) { + return nullptr; // layer out of range } - if (!placed_) { - if (!declared_ || length != static_cast(step_seq_ids_.size())) { - return nullptr; // no declaration, or a token count disagreeing with it - } - if (!extends(positions, length)) { - return nullptr; // nothing mutated yet, so the step can be re-placed - } - step_pos_.assign(positions, positions + length); - if (!place()) { + if (placed_) { + // Re-serve within the placed forward. Every layer of a forward places the + // same tokens, and a KV-shared layer re-serves its donor's id, so a repeat + // with the same positions returns the same step and claims no new cells. + // Different positions mean a new step that never declared, still refused. + if (length != static_cast(step_pos_.size()) || + !std::equal(positions, positions + length, step_pos_.begin())) { return nullptr; } - declared_ = false; // one declaration, one placement - placed_ = true; + return &step_for(windows_[layer]); + } + if (!declared_ || length != static_cast(step_seq_ids_.size())) { + return nullptr; // no declaration, or a token count disagreeing with it + } + if (!extends(positions, length)) { + return nullptr; // a position a sequence already holds + } + step_pos_.assign(positions, positions + length); + if (!place()) { + return nullptr; // out of cells } - served_[layer] = true; + declared_ = false; // one declaration, one placement + placed_ = true; return &step_for(windows_[layer]); } diff --git a/extension/llm/cache/cell_cache.h b/extension/llm/cache/cell_cache.h index 2c519d9d28c..256a43739c5 100644 --- a/extension/llm/cache/cell_cache.h +++ b/extension/llm/cache/cell_cache.h @@ -29,7 +29,7 @@ namespace cache { // Integer-only handoff to the byte layer, covering the whole forward: a cell // means the same token in every layer's pool. -struct CellStep { +struct ET_EXPERIMENTAL CellStep { int length; int read_len; // the window is cells [0, read_len) std::vector cells; // cell per query token @@ -40,34 +40,32 @@ struct CellStep { // every later layer reuses that placement. `layer` selects the window, which // decides the kind and mask, so a step is per policy and memoized for the // forward. The returned step is owned by the cache and valid until the next -// verb. nullptr = no declaration, a token count disagreeing with it, a position -// a sequence already holds, a layer out of range, or a layer served twice. -class CellStepper { +// verb. A layer may be served more than once per forward -- a KV-shared layer +// re-serves its donor's id -- provided the repeat passes the same positions; it +// returns the same step and claims no new cells. nullptr = no declaration, a +// token count disagreeing with it, a position a sequence already holds, a layer +// out of range, or a re-serve whose positions differ (a step that never +// declared). +class ET_EXPERIMENTAL CellStepper { public: + static constexpr const char* kFaceName = "et.cache.CellStepper"; + virtual ~CellStepper() = default; virtual const CellStep* place_step(int layer, const int32_t* positions, int length) = 0; }; -class CellCache : public CacheBase, public BatchControl, public CellStepper { +class ET_EXPERIMENTAL CellCache : public Cache, + public BatchControl, + public CellStepper { public: // One bit per sequence in the owner bitset. static constexpr int kMaxSeqs = 64; - // Precondition: valid(cfg). CacheBuilderRegistry::build enforces it for + // Precondition: valid(cfg). CacheFactory::build enforces it for // registry-created caches; direct construction must check first. explicit CellCache(const CacheConfig& cfg); - CacheBase* base() { - return this; - } - BatchControl* as_batch_control() override { - return this; - } - CellStepper* as_cell_stepper() override { - return this; - } - // -- CacheControl ------------------------------------------------------ bool can_extend(int n = 1) const override; @@ -92,6 +90,11 @@ class CellCache : public CacheBase, public BatchControl, public CellStepper { const CellStep* place_step(int layer, const int32_t* positions, int length) override; + protected: + void* face(FaceId id) override { + return expose(this, id); + } + private: struct SeqInfo { int count = 0; @@ -149,7 +152,6 @@ class CellCache : public CacheBase, public BatchControl, public CellStepper { std::vector step_seq_ids_; // set by declare_step std::vector step_pos_; // set when the step is placed std::vector cells_; // the step's placement, shared by every layer - std::vector served_; // layers this step has already answered std::vector windows_; // per layer; 0 = keeps all history // window -> step, memoized per forward. Node-based is required: a step // handed to one layer must survive another layer's insert. diff --git a/extension/llm/cache/reference_cache.py b/extension/llm/cache/reference_cache.py index 7abc5d5138a..c0e206811f5 100644 --- a/extension/llm/cache/reference_cache.py +++ b/extension/llm/cache/reference_cache.py @@ -20,9 +20,10 @@ The cache places K/V and returns the history plus an ``AttendSpec`` (a mask *semantic*). The attend mechanism (``attend`` below) is applied by the op/backend from that spec. -Two caches share the op: ``ContiguousReferenceCache`` (one sequence appended in -place) and ``CellReferenceCache`` (many sequences over a pool of per-token cells, -with sharing and eviction). Both store float KV. +Three caches share the op: ``SequenceReferenceCache`` (one sequence), +``BatchedSequenceReferenceCache`` (many private sequence caches), and +``CellReferenceCache`` (many sequences over a shared pool of per-token cells, +with sharing and eviction). All store float KV. """ from __future__ import annotations @@ -50,7 +51,17 @@ class MaskKind(Enum): @dataclass class AttendSpec: + """One attention: what to attend over, which queries do it, how to mask it. + + A step is answered with a list of these -- one for a cache holding a single + history, one per sequence for a cache holding a private history each. They + cover the query axis in order, so ``q_len`` alone places each. + """ + + k: torch.Tensor # [B, H_kv, total, head_dim] -- key history + v: torch.Tensor # [B, H_kv, total, v_head_dim] -- value history kind: MaskKind + q_len: int # query tokens this spec answers, following the one before it mask: Optional[torch.Tensor] = None # EXPLICIT only: bool, true = attend @@ -109,7 +120,7 @@ def policy_for(self, layer_id: int) -> LayerPolicy: @experimental( "update_and_attend KV cache is experimental and may change without notice." ) -class ContiguousReferenceCache: +class SequenceReferenceCache: """Per-layer contiguous float KV history for a single sequence.""" def __init__(self, config: CacheConfig): @@ -126,6 +137,36 @@ def __init__(self, config: CacheConfig): def used(self, layer_id: int) -> int: return self._used[layer_id] + def rewind(self, new_len: int) -> None: + """Drop everything from ``new_len`` on, in every layer. + + A windowed layer retains only its last ``window`` positions, so it + cannot go back further than that even though this reference keeps the + older ones -- the window is applied to the mask here and to the storage + in a byte layer, and a rewind past it would attend cells that layer no + longer holds. + """ + used = self._used[0] + if new_len < 0 or new_len > used: + raise ValueError(f"rewind to {new_len}: the history holds {used}") + floor = max( + ( + used - self.config.policy_for(layer_id).window + for layer_id in range(self.config.n_layers) + if self.config.policy_for(layer_id).window > 0 + ), + default=0, + ) + if new_len < floor: + raise ValueError( + f"rewind to {new_len}: a windowed layer retains only from {floor}" + ) + for layer_id in range(self.config.n_layers): + if self.config.sizing == CacheSizing.DYNAMIC: + self._k[layer_id] = self._k[layer_id][:, :, :new_len, :] + self._v[layer_id] = self._v[layer_id][:, :, :new_len, :] + self._used[layer_id] = new_len + def reset(self): self._used = [0] * self.config.n_layers if self.config.sizing == CacheSizing.DYNAMIC: @@ -144,8 +185,10 @@ def update_and_fetch( k: torch.Tensor, v: torch.Tensor, position: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, AttendSpec]: - """Place this step's K/V and return the full history + mask semantic. + ) -> List[AttendSpec]: + """Place this step's K/V and return what to attend over. + + One sequence, one history, so the list is always one long. Per the design, ``position`` is the cache's placement + masking input. This contiguous single-sequence cache appends at its used length, so the @@ -160,9 +203,8 @@ def update_and_fetch( position: ``[q_len, n_dims]`` int -- per-query-token positions. Returns: - ``(k_hist, v_hist, spec)`` -- history ``[B, H_kv, total, head_dim]`` / - ``[B, H_kv, total, v_head_dim]`` (``total`` = prior length + q_len) and - the AttendSpec mask semantic. + one AttendSpec over the whole history, ``total`` = prior length + + q_len. """ q_len = k.shape[-2] used = self._used[layer_id] @@ -188,10 +230,14 @@ def update_and_fetch( v_hist = self._v[layer_id] self._used[layer_id] = new_used - return k_hist, v_hist, self._spec(layer_id, q_len, new_used, k.device) + return [self._spec(layer_id, k_hist, v_hist, q_len)] def _spec( - self, layer_id: int, q_len: int, total: int, device: torch.device + self, + layer_id: int, + k_hist: torch.Tensor, + v_hist: torch.Tensor, + q_len: int, ) -> AttendSpec: """The mask semantic for q_len new cells at the tail of a total window. @@ -200,21 +246,25 @@ def _spec( ``i + total - q_len - window``. Whichever bound the fused kinds cannot express is what makes the step EXPLICIT. """ + total = k_hist.shape[-2] window = self.config.policy_for(layer_id).window windowed = 0 < window < total if q_len == 1 and not windowed: - return AttendSpec(kind=MaskKind.NONE) + return AttendSpec(k=k_hist, v=v_hist, kind=MaskKind.NONE, q_len=q_len) if q_len == total and not windowed: - return AttendSpec(kind=MaskKind.CAUSAL) + return AttendSpec(k=k_hist, v=v_hist, kind=MaskKind.CAUSAL, q_len=q_len) # torch's is_causal is upper-left and expresses no window, so the band # is handed back explicitly. + device = k_hist.device offsets = torch.arange(total, device=device) - torch.arange( q_len, device=device ).unsqueeze(-1) band = offsets <= total - q_len if windowed: band &= offsets > total - q_len - window - return AttendSpec(kind=MaskKind.EXPLICIT, mask=band) + return AttendSpec( + k=k_hist, v=v_hist, kind=MaskKind.EXPLICIT, q_len=q_len, mask=band + ) # A cell's owners are a bitset in a torch int64, so bit 63 (the sign bit) is out. @@ -242,7 +292,7 @@ def flatten_step( Returns: ``(tokens, positions, seq_ids, logits_indices)`` -- tokens concatenated on the token axis and ``positions`` (``[n_tok, 1]``) as model inputs, - ``seq_ids`` for ``begin_step``, and ``logits_indices`` selecting each + ``seq_ids`` for ``declare_step``, and ``logits_indices`` selecting each sequence's last token, the rows worth running the LM head on. """ tokens, positions, seq_ids, logits_indices = [], [], [], [] @@ -259,6 +309,180 @@ def flatten_step( ) +@dataclass(frozen=True) +class _SequenceSpan: + seq_id: int + start: int + length: int + + +@experimental( + "update_and_attend KV cache is experimental and may change without notice." +) +class BatchedSequenceReferenceCache: + """A private ``SequenceReferenceCache`` per sequence in a flat batch. + + Projections share one model forward over the flattened token axis. Attention + splits that axis into its declared sequence spans, runs independently over + each sequence's private history, then concatenates the outputs in input + order. No sequence attends another and no dense cross-sequence mask is built. + """ + + def __init__(self, config: CacheConfig): + if config.batch_size != 1: + raise ValueError( + "batched sequence cache is flat on the token axis: batch_size must be 1" + ) + self.config = config + self._sequences: Dict[int, SequenceReferenceCache] = {} + self._spans: List[_SequenceSpan] = [] + self._served: Set[int] = set() + self._declared = False + + def declare_step(self, seq_ids: Sequence[int]) -> None: + if not seq_ids: + raise ValueError("a step carries at least one token") + for seq_id in seq_ids: + self._check_seq_id(seq_id) + + spans: List[_SequenceSpan] = [] + start = 0 + while start < len(seq_ids): + seq_id = seq_ids[start] + end = start + 1 + while end < len(seq_ids) and seq_ids[end] == seq_id: + end += 1 + spans.append(_SequenceSpan(seq_id, start, end - start)) + start = end + + # capacity bounds the whole cache. Checked before anything is created so a refusal changes + # nothing. + held = sum(sequence.used(0) for sequence in self._sequences.values()) + if held + len(seq_ids) > self.config.capacity: + raise RuntimeError( + f"KV cache overflow: {held + len(seq_ids)} cells exceeds " + f"capacity {self.config.capacity}" + ) + + for span in spans: + if span.seq_id not in self._sequences: + self._sequences[span.seq_id] = SequenceReferenceCache(self.config) + + self._spans = spans + self._served.clear() + self._declared = True + + def update_and_fetch( + self, + layer_id: int, + k: torch.Tensor, + v: torch.Tensor, + position: torch.Tensor, + ) -> List[AttendSpec]: + """Place each span's K/V in its own sequence and return one spec each. + + The specs follow the declared spans, so they cover the query axis in + order and no sequence appears in another's window. + """ + if not self._declared: + raise RuntimeError( + "no step declared: declare_step must precede every forward" + ) + if layer_id in self._served: + raise RuntimeError( + f"layer {layer_id} served twice for one step: " + "declare_step must precede every forward" + ) + token_count = k.shape[-2] + if not position.shape[0] == token_count == v.shape[-2]: + raise ValueError("position, k, and v must have the same token count") + if token_count != sum(span.length for span in self._spans): + raise ValueError("the forward token count must match declare_step") + if position.shape[-1] != 1: + raise NotImplementedError( + "sequence placement needs one position per token, got " + f"{position.shape[-1]}" + ) + self._check_positions(layer_id, position.reshape(-1).tolist()) + + specs: List[AttendSpec] = [] + for span in self._spans: + end = span.start + span.length + sequence = self._sequences[span.seq_id] + specs.extend( + sequence.update_and_fetch( + layer_id, + k[:, :, span.start : end, :], + v[:, :, span.start : end, :], + position[span.start : end], + ) + ) + + self._served.add(layer_id) + return specs + + def reset(self) -> None: + self._sequences.clear() + self._spans.clear() + self._served.clear() + self._declared = False + + def seq_rm(self, seq_id: int, p0: int = 0, p1: Optional[int] = None) -> None: + """Drop seq_id's claim on positions [p0, p1); p1 = None runs to the end. + + A private contiguous history drops its tail but not its middle: it has + no per-token position map to reindex what would survive. Bounding + history from below is a layer policy, not a verb. + """ + self._check_seq_id(seq_id) + if p1 is not None: + raise NotImplementedError( + "a private history drops only its tail, so [p0, p1) with a " + "bounded end has nothing to reindex the remainder against" + ) + sequence = self._sequences.get(seq_id) + if sequence is not None: + if p0 == 0: + del self._sequences[seq_id] + else: + sequence.rewind(p0) + self._spans.clear() + self._served.clear() + self._declared = False + + def seq_len(self, seq_id: int) -> int: + self._check_seq_id(seq_id) + sequence = self._sequences.get(seq_id) + return sequence.used(0) if sequence is not None else 0 + + def _check_positions(self, layer_id: int, positions: List[int]) -> None: + """Every span continues its own sequence, from where that sequence ends. + + A private history appends at its used length and never reads + ``position``, so a step that declared the wrong one would still place + its tokens contiguously -- correct cells under the wrong names, and no + later step would notice. A sequence spanned twice in one step continues + across both. + """ + ends: Dict[int, int] = {} + for span in self._spans: + at = ends.get(span.seq_id, self._sequences[span.seq_id].used(layer_id)) + got = positions[span.start : span.start + span.length] + want = list(range(at, at + span.length)) + if got != want: + raise ValueError( + f"sequence {span.seq_id} holds {at} positions on layer " + f"{layer_id}: the step declares {got}, not {want}" + ) + ends[span.seq_id] = at + span.length + + @staticmethod + def _check_seq_id(seq_id: int) -> None: + # No upper bound: a sequence is a dict entry, not a bit in an owner set. + if seq_id < 0: + raise ValueError(f"seq_id must be non-negative, got {seq_id}") + + @dataclass class _CellStepPlan: """One step's allocation, shared by every layer of that forward. @@ -296,7 +520,7 @@ class CellReferenceCache: that, so the spec is always EXPLICIT. The batch is flat: tokens from every sequence sit on one axis with B = 1, - and sequence identity is supplied out-of-band. ``begin_step`` declares which + and sequence identity is supplied out-of-band. ``declare_step`` declares which sequence each of the next forward's tokens belongs to; the positions arrive with the forward itself, in the op's ``position`` tensor, so cells are allocated on the first layer of the step and memoized for the rest of it. @@ -332,7 +556,7 @@ def __init__(self, config: CacheConfig): for _ in range(config.n_layers) ] self._step_seq_ids: List[int] = [] - self._declared = False # set by begin_step, cleared by the step it authorizes + self._declared = False # set by declare_step, cleared by the step it authorizes self._plan: Optional[_CellStepPlan] = None self._served: Set[int] = set() @@ -354,7 +578,7 @@ def seq_len(self, seq_id: int) -> int: bit = 1 << seq_id return sum(1 for owners in self._owners if owners & bit) - def begin_step(self, seq_ids: Sequence[int]) -> None: + def declare_step(self, seq_ids: Sequence[int]) -> None: """Declare the sequence each of the next forward's tokens belongs to. Admission is decided here, before the forward: the token count is known @@ -430,17 +654,19 @@ def update_and_fetch( k: torch.Tensor, v: torch.Tensor, position: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, AttendSpec]: + ) -> List[AttendSpec]: """Scatter this step's K/V into its cells and return the read window. The first layer of a step allocates; the rest reuse that allocation, so the cells and the mask are computed once per forward, not once per - layer. Args are as ``ContiguousReferenceCache.update_and_fetch``. + layer. Every sequence reads the same window and the mask holds them + apart, so the list is always one long. Args are as + ``SequenceReferenceCache.update_and_fetch``. """ if layer_id in self._served: raise RuntimeError( f"layer {layer_id} served twice for one step: " - "begin_step must precede every forward" + "declare_step must precede every forward" ) if self._plan is None: self._plan = self._allocate(position) @@ -451,14 +677,15 @@ def update_and_fetch( cells = self._plan.cells self._k[layer_id][:, :, cells, :] = k.to(self.config.dtype) self._v[layer_id][:, :, cells, :] = v.to(self.config.dtype) - return ( - self._k[layer_id][:, :, :read_len, :], - self._v[layer_id][:, :, :read_len, :], + return [ AttendSpec( + k=self._k[layer_id][:, :, :read_len, :], + v=self._v[layer_id][:, :, :read_len, :], kind=MaskKind.EXPLICIT, + q_len=len(cells), mask=self._plan.mask_for(self.config.policy_for(layer_id).window), - ), - ) + ) + ] # -- internals ---------------------------------------------------------- @@ -467,18 +694,17 @@ def _allocate(self, position: torch.Tensor) -> _CellStepPlan: device = self._k[0].device if not self._declared: raise RuntimeError( - "no step declared: begin_step must precede every forward" + "no step declared: declare_step must precede every forward" ) self._declared = False # one declaration, one attempt at allocating it if position.shape[-1] != 1: raise NotImplementedError( - "cell placement needs one position per token, got " - f"{position.shape[-1]}" + f"cell placement needs one position per token, got {position.shape[-1]}" ) positions = position.reshape(-1).tolist() if len(positions) != len(self._step_seq_ids): raise ValueError( - f"begin_step declared {len(self._step_seq_ids)} tokens, " + f"declare_step declared {len(self._step_seq_ids)} tokens, " f"the forward carries {len(positions)}" ) cells = [ @@ -537,7 +763,7 @@ def _claim(self, pos: int, owners: int) -> int: self._owners[i] = owners self._used_end = max(self._used_end, i + 1) return i - raise RuntimeError("no free cell") # begin_step admitted the step + raise RuntimeError("no free cell") # declare_step admitted the step def _shrink(self): while self._used_end > 0 and self._pos[self._used_end - 1] < 0: @@ -546,7 +772,7 @@ def _shrink(self): def _invalidate_plan(self): # A mutated cell table leaves a built plan's cells and mask stale. The # step protocol state is deliberately left alone: a mutation must not - # disguise a forward that skipped begin_step. + # disguise a forward that skipped declare_step. self._plan = None @staticmethod @@ -563,8 +789,6 @@ def _in_range(pos: int, p0: int, p1: Optional[int]) -> bool: def attend( q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, spec: AttendSpec, scale: float, out_dtype: torch.dtype, @@ -579,17 +803,17 @@ def attend( so a cache must declare EXPLICIT for a chunked or multi-turn step. Args (BHSD): - q: ``[B, H_q, q_len, head_dim]`` -- queries (already RoPE-rotated). - k: ``[B, H_kv, total, head_dim]`` -- key history. - v: ``[B, H_kv, total, v_head_dim]`` -- value history. - spec: mask semantic (NONE = attend all; CAUSAL = causal; EXPLICIT = the - spec's bool mask). + q: ``[B, H_q, q_len, head_dim]`` -- queries (already RoPE-rotated), the + ones this spec answers. + spec: the K/V history to attend over and its mask semantic (NONE = + attend all; CAUSAL = causal; EXPLICIT = the spec's bool mask). scale: attention softmax scale. out_dtype: output dtype. Returns: ``[B, H_q, q_len, v_head_dim]`` attention output, in ``out_dtype``. """ + k, v = spec.k, spec.v n_q_heads = q.shape[1] n_kv_heads = k.shape[1] if n_q_heads != n_kv_heads: diff --git a/extension/llm/cache/sequence_cache.h b/extension/llm/cache/sequence_cache.h index 5af756e6e24..75aa02f8987 100644 --- a/extension/llm/cache/sequence_cache.h +++ b/extension/llm/cache/sequence_cache.h @@ -12,6 +12,11 @@ // policies (FlatPolicy / RingPolicy). SequenceCache owns the one logical length // for the whole model and dispatches per-layer layout to a policy, so a mixed // flat/ring model (gemma4) stays coherent. Tensor-free / ET-independent. +// +// The backend-facing planner face (SequencePlanner) and the types it hands over +// live here rather than in cache.h: only this layout implements them, and only +// a byte layer that already includes this header calls them. cell_cache.h holds +// CellStepper for the same reason. #include #include @@ -26,8 +31,52 @@ namespace extension { namespace llm { namespace cache { +// A contiguous span of physical rows in a layer's pool. +struct ET_EXPERIMENTAL Run { + int start; + int len; +}; + +// Integer-only handoff to the backend byte layer. Runs are in logical order +// (oldest -> newest); a flat layer uses one, a ring layer two when it wraps. +// read_base_pos is the logical position of read[0].start. +struct ET_EXPERIMENTAL SeqStepPlan { + Run write[2]; + int n_write; + Run read[2]; + int n_read; + int read_base_pos; +}; + +// Backend face. plan() is const: it computes a layer's layout without changing +// state, and commit() advances the shared logical length. nullopt = the step +// exceeds capacity, or `layer` is out of range. +class ET_EXPERIMENTAL SequencePlanner { + public: + static constexpr const char* kFaceName = "et.cache.SequencePlanner"; + + virtual ~SequencePlanner() = default; + virtual std::optional plan(int layer, int position, int T) + const = 0; + // Advance the logical length past this step. Idempotent, so once per step + // suffices. + virtual void commit(const SeqStepPlan& plan) = 0; +}; + +// Per-layer layout: flat keeps all history, ring slides a window. Stateless. +class ET_EXPERIMENTAL LayoutPolicy { + public: + virtual ~LayoutPolicy() = default; + // Write/read runs for T cells at logical `position`. Precondition: T fits the + // policy's window. + virtual SeqStepPlan plan(int position, int T) const = 0; + // Oldest logical position still retained at this length: 0 for flat, + // length - window for ring. + virtual int retained_from(int length) const = 0; +}; + // Full history [0, length): one contiguous write run, read over all history. -class FlatPolicy final : public LayoutPolicy { +class ET_EXPERIMENTAL FlatPolicy final : public LayoutPolicy { public: int retained_from(int /*length*/) const override { return 0; // keeps all history @@ -48,7 +97,7 @@ class FlatPolicy final : public LayoutPolicy { // slots. The ring is oversized so a step of up to max_write tokens fits without // overwriting cells earlier queries in the same step still attend to; the // backend masks each query to its own window within the read span. -class RingPolicy final : public LayoutPolicy { +class ET_EXPERIMENTAL RingPolicy final : public LayoutPolicy { public: RingPolicy(int window, int max_write) : window_(window), ring_size_(window + max_write - 1) {} @@ -92,9 +141,9 @@ class RingPolicy final : public LayoutPolicy { // rewind; dispatches per-layer layout to a shared LayoutPolicy. Policies are // deduped by (kind, window), so a uniform or two-kind (gemma4) model holds one // or two policy objects. -class SequenceCache : public CacheBase, - public SequenceControl, - public SequencePlanner { +class ET_EXPERIMENTAL SequenceCache : public Cache, + public SequenceControl, + public SequencePlanner { public: explicit SequenceCache(const CacheConfig& cfg) : capacity_(cfg.capacity), max_write_(cfg.max_write) { @@ -108,14 +157,6 @@ class SequenceCache : public CacheBase, } } - // CacheBase: face recovery without RTTI. - SequenceControl* as_control() override { - return this; - } - SequencePlanner* as_planner() override { - return this; - } - // SequenceControl. bool can_extend(int n = 1) const override { return length_ + n <= @@ -173,6 +214,11 @@ class SequenceCache : public CacheBase, length_ = std::max(length_, end); } + protected: + void* face(FaceId id) override { + return expose(this, id); + } + private: int policy_index(const LayerPolicy& lp) { for (std::size_t i = 0; i < specs_.size(); ++i) { diff --git a/extension/llm/cache/test/cache_test.cpp b/extension/llm/cache/test/cache_test.cpp index e41e91a6089..654db1a5b72 100644 --- a/extension/llm/cache/test/cache_test.cpp +++ b/extension/llm/cache/test/cache_test.cpp @@ -7,7 +7,6 @@ */ #include -#include #include #include #include @@ -19,31 +18,53 @@ #include #include +#include #include using executorch::extension::llm::cache::BatchControl; -using executorch::extension::llm::cache::CacheBase; -using executorch::extension::llm::cache::CacheBuilderRegistry; +using executorch::extension::llm::cache::Cache; +using executorch::extension::llm::cache::CacheBuilder; using executorch::extension::llm::cache::CacheConfig; +using executorch::extension::llm::cache::CacheFactory; using executorch::extension::llm::cache::CacheRegistry; -using executorch::extension::llm::cache::CacheSession; using executorch::extension::llm::cache::CellCache; using executorch::extension::llm::cache::CellStep; using executorch::extension::llm::cache::CellStepper; +using executorch::extension::llm::cache::InstallGuard; +using executorch::extension::llm::cache::SequenceControl; +using executorch::extension::llm::cache::SequencePlanner; +namespace kind = executorch::extension::llm::cache::kind; using executorch::extension::llm::cache::LayerConfig; using executorch::extension::llm::cache::LayerPolicy; -using executorch::extension::llm::cache::make_unique_key; using executorch::extension::llm::cache::SequenceCache; +using executorch::runtime::BackendOptions; using executorch::runtime::Error; -namespace et = executorch::extension::llm::cache::et; namespace { +std::string installed_key(const InstallGuard& guard) { + BackendOptions<1> options; + if (guard.set_option(options) != Error::Ok) { + return {}; + } + const char* key = nullptr; + if (options.get_option( + executorch::extension::llm::cache::kCacheKeyOption, key) != + Error::Ok) { + return {}; + } + return key; +} + LayerConfig flat_layer() { return LayerConfig{LayerPolicy{LayerPolicy::Kind::Flat, 0}, 2, 8}; } LayerConfig ring_layer(int window) { return LayerConfig{LayerPolicy{LayerPolicy::Kind::Ring, window}, 2, 8}; } + +struct UnsupportedFace { + static constexpr const char* kFaceName = "test.UnsupportedFace"; +}; } // namespace // Initializes the ExecuTorch PAL so the ET adapter's error paths (which ET_LOG) @@ -167,92 +188,169 @@ TEST_F(CacheTest, RewindBoundedByRingWindow) { EXPECT_FALSE(cache.rewind(11)); // cannot grow } -// ---- Faces / registry / session -------------------------------------------- +// ---- Faces / registry / lease ---------------------------------------------- TEST_F(CacheTest, FaceRecoveryReturnsSameObject) { SequenceCache cache(CacheConfig{4, 1, {flat_layer()}}); - CacheBase* base = &cache; - ASSERT_NE(base->as_control(), nullptr); - ASSERT_NE(base->as_planner(), nullptr); - EXPECT_TRUE(base->as_control()->can_extend(4)); - auto plan = base->as_planner()->plan(0, 0, 1); + Cache* base = &cache; + ASSERT_NE(base->as(), nullptr); + ASSERT_NE(base->as(), nullptr); + EXPECT_EQ(base->as(), nullptr); + EXPECT_TRUE(base->as()->can_extend(4)); + auto plan = base->as()->plan(0, 0, 1); ASSERT_TRUE(plan.has_value()); EXPECT_EQ(plan->read[0].len, 1); } -TEST_F(CacheTest, RegistryInstallGetErase) { - auto& reg = CacheRegistry::global(); - const std::string key = make_unique_key(); - EXPECT_EQ(reg.get(key), nullptr); - - std::shared_ptr cache = - std::make_shared(CacheConfig{16, 1, {flat_layer()}}); - reg.install(key, cache); - EXPECT_EQ(reg.get(key), cache); - EXPECT_TRUE(reg.get(key)->as_control()->can_extend(16)); - - reg.erase(key); - EXPECT_EQ(reg.get(key), nullptr); +TEST_F(CacheTest, NullFaceIdsNeverMatch) { + using executorch::extension::llm::cache::same_face; + EXPECT_FALSE(same_face(nullptr, nullptr)); + EXPECT_FALSE(same_face(nullptr, SequenceControl::kFaceName)); + EXPECT_FALSE(same_face(SequenceControl::kFaceName, nullptr)); } -TEST_F(CacheTest, UniqueKeysDoNotCollide) { - EXPECT_NE(make_unique_key(), make_unique_key()); +TEST_F(CacheTest, LiveGuardsDoNotCollideOnKeys) { + auto a = std::make_shared(CacheConfig{4, 1, {flat_layer()}}); + auto b = std::make_shared(CacheConfig{4, 1, {flat_layer()}}); + InstallGuard ga(a); + InstallGuard gb(b); + + const std::string a_key = installed_key(ga); + const std::string b_key = installed_key(gb); + ASSERT_FALSE(a_key.empty()); + ASSERT_FALSE(b_key.empty()); + EXPECT_NE(a_key, b_key); + // Both entries are live: neither guard displaced the other. + EXPECT_EQ(CacheRegistry::global().get(a_key).get(), a.get()); + EXPECT_EQ(CacheRegistry::global().get(b_key).get(), b.get()); } TEST_F(CacheTest, BuilderBuildsRegisteredKindElseError) { - auto& reg = CacheBuilderRegistry::global(); - reg.register_builder("TestBackend", "seq", [](const CacheConfig& cfg) { - return std::static_pointer_cast( - std::make_shared(cfg)); - }); + // Its own factory: a builder registered into the global one would + // outlive this test and be visible to every case after it. + CacheFactory reg; + EXPECT_EQ( + reg.register_builder( + "TestBackend", + kind::kSingle, + [](const CacheConfig& cfg) { + return std::static_pointer_cast( + std::make_shared(cfg)); + }), + Error::Ok); CacheConfig cfg{32, 1, {flat_layer()}}; - auto cache = reg.build("TestBackend", "seq", cfg); + auto cache = reg.build("TestBackend", kind::kSingle, cfg); ASSERT_TRUE(cache.ok()); - EXPECT_EQ(cache.get()->as_control()->capacity(), 32); + EXPECT_EQ(cache.get()->as()->capacity(), 32); EXPECT_EQ(reg.build("TestBackend", "missing", cfg).error(), Error::NotFound); // A layers list that is neither size 1 nor n_layers would be indexed past // the end, so build refuses it before the cache is constructed. EXPECT_EQ( - reg.build("TestBackend", "seq", CacheConfig{32, 3, {}}).error(), + reg.build("TestBackend", kind::kSingle, CacheConfig{32, 3, {}}).error(), Error::InvalidArgument); EXPECT_EQ( reg.build( "TestBackend", - "seq", + kind::kSingle, CacheConfig{32, 3, {flat_layer(), flat_layer()}}) .error(), Error::InvalidArgument); } -TEST_F(CacheTest, SessionInstallsOnCtorErasesOnDtor) { - const std::string key = make_unique_key(); +TEST_F(CacheTest, BuilderRegistrationRejectsInvalidEntries) { + CacheFactory factory; + CacheConfig cfg{32, 1, {flat_layer()}}; + + EXPECT_EQ( + factory.register_builder("TestBackend", "empty", CacheBuilder{}), + Error::InvalidArgument); + EXPECT_EQ( + factory.build("TestBackend", "empty", cfg).error(), Error::NotFound); + + EXPECT_EQ( + factory.register_builder( + "TestBackend", + "duplicate", + [](const CacheConfig& config) { + return std::static_pointer_cast( + std::make_shared(config)); + }), + Error::Ok); + EXPECT_EQ( + factory.register_builder( + "TestBackend", + "duplicate", + [](const CacheConfig&) { return std::shared_ptr{}; }), + Error::InvalidArgument); + auto original = factory.build("TestBackend", "duplicate", cfg); + ASSERT_TRUE(original.ok()); + EXPECT_NE(original.get()->as(), nullptr); +} + +TEST_F(CacheTest, BuilderReturningNullIsAnError) { + CacheFactory factory; + EXPECT_EQ( + factory.register_builder( + "TestBackend", + "null", + [](const CacheConfig&) { return std::shared_ptr{}; }), + Error::Ok); + EXPECT_EQ( + factory.build("TestBackend", "null", CacheConfig{32, 1, {flat_layer()}}) + .error(), + Error::Internal); +} + +TEST_F(CacheTest, NullCacheCannotBeInstalled) { + ET_EXPECT_DEATH( + { InstallGuard guard{std::shared_ptr{}}; }, + "Cannot install a null cache"); +} + +TEST_F(CacheTest, GuardInstallsOnCtorErasesOnDtor) { + auto cache = + std::make_shared(CacheConfig{4, 1, {flat_layer()}}); + std::string key; { - CacheSession session( - key, - std::make_shared(CacheConfig{4, 1, {flat_layer()}})); + InstallGuard guard(cache); + key = installed_key(guard); + ASSERT_FALSE(key.empty()); + // Publishing is the guard's alone, since CacheRegistry::install is + // private: an entry cannot outlive its owner or be clobbered. EXPECT_NE(CacheRegistry::global().get(key), nullptr); - EXPECT_TRUE(session.control()->can_extend(4)); + // The published entry is the same object the caller still holds. + EXPECT_EQ(CacheRegistry::global().get(key).get(), cache.get()); + EXPECT_TRUE(cache->as()->can_extend(4)); } EXPECT_EQ(CacheRegistry::global().get(key), nullptr); + // The guard held only the registry entry; the cache outlives it. + EXPECT_TRUE(cache->as()->can_extend(4)); } -// ---- ET adapter (maps core bool/optional to Error/Result) ------------------ - -TEST_F(CacheTest, EtAdapterMapsResultsAndCodes) { - SequenceCache cache(CacheConfig{2, 1, {flat_layer()}}); - auto ok = et::plan(cache, /*layer=*/0, /*position=*/0, /*T=*/2); - ASSERT_TRUE(ok.ok()); - EXPECT_EQ(ok->read[0].len, 2); - cache.commit(ok.get()); // accept the step so rewind has history to truncate - EXPECT_EQ( - et::plan(cache, 0, 2, 1).error(), Error::OutOfResources); // over capacity - EXPECT_FALSE(et::plan(cache, 5, 0, 1).ok()); // bad layer +TEST_F(CacheTest, AcquiredCacheOutlivesRegistryEntry) { + std::weak_ptr weak; + std::shared_ptr acquired; + std::string key; + { + auto cache = + std::make_shared(CacheConfig{4, 1, {flat_layer()}}); + weak = cache; + InstallGuard guard(cache); + key = installed_key(guard); + ASSERT_FALSE(key.empty()); + acquired = CacheRegistry::global().get(key); + ASSERT_NE(acquired, nullptr); + cache.reset(); + } - EXPECT_EQ(et::rewind(cache, 9), Error::InvalidArgument); // cannot grow - EXPECT_EQ(et::rewind(cache, 1), Error::Ok); + EXPECT_EQ(CacheRegistry::global().get(key), nullptr); + ASSERT_FALSE(weak.expired()); + EXPECT_TRUE(acquired->as()->can_extend(4)); + acquired.reset(); + EXPECT_TRUE(weak.expired()); } // ---- Cell layout ----------------------------------------------------------- @@ -291,8 +389,8 @@ struct Cells { int capacity, std::vector layers = {flat_layer(), flat_layer()}) : cache(CacheConfig{capacity, static_cast(layers.size()), layers}), - ctl(cache.as_batch_control()), - stepper(cache.as_cell_stepper()) {} + ctl(cache.as()), + stepper(cache.as()) {} // Ids come from the cache, so a test names sequences by allocating them. // Unlike the face's, this one unwraps and fails the test if none is free. @@ -434,7 +532,7 @@ TEST_F(CacheTest, CellPlacementIsSharedByEveryLayerOfTheStep) { const auto* first = c.place(0, args.positions); // layer 0 places the cells ASSERT_NE(first, nullptr); EXPECT_EQ(c.place(1, args.positions), first); // later layers reuse them - EXPECT_EQ(c.place(0, args.positions), nullptr); // asking twice is a new step + EXPECT_EQ(c.place(0, args.positions), first); // re-serving repeats the step EXPECT_EQ(c.cache.free_cells(), 14); // placed once, not once per layer } @@ -678,3 +776,22 @@ TEST_F(CacheTest, CellClearReturnsEveryCell) { EXPECT_EQ(c.ctl->next_pos(s0), 0); // the sequence is gone EXPECT_EQ(c.place(0, {0}), nullptr); // and the step went with it } + +TEST_F(CacheTest, KvSharedLayerReservesIdempotently) { + // A KV-shared layer re-serves its donor's id with the same tokens: the repeat + // returns the donor's step and claims no new cells. A re-serve with different + // tokens is a new step that never declared, and is refused. + Cells c(16, {flat_layer(), flat_layer()}); + const int32_t s0 = c.seq_new(); + const auto args = flatten_step({{s0, 0, 3}}); + ASSERT_TRUE(c.ctl->declare_step(args.seq_ids)); + + const auto* donor = c.place(0, args.positions); + ASSERT_NE(donor, nullptr); + const int free_after_place = c.cache.free_cells(); + + EXPECT_EQ(c.place(0, args.positions), donor); // same tokens -> same step + EXPECT_EQ(c.cache.free_cells(), free_after_place); // no new cells claimed + + EXPECT_EQ(c.place(0, {7, 8, 9}), nullptr); // different tokens, never declared +} diff --git a/extension/llm/cache/test_update_and_attend.py b/extension/llm/cache/test_update_and_attend.py index 96933a43b01..eed13f5b20c 100644 --- a/extension/llm/cache/test_update_and_attend.py +++ b/extension/llm/cache/test_update_and_attend.py @@ -11,15 +11,16 @@ from executorch.extension.llm.cache.reference_cache import ( attend, AttendSpec, + BatchedSequenceReferenceCache, CacheConfig, CacheSizing, CellReferenceCache, - ContiguousReferenceCache, flatten_step, LayerKind, LayerPolicy, MaskKind, MAX_SEQS, + SequenceReferenceCache, ) from executorch.extension.llm.cache.update_and_attend import REGISTRY, update_and_attend @@ -175,7 +176,7 @@ def test_prefill_matches_baseline(self): (CacheSizing.STATIC, seq_len), ]: with self.subTest(sizing=sizing): - cache = ContiguousReferenceCache(self._config(sizing, cap)) + cache = SequenceReferenceCache(self._config(sizing, cap)) REGISTRY.install(self.cache_key, cache) with REGISTRY.active(self.cache_key): out = ep.module()(x, _positions(0, seq_len), torch.arange(seq_len)) @@ -195,7 +196,7 @@ def test_incremental_decode_matches_baseline(self): (CacheSizing.STATIC, total), ]: with self.subTest(sizing=sizing): - cache = ContiguousReferenceCache(self._config(sizing, cap)) + cache = SequenceReferenceCache(self._config(sizing, cap)) REGISTRY.install(self.cache_key, cache) with REGISTRY.active(self.cache_key): ep_prefill.module()( @@ -221,7 +222,7 @@ def test_chunked_prefill_matches_baseline(self): ref = self.model.reference_forward(x, torch.arange(total)) ep = self._export(chunk) - cache = ContiguousReferenceCache(self._config(CacheSizing.DYNAMIC, total)) + cache = SequenceReferenceCache(self._config(CacheSizing.DYNAMIC, total)) REGISTRY.install(self.cache_key, cache) with REGISTRY.active(self.cache_key): for start in range(0, total, chunk): @@ -236,13 +237,35 @@ def test_chunked_prefill_matches_baseline(self): def test_static_overflow_raises(self): ep = self._export(seq_len=5) - cache = ContiguousReferenceCache(self._config(CacheSizing.STATIC, capacity=3)) + cache = SequenceReferenceCache(self._config(CacheSizing.STATIC, capacity=3)) REGISTRY.install(self.cache_key, cache) with self.assertRaises(RuntimeError), REGISTRY.active(self.cache_key): ep.module()( torch.randn(1, 5, self.hidden), _positions(0, 5), torch.arange(5) ) + def test_the_specs_must_cover_every_query_token(self): + # Each spec's queries are placed by the running total of the ones + # before it, so a cache that miscounts would attend the wrong slice + # rather than fail. Only this check separates the two. + class Miscounting: + def __init__(self, q_len): + self.q_len = q_len + + def update_and_fetch(self, layer_id, k, v, position): + return [AttendSpec(k=k, v=v, kind=MaskKind.NONE, q_len=self.q_len)] + + q = torch.randn(1, self.n_heads, 3, self.head_dim) + kv = torch.randn(1, self.n_kv_heads, 3, self.head_dim) + for q_len in (2, 4): # answering too few, and claiming too many + with self.subTest(q_len=q_len): + REGISTRY.install(self.cache_key, Miscounting(q_len)) + with self.assertRaisesRegex(ValueError, "of 3 query tokens"): + with REGISTRY.active(self.cache_key): + update_and_attend( + q, kv, kv, _positions(0, 3), 0, 0.125, torch.float32 + ) + def test_output_shape_uses_value_head_dim(self): # The output's last dim comes from v, which may differ from q's head dim # (e.g. MLA). Export (fake kernel only) and check the op node's meta. @@ -263,6 +286,295 @@ def forward(self, q, k, v, position): self.assertEqual(tuple(node.meta["val"].shape), (1, 4, 3, 5)) +class BatchedSequenceCacheTest(unittest.TestCase): + def setUp(self): + torch.manual_seed(0) + self.n_layers, self.hidden = 2, 16 + self.n_heads, self.n_kv_heads, self.head_dim = 4, 2, 8 + self.model = TinyAttentionModel( + self.n_layers, + self.hidden, + self.n_heads, + self.n_kv_heads, + self.head_dim, + 40, + ).eval() + self.cache_key = "batched-sequences" + + def tearDown(self): + REGISTRY.uninstall(self.cache_key) + + def _cache(self, capacity=16, layers=None): + cache = BatchedSequenceReferenceCache( + CacheConfig( + n_layers=self.n_layers, + n_kv_heads=self.n_kv_heads, + head_dim=self.head_dim, + capacity=capacity, + layers=[LayerPolicy.flat()] if layers is None else layers, + ) + ) + REGISTRY.install(self.cache_key, cache) + return cache + + def _step(self, cache, x, positions, seq_ids): + cache.declare_step(seq_ids) + with REGISTRY.active(self.cache_key): + return self.model(x, positions, torch.arange(x.shape[1])) + + def _attention_inputs(self, length): + return ( + torch.randn(1, self.n_heads, length, self.head_dim), + torch.randn(1, self.n_kv_heads, length, self.head_dim), + torch.randn(1, self.n_kv_heads, length, self.head_dim), + _positions(0, length), + ) + + def _attend(self, cache, inputs, layer_id=0): + # What the op does: fetch one spec per span, attend each over the query + # tokens it answers, rejoin. + q, k, v, positions = inputs + specs = cache.update_and_fetch(layer_id, k, v, positions) + outputs, start = [], 0 + for spec in specs: + end = start + spec.q_len + outputs.append( + attend( + q[:, :, start:end, :], + spec, + self.head_dim**-0.5, + torch.float32, + ) + ) + start = end + return outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=2) + + def test_single_span_matches_single_sequence(self): + x = torch.randn(1, 5, self.hidden) + out = self._step(self._cache(), x, _positions(0, 5), [3] * 5) + torch.testing.assert_close( + out, + self.model.reference_forward(x, torch.arange(5)), + atol=1e-4, + rtol=1e-4, + ) + + def test_multiple_sequence_spans_match_separate_runs(self): + a = torch.randn(1, 4, self.hidden) + b = torch.randn(1, 3, self.hidden) + tokens, positions, seq_ids, _ = flatten_step({2: (a, 0), 7: (b, 0)}) + + out = self._step(self._cache(), tokens, positions, seq_ids) + + torch.testing.assert_close( + out[:, :4], + self.model.reference_forward(a, torch.arange(4)), + atol=1e-4, + rtol=1e-4, + ) + torch.testing.assert_close( + out[:, 4:], + self.model.reference_forward(b, torch.arange(3)), + atol=1e-4, + rtol=1e-4, + ) + + def test_repeated_sequence_spans_preserve_input_order(self): + a = torch.randn(1, 3, self.hidden) + b = torch.randn(1, 1, self.hidden) + tokens = torch.cat([a[:, :2], b, a[:, 2:]], dim=1) + positions = torch.tensor([[0], [1], [0], [2]], dtype=torch.long) + out = self._step(self._cache(), tokens, positions, [2, 2, 7, 2]) + + a_out = self.model.reference_forward(a, torch.arange(3)) + b_out = self.model.reference_forward(b, torch.arange(1)) + torch.testing.assert_close(out[:, [0, 1, 3]], a_out, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(out[:, 2:3], b_out, atol=1e-4, rtol=1e-4) + + def test_decode_continues_each_private_sequence(self): + a = torch.randn(1, 4, self.hidden) + b = torch.randn(1, 3, self.hidden) + cache = self._cache() + + tokens, positions, seq_ids, _ = flatten_step( + {2: (a[:, :3], 0), 7: (b[:, :2], 0)} + ) + self._step(cache, tokens, positions, seq_ids) + + tokens, positions, seq_ids, _ = flatten_step( + {2: (a[:, 3:], 3), 7: (b[:, 2:], 2)} + ) + out = self._step(cache, tokens, positions, seq_ids) + + torch.testing.assert_close( + out[:, 0], + self.model.reference_forward(a, torch.arange(4))[:, -1], + atol=1e-4, + rtol=1e-4, + ) + torch.testing.assert_close( + out[:, 1], + self.model.reference_forward(b, torch.arange(3))[:, -1], + atol=1e-4, + rtol=1e-4, + ) + + def test_a_span_must_continue_its_own_sequence(self): + # A private history appends at its length, so a wrong position would + # still land contiguously. Only this check separates the two. + cache = self._cache() + q, k, v, _ = self._attention_inputs(2) + cache.declare_step([5, 5]) + + with self.assertRaisesRegex(ValueError, r"declares \[1, 2\], not \[0, 1\]"): + self._attend(cache, (q, k, v, _positions(1, 2))) + self.assertEqual(cache.seq_len(5), 0) # a refusal writes nothing + + # Ascending is not enough; a span is a consecutive run. + gapped = torch.tensor([[0], [2]], dtype=torch.long) + with self.assertRaisesRegex(ValueError, r"declares \[0, 2\], not \[0, 1\]"): + self._attend(cache, (q, k, v, gapped)) + self.assertEqual(cache.seq_len(5), 0) + + # The declaration still stands, so the same layer can be retried. + self._attend(cache, (q, k, v, _positions(0, 2))) + self.assertEqual(cache.seq_len(5), 2) + + def test_a_sequence_spanned_twice_continues_across_both(self): + cache = self._cache() + q, k, v, _ = self._attention_inputs(3) + + # Tokens 0 and 2 are seq 4, token 1 is seq 9; seq 4's second span picks + # up where its first left off rather than at its prior length. + cache.declare_step([4, 9, 4]) + self._attend(cache, (q, k, v, torch.tensor([[0], [0], [1]]))) + self.assertEqual(cache.seq_len(4), 2) + self.assertEqual(cache.seq_len(9), 1) + + # The bad position is in the last span, so a per-span check would have + # written the first two before refusing. + cache.declare_step([4, 9, 4]) + with self.assertRaisesRegex(ValueError, r"holds 3 .*declares \[4\], not \[3\]"): + self._attend(cache, (q, k, v, torch.tensor([[2], [1], [4]]))) + self.assertEqual(cache.seq_len(4), 2) + self.assertEqual(cache.seq_len(9), 1) + + def test_requires_one_declared_step_per_forward(self): + cache = self._cache() + inputs = self._attention_inputs(1) + + with self.assertRaisesRegex(RuntimeError, "no step declared"): + self._attend(cache, inputs) + + cache.declare_step([2]) + self._attend(cache, inputs) + with self.assertRaisesRegex(RuntimeError, "served twice"): + self._attend(cache, inputs) + + def test_declaration_and_sequence_verbs_validate_ids(self): + cache = self._cache() + with self.assertRaisesRegex(ValueError, "at least one token"): + cache.declare_step([]) + + for call in ( + lambda: cache.declare_step([-1]), + lambda: cache.seq_rm(-1), + lambda: cache.seq_len(-1), + ): + with self.subTest(call=call), self.assertRaises(ValueError): + call() + + # Private histories are dict entries, so nothing caps the id. + cache.declare_step([9999]) + self.assertEqual(cache.seq_len(9999), 0) + + def test_capacity_is_the_pool_total_and_refusal_changes_nothing(self): + cache = self._cache(capacity=4) + with self.assertRaisesRegex(RuntimeError, "exceeds capacity"): + cache.declare_step([2] * 5) + self.assertEqual(cache.seq_len(2), 0) + + # Two sequences share the budget rather than each getting one. + a = torch.randn(1, 2, self.hidden) + b = torch.randn(1, 2, self.hidden) + tokens, positions, seq_ids, _ = flatten_step({2: (a, 0), 7: (b, 0)}) + self._step(cache, tokens, positions, seq_ids) + self.assertEqual(cache.seq_len(2), 2) + self.assertEqual(cache.seq_len(7), 2) + + # Either sequence is now blocked by what the other holds. + with self.assertRaisesRegex(RuntimeError, "exceeds capacity"): + cache.declare_step([2]) + self.assertEqual(cache.seq_len(2), 2) + self.assertEqual(cache.seq_len(7), 2) + + def test_forward_width_must_match_tensors_and_declaration(self): + one = self._attention_inputs(1) + two = self._attention_inputs(2) + cases = ( + ("position", [2], (one[0], one[1], one[2], two[3]), "same token count"), + ("q/k", [2, 2], (two[0], one[1], one[2], two[3]), "same token count"), + ("k/v", [2], (one[0], one[1], two[2], one[3]), "same token count"), + ("declaration", [2, 2], one, "must match declare_step"), + ) + for name, seq_ids, inputs, message in cases: + with self.subTest(name=name): + cache = self._cache() + cache.declare_step(seq_ids) + with self.assertRaisesRegex(ValueError, message): + self._attend(cache, inputs) + + def test_sequence_removal_invalidates_a_declared_step(self): + cache = self._cache() + cache.declare_step([2]) + cache.seq_rm(2) + + self.assertEqual(cache.seq_len(2), 0) + with self.assertRaisesRegex(RuntimeError, "no step declared"): + self._attend(cache, self._attention_inputs(1)) + + def test_seq_rm_truncates_or_drops_and_refuses_a_bounded_range(self): + cache = self._cache() + self._step(cache, torch.randn(1, 4, self.hidden), _positions(0, 4), [1] * 4) + self._step(cache, torch.randn(1, 2, self.hidden), _positions(0, 2), [6] * 2) + + with self.assertRaises(NotImplementedError): + cache.seq_rm(1, 0, 2) + self.assertEqual(cache.seq_len(1), 4) + + cache.seq_rm(1, 2) # keep positions 0..1 + self.assertEqual(cache.seq_len(1), 2) + self.assertEqual(cache.seq_len(6), 2) # its neighbour is untouched + + cache.seq_rm(1) # the whole sequence + self.assertEqual(cache.seq_len(1), 0) + self.assertEqual(cache.seq_len(6), 2) + + def test_rewinding_then_continuing_matches_an_unbroken_run(self): + x = torch.randn(1, 5, self.hidden) + ref = self.model.reference_forward(x, torch.arange(5)) + + cache = self._cache() + self._step(cache, x[:, :4], _positions(0, 4), [3] * 4) + cache.seq_rm(3, 2) # discard positions 2..3 + out = self._step(cache, x[:, 2:], _positions(2, 3), [3] * 3) + + torch.testing.assert_close(out, ref[:, 2:], atol=1e-4, rtol=1e-4) + + def test_rewind_refuses_to_grow_or_pass_a_window(self): + cache = self._cache(layers=[LayerPolicy.ring(2)]) + self._step(cache, torch.randn(1, 5, self.hidden), _positions(0, 5), [0] * 5) + + with self.assertRaisesRegex(ValueError, "the history holds 5"): + cache.seq_rm(0, 6) + # A windowed layer keeps only its last two positions, so 3 is the floor + # even though this reference still holds the older ones. + with self.assertRaisesRegex(ValueError, "retains only from 3"): + cache.seq_rm(0, 1) + cache.seq_rm(0, 3) + self.assertEqual(cache.seq_len(0), 3) + + class CellCacheTest(unittest.TestCase): # Many sequences over one pool of per-token cells, flat on the token axis. # The baseline throughout is the cacheless model: whatever a sequence would @@ -305,7 +617,7 @@ def _cache( def _step(self, cache, x, positions, seqs): """One forward carrying `x`, whose tokens have these positions/seqs.""" - cache.begin_step(seqs) + cache.declare_step(seqs) pos = torch.tensor(positions, dtype=torch.long).unsqueeze(-1) with REGISTRY.active(self.cache_key): return self.model(x, pos, torch.arange(x.shape[1])) @@ -325,7 +637,7 @@ def test_batched_sequences_match_separate_runs(self): # {seq_id: (tokens, start_pos)} -> the step's parallel arrays tokens, positions, seq_ids, _ = flatten_step({0: (a, 0), 1: (b, 0)}) - cache.begin_step(seq_ids) + cache.declare_step(seq_ids) with REGISTRY.active(self.cache_key): # every row, not one per sequence: each token is compared below out = self.model(tokens, positions, torch.arange(tokens.shape[1])) @@ -352,14 +664,14 @@ def test_batched_decode_continues_each_sequence(self): tokens, positions, seq_ids, logits_indices = flatten_step( {0: (a[:, :3], 0), 1: (b[:, :2], 0)} ) - cache.begin_step(seq_ids) + cache.declare_step(seq_ids) with REGISTRY.active(self.cache_key): self.model(tokens, positions, logits_indices) tokens, positions, seq_ids, logits_indices = flatten_step( {0: (a[:, 3:], 3), 1: (b[:, 2:], 2)} ) - cache.begin_step(seq_ids) + cache.declare_step(seq_ids) with REGISTRY.active(self.cache_key): out = self.model(tokens, positions, logits_indices) @@ -433,18 +745,18 @@ def test_fork_at_a_position_shares_only_the_prefix(self): def test_freeing_the_tail_shrinks_the_read_window(self): cache = self._cache() kv = torch.randn(1, self.n_kv_heads, 4, self.head_dim) - cache.begin_step([0] * 4) - k, _, _ = cache.update_and_fetch(0, kv, kv, _positions(0, 4)) - self.assertEqual(k.shape[2], 4) # four cells held, so a window of four + cache.declare_step([0] * 4) + spec = cache.update_and_fetch(0, kv, kv, _positions(0, 4))[0] + self.assertEqual(spec.k.shape[2], 4) # four cells held, so a window of four cache.seq_rm(0) # frees all four, so used_end walks back to 0 self.assertEqual(cache.free_cells(), self.CAPACITY) # one token reclaims cell 0, so the window is its own single cell kv = torch.randn(1, self.n_kv_heads, 1, self.head_dim) - cache.begin_step([1]) - k, _, spec = cache.update_and_fetch(0, kv, kv, torch.tensor([[0]])) - self.assertEqual(k.shape[2], 1) # the window length is 1, not the old 4 + cache.declare_step([1]) + spec = cache.update_and_fetch(0, kv, kv, torch.tensor([[0]]))[0] + self.assertEqual(spec.k.shape[2], 1) # the window length is 1, not the old 4 self.assertEqual(spec.mask.shape[-1], 1) def test_seq_rm_over_a_range_frees_only_that_window(self): @@ -464,7 +776,7 @@ def test_every_verb_range_checks_the_seq_id(self): # much later as an overflow while building the mask. cache = self._cache() for call in ( - lambda: cache.begin_step([MAX_SEQS]), + lambda: cache.declare_step([MAX_SEQS]), lambda: cache.seq_cp(0, MAX_SEQS), lambda: cache.seq_cp(MAX_SEQS, 0), lambda: cache.seq_rm(MAX_SEQS), @@ -491,8 +803,8 @@ def test_layer_policy_rejects_a_mismatched_window(self): def test_window_narrows_each_query_without_crossing_sequences(self): cache = self._cache(layers=[LayerPolicy.ring(2)]) kv = torch.randn(1, self.n_kv_heads, 4, self.head_dim) - cache.begin_step([0] * 4) - spec = cache.update_and_fetch(0, kv, kv, _positions(0, 4))[2] + cache.declare_step([0] * 4) + spec = cache.update_and_fetch(0, kv, kv, _positions(0, 4))[0] # offsets[i][j] = j - i, so <= 0 is causal and > -2 keeps the newest # two: a band whose row 2 drops key 0, which plain causal would keep. @@ -501,9 +813,9 @@ def test_window_narrows_each_query_without_crossing_sequences(self): # a second sequence is bounded the same way, and still sees none of # the first's cells even though they are inside its window - cache.begin_step([1, 1]) + cache.declare_step([1, 1]) kv = torch.randn(1, self.n_kv_heads, 2, self.head_dim) - spec = cache.update_and_fetch(0, kv, kv, _positions(0, 2))[2] + spec = cache.update_and_fetch(0, kv, kv, _positions(0, 2))[0] expected = torch.zeros(2, 6, dtype=torch.bool) expected[0, 4] = expected[1, 4] = expected[1, 5] = True torch.testing.assert_close(spec.mask, expected) @@ -513,10 +825,10 @@ def test_layers_can_window_independently(self): # per policy, not per layer, so a mixed model costs one extra mask. cache = self._cache(layers=[LayerPolicy.flat(), LayerPolicy.ring(2)]) kv = torch.randn(1, self.n_kv_heads, 4, self.head_dim) - cache.begin_step([0] * 4) + cache.declare_step([0] * 4) pos = _positions(0, 4) - flat = cache.update_and_fetch(0, kv, kv, pos)[2].mask - windowed = cache.update_and_fetch(1, kv, kv, pos)[2].mask + flat = cache.update_and_fetch(0, kv, kv, pos)[0].mask + windowed = cache.update_and_fetch(1, kv, kv, pos)[0].mask offsets = torch.arange(4) - torch.arange(4).unsqueeze(-1) torch.testing.assert_close(flat, offsets <= 0) @@ -530,11 +842,11 @@ def test_layers_sharing_a_window_share_one_mask(self): n_layers=3, ) kv = torch.randn(1, self.n_kv_heads, 4, self.head_dim) - cache.begin_step([0] * 4) + cache.declare_step([0] * 4) pos = _positions(0, 4) - flat = cache.update_and_fetch(0, kv, kv, pos)[2].mask - first = cache.update_and_fetch(1, kv, kv, pos)[2].mask - second = cache.update_and_fetch(2, kv, kv, pos)[2].mask + flat = cache.update_and_fetch(0, kv, kv, pos)[0].mask + first = cache.update_and_fetch(1, kv, kv, pos)[0].mask + second = cache.update_and_fetch(2, kv, kv, pos)[0].mask self.assertIs(first, second) self.assertIsNot(flat, first) @@ -543,22 +855,25 @@ def test_windowed_decode_attends_only_the_retained_cells(self): window = 2 cache = self._cache(layers=[LayerPolicy.ring(window)]) kv = torch.randn(1, self.n_kv_heads, 4, self.head_dim) - cache.begin_step([0] * 4) + cache.declare_step([0] * 4) cache.update_and_fetch(0, kv, kv, _positions(0, 4)) - cache.begin_step([0]) + cache.declare_step([0]) kv = torch.randn(1, self.n_kv_heads, 1, self.head_dim) - k, v, spec = cache.update_and_fetch(0, kv, kv, _positions(4, 1)) + spec = cache.update_and_fetch(0, kv, kv, _positions(4, 1))[0] q = torch.randn(1, self.n_heads, 1, self.head_dim) scale = self.head_dim**-0.5 torch.testing.assert_close( - attend(q, k, v, spec, scale, torch.float32), + attend(q, spec, scale, torch.float32), attend( # the last `window` cells of its sequence, unmasked q, - k[:, :, -window:, :], - v[:, :, -window:, :], - AttendSpec(kind=MaskKind.NONE), + AttendSpec( + k=spec.k[:, :, -window:, :], + v=spec.v[:, :, -window:, :], + kind=MaskKind.NONE, + q_len=1, + ), scale, torch.float32, ), @@ -568,7 +883,7 @@ def test_admission_fails_before_the_forward(self): cache = self._cache(capacity=4) self.assertFalse(cache.can_extend(5)) with self.assertRaises(RuntimeError): - cache.begin_step([0] * 5) + cache.declare_step([0] * 5) def test_step_protocol_is_enforced(self): cache = self._cache() @@ -576,17 +891,17 @@ def test_step_protocol_is_enforced(self): pos = torch.tensor([[0]]) with self.assertRaises(ValueError): # a step with no tokens - cache.begin_step([]) + cache.declare_step([]) - cache.begin_step([0, 0]) # declares two tokens, forward carries one + cache.declare_step([0, 0]) # declares two tokens, forward carries one with self.assertRaises(ValueError): cache.update_and_fetch(0, kv, kv, pos) with self.assertRaises(RuntimeError): # the failed attempt still cleared it cache.update_and_fetch(0, kv, kv, torch.tensor([[0], [1]])) - cache.begin_step([0]) + cache.declare_step([0]) cache.update_and_fetch(0, kv, kv, pos) - with self.assertRaises(RuntimeError): # a second step, no begin_step + with self.assertRaises(RuntimeError): # a second step, no declare_step cache.update_and_fetch(0, kv, kv, pos) def test_growth_keeps_cell_indices_and_bytes(self): @@ -595,18 +910,19 @@ def test_growth_keeps_cell_indices_and_bytes(self): # move history without anything noticing. cache = self._cache() first = torch.randn(1, self.n_kv_heads, 2, self.head_dim) - cache.begin_step([0, 0]) - k, _, _ = cache.update_and_fetch(0, first, first, torch.tensor([[0], [1]])) - self.assertEqual(k.shape[2], 2) # a short session reserves a short pool + cache.declare_step([0, 0]) + spec = cache.update_and_fetch(0, first, first, torch.tensor([[0], [1]]))[0] + self.assertEqual(spec.k.shape[2], 2) # a short session reserves a short pool rest = torch.randn(1, self.n_kv_heads, 6, self.head_dim) - cache.begin_step([0] * 6) - k, v, _ = cache.update_and_fetch( + cache.declare_step([0] * 6) + spec = cache.update_and_fetch( 0, rest, rest, torch.tensor([[p] for p in range(2, 8)]) - ) - self.assertEqual(k.shape[2], 8) - torch.testing.assert_close(k[:, :, :2, :], first) # cells 0,1 unmoved - torch.testing.assert_close(v[:, :, 2:, :], rest) + )[0] + self.assertEqual(spec.k.shape[2], 8) + # cells 0,1 unmoved + torch.testing.assert_close(spec.k[:, :, :2, :], first) + torch.testing.assert_close(spec.v[:, :, 2:, :], rest) def test_sizings_agree(self): x = torch.randn(1, 5, self.hidden) @@ -616,14 +932,14 @@ def test_sizings_agree(self): ] torch.testing.assert_close(out[0], out[1]) - def test_a_verb_does_not_hide_a_missing_begin_step(self): + def test_a_verb_does_not_hide_a_missing_declare_step(self): # A sequence verb drops the memoized plan, which must not be mistaken # for the start of a step -- that would silently reuse the previous # step's sequence assignment for the new tokens. cache = self._cache() kv = torch.randn(1, self.n_kv_heads, 2, self.head_dim) pos = torch.tensor([[0], [0]]) - cache.begin_step([0, 1]) + cache.declare_step([0, 1]) cache.update_and_fetch(0, kv, kv, pos) cache.seq_rm(2) # any verb; a no-op here beyond dropping the plan @@ -633,18 +949,18 @@ def test_a_verb_does_not_hide_a_missing_begin_step(self): cache.update_and_fetch(1, kv, kv, pos) -class ContiguousSpecTest(unittest.TestCase): +class SequenceSpecTest(unittest.TestCase): # Which mask semantic the cache declares for each shape of step. def setUp(self): torch.manual_seed(0) - self.cache = ContiguousReferenceCache( + self.cache = SequenceReferenceCache( CacheConfig(n_layers=1, n_kv_heads=2, head_dim=4, capacity=8) ) def _update(self, start, q_len): kv = torch.randn(1, 2, q_len, 4) - return self.cache.update_and_fetch(0, kv, kv, _positions(start, q_len))[2] + return self.cache.update_and_fetch(0, kv, kv, _positions(start, q_len))[0] def test_decode_is_unmasked(self): self.assertEqual(self._update(0, 1).kind, MaskKind.NONE) @@ -674,7 +990,7 @@ def setUp(self): self.scale = self.DIM**-0.5 def _cache(self, policy, sizing=CacheSizing.DYNAMIC): - return ContiguousReferenceCache( + return SequenceReferenceCache( CacheConfig( n_layers=1, n_kv_heads=self.HEADS, @@ -689,50 +1005,51 @@ def _update(self, cache, n, start): kv = torch.randn(1, self.HEADS, n, self.DIM) return cache.update_and_fetch(0, kv, kv, _positions(start, n)) - def _attend(self, q, k, v, spec): - return attend(q, k, v, spec, self.scale, torch.float32) + def _attend(self, q, spec): + return attend(q, spec, self.scale, torch.float32) + + def _unmasked(self, q, k, v): + # The same queries over a hand-picked window, for the spec to match. + spec = AttendSpec(k=k, v=v, kind=MaskKind.NONE, q_len=q.shape[-2]) + return attend(q, spec, self.scale, torch.float32) def test_decode_attends_only_the_window(self): window = 3 cache = self._cache(LayerPolicy.ring(window)) self._update(cache, 5, 0) - k, v, spec = self._update(cache, 1, 5) + spec = self._update(cache, 1, 5)[0] q = torch.randn(1, self.HEADS, 1, self.DIM) torch.testing.assert_close( - self._attend(q, k, v, spec), - self._attend( # the last `window` cells, unmasked - q, - k[:, :, -window:, :], - v[:, :, -window:, :], - AttendSpec(kind=MaskKind.NONE), + self._attend(q, spec), + self._unmasked( # the last `window` cells, unmasked + q, spec.k[:, :, -window:, :], spec.v[:, :, -window:, :] ), ) def test_each_prefill_query_attends_its_own_window(self): window = 2 cache = self._cache(LayerPolicy.ring(window)) - k, v, spec = self._update(cache, 4, 0) + spec = self._update(cache, 4, 0)[0] self.assertEqual(spec.kind, MaskKind.EXPLICIT) q = torch.randn(1, self.HEADS, 4, self.DIM) - out = self._attend(q, k, v, spec) + out = self._attend(q, spec) for i in range(4): # query at position i sees (i - window, i] lo = max(0, i - window + 1) torch.testing.assert_close( out[:, :, i : i + 1, :], - self._attend( + self._unmasked( q[:, :, i : i + 1, :], - k[:, :, lo : i + 1, :], - v[:, :, lo : i + 1, :], - AttendSpec(kind=MaskKind.NONE), + spec.k[:, :, lo : i + 1, :], + spec.v[:, :, lo : i + 1, :], ), ) def test_layers_can_window_independently(self): # gemma-style: only some layers are windowed, so one step yields two # different semantics from the same cache. - cache = ContiguousReferenceCache( + cache = SequenceReferenceCache( CacheConfig( n_layers=2, n_kv_heads=self.HEADS, @@ -743,8 +1060,8 @@ def test_layers_can_window_independently(self): ) kv = torch.randn(1, self.HEADS, 4, self.DIM) pos = _positions(0, 4) - flat = cache.update_and_fetch(0, kv, kv, pos)[2] - windowed = cache.update_and_fetch(1, kv, kv, pos)[2] + flat = cache.update_and_fetch(0, kv, kv, pos)[0] + windowed = cache.update_and_fetch(1, kv, kv, pos)[0] self.assertEqual(flat.kind, MaskKind.CAUSAL) self.assertEqual(windowed.kind, MaskKind.EXPLICIT) @@ -757,7 +1074,7 @@ def test_windowed_continuation_bounds_the_band_at_both_ends(self): window = 2 cache = self._cache(LayerPolicy.ring(window)) self._update(cache, 4, 0) - _, _, spec = self._update(cache, 3, 4) + spec = self._update(cache, 3, 4)[0] self.assertEqual(spec.kind, MaskKind.EXPLICIT) q_len, total = 3, 7 @@ -773,8 +1090,8 @@ def test_window_equal_to_history_stays_fused(self): # later the band appears. window = 4 cache = self._cache(LayerPolicy.ring(window)) - self.assertEqual(self._update(cache, window, 0)[2].kind, MaskKind.CAUSAL) - self.assertEqual(self._update(cache, 1, window)[2].kind, MaskKind.EXPLICIT) + self.assertEqual(self._update(cache, window, 0)[0].kind, MaskKind.CAUSAL) + self.assertEqual(self._update(cache, 1, window)[0].kind, MaskKind.EXPLICIT) def test_static_sizing_windows_like_dynamic(self): # STATIC writes into a preallocated buffer and slices it; the window is @@ -782,13 +1099,13 @@ def test_static_sizing_windows_like_dynamic(self): window = 2 torch.manual_seed(0) static = self._cache(LayerPolicy.ring(window), sizing=CacheSizing.STATIC) - sk, sv, s_spec = self._update(static, 4, 0) + s_spec = self._update(static, 4, 0)[0] torch.manual_seed(0) dynamic = self._cache(LayerPolicy.ring(window)) - dk, dv, d_spec = self._update(dynamic, 4, 0) + d_spec = self._update(dynamic, 4, 0)[0] - torch.testing.assert_close(sk, dk) - torch.testing.assert_close(sv, dv) + torch.testing.assert_close(s_spec.k, d_spec.k) + torch.testing.assert_close(s_spec.v, d_spec.v) self.assertEqual(s_spec.kind, d_spec.kind) torch.testing.assert_close(s_spec.mask, d_spec.mask) @@ -796,8 +1113,8 @@ def test_window_wider_than_history_stays_fused(self): # Nothing to bound from below, so the window must not force a mask. for policy in (LayerPolicy.flat(), LayerPolicy.ring(64)): cache = self._cache(policy) - self.assertEqual(self._update(cache, 4, 0)[2].kind, MaskKind.CAUSAL) - self.assertEqual(self._update(cache, 1, 4)[2].kind, MaskKind.NONE) + self.assertEqual(self._update(cache, 4, 0)[0].kind, MaskKind.CAUSAL) + self.assertEqual(self._update(cache, 1, 4)[0].kind, MaskKind.NONE) class AttendExplicitTest(unittest.TestCase): @@ -812,16 +1129,21 @@ def setUp(self): self.v = torch.randn(1, 2, self.total, self.head_dim) self.scale = self.head_dim**-0.5 - def _attend(self, spec, k=None, v=None): - k = self.k if k is None else k - v = self.v if v is None else v - return attend(self.q, k, v, spec, self.scale, torch.float32) + def _attend(self, kind, mask=None, k=None, v=None): + spec = AttendSpec( + k=self.k if k is None else k, + v=self.v if v is None else v, + kind=kind, + q_len=self.q_len, + mask=mask, + ) + return attend(self.q, spec, self.scale, torch.float32) def test_causal_rejects_a_non_square_window(self): # torch's is_causal is upper-left, so it cannot serve a continuation; # a cache must declare EXPLICIT there rather than CAUSAL. with self.assertRaises(ValueError): - self._attend(AttendSpec(kind=MaskKind.CAUSAL)) + self._attend(MaskKind.CAUSAL) def test_explicit_attends_the_true_cells(self): # Polarity: masking to cells {0, 2} must equal attending over just those @@ -830,11 +1152,11 @@ def test_explicit_attends_the_true_cells(self): mask = torch.zeros(self.q_len, self.total, dtype=torch.bool) mask[:, keep] = True torch.testing.assert_close( - self._attend(AttendSpec(kind=MaskKind.EXPLICIT, mask=mask)), + self._attend(MaskKind.EXPLICIT, mask=mask), self._attend( - AttendSpec(kind=MaskKind.NONE), - self.k.index_select(2, keep), - self.v.index_select(2, keep), + MaskKind.NONE, + k=self.k.index_select(2, keep), + v=self.v.index_select(2, keep), ), ) diff --git a/extension/llm/cache/update_and_attend.py b/extension/llm/cache/update_and_attend.py index 9feaabc4038..227f0aebbe6 100644 --- a/extension/llm/cache/update_and_attend.py +++ b/extension/llm/cache/update_and_attend.py @@ -95,8 +95,16 @@ def update_and_attend( Returns: ``[B, H_q, q_len, v_head_dim]`` attention output, in ``out_dtype``. """ - k_hist, v_hist, spec = REGISTRY.current().update_and_fetch(layer_id, k, v, position) - return attend(q, k_hist, v_hist, spec, scale, out_dtype) + specs = REGISTRY.current().update_and_fetch(layer_id, k, v, position) + outputs = [] + start = 0 + for spec in specs: + end = start + spec.q_len + outputs.append(attend(q[:, :, start:end, :], spec, scale, out_dtype)) + start = end + if start != q.shape[-2]: + raise ValueError(f"the cache answered {start} of {q.shape[-2]} query tokens") + return outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=2) @update_and_attend.register_fake diff --git a/extension/llm/custom_ops/CMakeLists.txt b/extension/llm/custom_ops/CMakeLists.txt index c640cad6706..63ff2f23e69 100644 --- a/extension/llm/custom_ops/CMakeLists.txt +++ b/extension/llm/custom_ops/CMakeLists.txt @@ -20,7 +20,7 @@ set(_common_compile_options $<$:/wd4996> $<$>:-Wno-deprecated-declarations -fPIC> ) -if(CMAKE_SYSTEM_PROCESSOR MATCHES "arm64|aarch64") +if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(arm64|aarch64)$") list(APPEND _common_compile_options "$<$>:-march=armv8.2-a+dotprod>" ) @@ -86,28 +86,72 @@ target_link_libraries(custom_ops PUBLIC ${custom_ops_libs} executorch_core) # The MoE kernel always compiles with a reference fallback (unpack + dequant + # cpublas::gemm) using the torchao weight_packing headers from third-party/ao # (already on the include path). Pass -# -DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON to additionally link the -# optimized torchao linear_operator (fused INT8-dyn-act GEMM, aarch64 NEON +# -DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED=ON to additionally link +# the optimized torchao linear_operator (fused INT8-dyn-act GEMM, aarch64 NEON # dotprod). -option(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE +option(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED "Link the optimized torchao linear kernel for llama::quantized_moe_ffn" OFF ) -if(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE) - if(NOT TARGET torchao_ops_linear_8bit_act_xbit_weight_executorch) + +function(_target_enable_quantized_moe_torchao target) + target_compile_definitions( + ${target} + PRIVATE TORCHAO_BUILD_CPU_AARCH64=1 TORCHAO_ENABLE_ARM_NEON_DOT=1 + TORCHAO_PARALLEL_EXECUTORCH=1 + TORCHAO_SHARED_KERNELS_BUILD_EXECUTORCH=1 + ) +endfunction() + +if(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED) + if(NOT CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$") message( FATAL_ERROR - "EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON requires target " - "torchao_ops_linear_8bit_act_xbit_weight_executorch, which is not " - "defined. Build the torchao ops or set this option OFF." + "EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED=ON is supported only on " + "aarch64 and arm64." ) endif() - # Compile definition and link must be gated on the same condition, else the - # ENABLE_QUANTIZED_MOE_FFN path compiles without the library that defines it. - target_compile_definitions(custom_ops PUBLIC ENABLE_QUANTIZED_MOE_FFN=1) - target_link_libraries( - custom_ops PUBLIC torchao_ops_linear_8bit_act_xbit_weight_executorch + + # The selected target is also reused by custom_ops_aot_lib below. + if(EXECUTORCH_BUILD_KERNELS_TORCHAO) + if(NOT TARGET torchao_ops_executorch) + message(FATAL_ERROR "EXECUTORCH_BUILD_KERNELS_TORCHAO=ON but target " + "torchao_ops_executorch is not defined." + ) + endif() + set(quantized_moe_torchao_target torchao_ops_executorch) + else() + set(quantized_moe_torchao_target torchao_moe_linear) + add_library( + torchao_moe_linear STATIC + ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight/linear_8bit_act_xbit_weight.cpp + ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/torch_free_kernels/aarch64/quantization/quantize.cpp + ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/torch_free_kernels/aarch64/reduction/compute_sum.cpp + ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/torch_free_kernels/aarch64/reduction/find_min_and_max.cpp + ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/torch_free_kernels/aarch64/valpacking/interleave.cpp + ) + target_include_directories( + torchao_moe_linear PRIVATE ${EXECUTORCH_ROOT}/third-party/ao + ) + _target_enable_quantized_moe_torchao(torchao_moe_linear) + target_link_libraries( + torchao_moe_linear PRIVATE cpuinfo executorch_core extension_threadpool + ) + target_compile_options( + torchao_moe_linear PRIVATE ${_common_compile_options} + ) + install( + TARGETS torchao_moe_linear + EXPORT ExecuTorchTargets + DESTINATION ${CMAKE_INSTALL_LIBDIR} + ) + endif() + + target_compile_definitions( + custom_ops PRIVATE EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO=1 ) + _target_enable_quantized_moe_torchao(custom_ops) + target_link_libraries(custom_ops PUBLIC ${quantized_moe_torchao_target}) endif() target_compile_options(custom_ops PUBLIC ${_common_compile_options}) @@ -190,21 +234,13 @@ if(EXECUTORCH_BUILD_KERNELS_LLM_AOT) custom_ops_aot_lib PUBLIC cpublas torch extension_tensor extension_threadpool ) - if(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE) - if(NOT TARGET torchao_ops_linear_8bit_act_xbit_weight_executorch) - message( - FATAL_ERROR - "EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON requires target " - "torchao_ops_linear_8bit_act_xbit_weight_executorch, which is not " - "defined. Build the torchao ops or set this option OFF." - ) - endif() + if(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED) target_compile_definitions( - custom_ops_aot_lib PUBLIC ENABLE_QUANTIZED_MOE_FFN=1 + custom_ops_aot_lib PRIVATE EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO=1 ) + _target_enable_quantized_moe_torchao(custom_ops_aot_lib) target_link_libraries( - custom_ops_aot_lib - PUBLIC torchao_ops_linear_8bit_act_xbit_weight_executorch + custom_ops_aot_lib PUBLIC ${quantized_moe_torchao_target} ) endif() if(WIN32) diff --git a/extension/llm/custom_ops/op_moe.cpp b/extension/llm/custom_ops/op_moe.cpp index bcb0e88df10..e218d58b35a 100644 --- a/extension/llm/custom_ops/op_moe.cpp +++ b/extension/llm/custom_ops/op_moe.cpp @@ -9,19 +9,28 @@ #include #include +#if defined(EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO) && \ + !defined(TORCHAO_PARALLEL_EXECUTORCH) #include +#endif #include #include +#if defined(EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO) && \ + !defined(TORCHAO_PARALLEL_EXECUTORCH) #include +#endif #include #include -#ifdef ENABLE_QUANTIZED_MOE_FFN -#include +#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO +#include +#include #include +#include +#include #include // std::nullopt, used only by the optimized aarch64 path -#endif // ENABLE_QUANTIZED_MOE_FFN +#endif // EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO #include #include @@ -37,6 +46,49 @@ namespace { using ::executorch::aten::string_view; +#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO +template +const torchao::ops::linear_8bit_act_xbit_weight::UKernelConfig& +universal_ukernel_config() { + using torchao::ops::linear_8bit_act_xbit_weight::UKernelConfig; + namespace kernel = torchao::kernels::cpu::aarch64::linear:: + channelwise_8bit_activation_groupwise_lowbit_weight; + + static const auto config = [] { + ET_CHECK_MSG( + cpuinfo_initialize() && cpuinfo_has_arm_neon_dot(), + "quantized_moe_ffn optimized path requires Arm NEON dot product"); + auto result = UKernelConfig::make( + /*preferred_alignment=*/16, + /*n_step=*/8, + /*nr=*/8, + /*kr=*/16, + /*sr=*/2, + kWeightNbit, + /*has_weight_zeros=*/false, + /*has_bias=*/false, + &torchao::weight_packing::packed_weights_size, + &torchao::weight_packing::packed_weights_offset, + &torchao::weight_packing::pack_weights, + {}); + result.linear_configs[0] = UKernelConfig::linear_config_type({ + /*m_step=*/1, + /*mr=*/1, + &kernel::packed_activations_size, + &kernel::packed_activations_offset, + &kernel::pack_activations<1, 16, 2>, + &kernel::kernel_1x8x16_f32_neondot< + kWeightNbit, + /*has_weight_zeros=*/false, + /*has_lut=*/false>, + }); + result.validate(); + return result; + }(); + return config; +} +#endif + // Numerically-stable sigmoid. Branching on sign keeps exp()'s argument // non-positive on both sides, so it can never overflow. inline float stable_sigmoid(float v) { @@ -189,7 +241,7 @@ inline void reference_linear( // Dispatch a single per-expert grouped GEMM through torchao's // linear_operator (optimized, aarch64) or reference unpack+dequant+gemm. -#ifdef ENABLE_QUANTIZED_MOE_FFN +#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO template inline void torchao_linear( const uint8_t* packed_w_blob, @@ -205,21 +257,22 @@ inline void torchao_linear( static_cast(torchao::ops::PackedWeightsHeader::size()), "torchao packed blob too small to contain header"); auto header = torchao::ops::PackedWeightsHeader::read(packed_w_blob); - // Select the ukernel from the format declared in the header. This resolves - // the universal or kleidi packing automatically; a format whose kernels are - // not compiled into this build (e.g. kleidi when TORCHAO_ENABLE_KLEIDI is - // unset) throws here instead of being silently mis-read. - // TODO: enable KleidiAI here — build this op on xplat arm64 with - // -DTORCHAO_ENABLE_KLEIDI (+ -DTORCHAO_ENABLE_ARM_I8MM) and link the kleidi - // kernel target so a kleidi header actually resolves to a kleidi ukernel. - // Must be coordinated with the AoT packer emitting kleidi headers (see - // targets.bzl). - auto uk = torchao::ops::linear_8bit_act_xbit_weight::select_ukernel_config< - kWeightNbit>(header); - - // Validate the blob against the *selected* format's layout. nr/kr/sr and the - // size formula differ between universal and kleidi, so derive them from the - // chosen config rather than assuming a fixed layout. + ET_CHECK_MSG( + header.type == + torchao::ops::PackedWeightsType:: + linear_8bit_act_xbit_weight_universal, + "quantized_moe_ffn requires universal torchao packed weights"); + const auto format = torchao::ops::linear_8bit_act_xbit_weight:: + PackedWeightsFormat::from_packed_weights_header(header); + ET_CHECK_MSG( + format.weight_nbit == kWeightNbit && !format.has_weight_zeros && + !format.has_bias && format.nr == 8 && format.kr == 16 && + format.sr == 2, + "quantized_moe_ffn received an unsupported universal weight format"); + const auto& uk = universal_ukernel_config(); + + // Validate the blob against the universal config's layout without + // duplicating its packed-weight size formula. const int64_t required_bytes = static_cast(torchao::ops::PackedWeightsHeader::size()) + static_cast(uk.packed_weights_size( @@ -258,7 +311,7 @@ inline void torchao_linear( /*clamp_min=*/0.0f, /*clamp_max=*/0.0f); } -#endif // ENABLE_QUANTIZED_MOE_FFN +#endif // EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO inline void expert_linear_dispatch( int64_t weight_nbit, @@ -270,12 +323,21 @@ inline void expert_linear_dispatch( int64_t k, int64_t group_size, float* out) { -#ifndef ENABLE_QUANTIZED_MOE_FFN +#ifndef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO // Reference path only: it unpacks the universal layout, so validate the blob // holds the header plus the universal packed weight-data bytes for the // claimed dims before any path dereferences it. The torchao path validates - // against its own selected format (universal or kleidi) inside - // torchao_linear. + // the same required universal format inside torchao_linear. + ET_CHECK_MSG( + packed_blob_bytes >= + static_cast(torchao::ops::PackedWeightsHeader::size()), + "torchao packed blob too small to contain header"); + const auto header = torchao::ops::PackedWeightsHeader::read(packed_w_blob); + ET_CHECK_MSG( + header.type == + torchao::ops::PackedWeightsType:: + linear_8bit_act_xbit_weight_universal, + "quantized_moe_ffn requires universal torchao packed weights"); constexpr int kNr = 8, kKr = 16, kSr = 2; const int64_t required_bytes = static_cast(torchao::ops::PackedWeightsHeader::size()) + @@ -299,10 +361,10 @@ inline void expert_linear_dispatch( static_cast(k), static_cast(group_size), static_cast(weight_nbit)); -#endif // !ENABLE_QUANTIZED_MOE_FFN +#endif // !EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO switch (weight_nbit) { case 4: -#ifdef ENABLE_QUANTIZED_MOE_FFN +#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO torchao_linear<4>( packed_w_blob, packed_blob_bytes, x, m, n, k, group_size, out); #else @@ -311,7 +373,7 @@ inline void expert_linear_dispatch( #endif return; case 8: -#ifdef ENABLE_QUANTIZED_MOE_FFN +#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO torchao_linear<8>( packed_w_blob, packed_blob_bytes, x, m, n, k, group_size, out); #else @@ -634,26 +696,20 @@ Tensor& quantized_moe_ffn_out( } }; -#ifdef ENABLE_QUANTIZED_MOE_FFN - // torchao linear path (perf-sensitive). The kernel threads internally on the - // shared pool in the common config, or runs single-threaded when only the - // thread-pool-free variant is linked. Distribute experts across the pool - // ourselves only when the kernel won't and the pool has more than one thread - // -- running both would nest on one pthreadpool and deadlock. +#if defined(EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO) && \ + !defined(TORCHAO_PARALLEL_EXECUTORCH) const bool parallelize_experts = torchao::ops::linear_8bit_act_xbit_weight:: linear_operator_num_threads() == 1 && ::executorch::extension::threadpool::get_threadpool() ->get_thread_count() > 1; -#else - // Portable reference path: prefer simplicity over speed and run experts - // serially. - const bool parallelize_experts = false; -#endif if (parallelize_experts) { torch::executor::parallel_for(0, E, /*grain_size=*/1, run_experts); } else { run_experts(0, E); } +#else + run_experts(0, E); +#endif // ----- 7. Weighted scatter-add unpermute (cross-expert reduction) ----- // Each token sums the contributions of its top-k experts; run serially to diff --git a/extension/llm/custom_ops/op_sdpa_impl.h b/extension/llm/custom_ops/op_sdpa_impl.h index f6ed378ec03..7e7275c0427 100644 --- a/extension/llm/custom_ops/op_sdpa_impl.h +++ b/extension/llm/custom_ops/op_sdpa_impl.h @@ -1045,11 +1045,17 @@ void cpu_flash_attention( is_causal ? std::min(m + start_pos + qBlockSize, kvSize) : kvSize; int64_t m_start_pos = m + start_pos; auto j_kv = j / num_reps; - fill_stub(dst_data, static_cast(0), qSplitSize * headSize); + fill_stub(dst_data, static_cast(0), qBlockSize * headSize); for (int64_t n = 0; n < num_keys; n += kvSplitSize) { - int64_t kvBlockSize = std::min(kvSplitSize, kvSize - n); + // Only the first num_keys columns are causally attendable; the rest + // would be masked to -inf and contribute exactly zero, so clamping + // here skips their gemm, softmax and v-multiply. This matters for the + // leading query blocks of a prefill, where num_keys is much smaller + // than the key-cache extent. Not bit-exact: shortening the reduction + // moves the vector-lane partition, so the accumulation order changes. + int64_t kvBlockSize = std::min(kvSplitSize, num_keys - n); // Calculate scale * q @ k.T - fill_stub(qk_data, static_cast(0), qSplitSize * kvSplitSize); + fill_stub(qk_data, static_cast(0), qBlockSize * kvBlockSize); const void* q_sub_matrix_data_ptr; const void* k_sub_matrix_data_ptr; @@ -1147,10 +1153,10 @@ void cpu_flash_attention( take care of this case because the loop for (int64_t n = 0; n < num_keys; n += kvSplitSize) will exit before that. */ - if (is_causal && m_start_pos <= n + kvSplitSize) { + if (is_causal && m_start_pos <= n + kvBlockSize) { // For this fn to work k_split_size > q_split_size for (int32_t row = 0; - row < qBlockSize && (m_start_pos + row < n + (kvSplitSize - 1)); + row < qBlockSize && (m_start_pos + row < n + (kvBlockSize - 1)); ++row) { // When last_col is 0, it means that the entire row is not attended // to because m_pos is smaller than n_pos. So everything in n is for diff --git a/extension/llm/custom_ops/targets.bzl b/extension/llm/custom_ops/targets.bzl index d6f251b38bd..4c2efb9adaa 100644 --- a/extension/llm/custom_ops/targets.bzl +++ b/extension/llm/custom_ops/targets.bzl @@ -53,17 +53,13 @@ def _get_quantized_moe_preproc_flags(): if runtime.is_oss: return [] if is_xplat(): - # TODO: enable KleidiAI for the runtime here by adding - # -DTORCHAO_ENABLE_KLEIDI (+ -DTORCHAO_ENABLE_ARM_I8MM=1) on arm64 and - # linking the kleidi kernel target in _get_quantized_moe_deps(). The - # runtime (op_moe.cpp) already selects the ukernel from the header, so - # kleidi headers resolve automatically once the kernels are compiled. - # Must be paired with the AoT packer emitting kleidi headers (see - # _get_quantized_moe_aot_packer_deps()). + # TODO: enable KleidiAI by adding its runtime config to op_moe.cpp, + # compiling and linking its kernels here, and pairing it with an AoT + # packer that emits Kleidi headers. return select({ "DEFAULT": [], "ovr_config//cpu:arm64": [ - "-DENABLE_QUANTIZED_MOE_FFN", + "-DEXECUTORCH_QUANTIZED_MOE_USE_TORCHAO", "-DTORCHAO_BUILD_CPU_AARCH64=1", "-DTORCHAO_ENABLE_ARM_NEON_DOT=1", ], diff --git a/extension/llm/custom_ops/test_op_moe.cpp b/extension/llm/custom_ops/test_op_moe.cpp index 30e05734bf0..81fb15185f4 100644 --- a/extension/llm/custom_ops/test_op_moe.cpp +++ b/extension/llm/custom_ops/test_op_moe.cpp @@ -43,10 +43,8 @@ TEST(OpQuantizedMoeFfnTest, RegistrationSmokeTest) { Tensor gate = tff.zeros({E, D}); Tensor expert_bias = tff.zeros({0}); - // Use empty packed buffers; the kernel will fail loudly if it tries to - // dereference them. With ENABLE_QUANTIZED_MOE_FFN unset (CI x86 build - // without torchao linkage) the kernel ET_CHECK_MSGs out before doing - // any real work, which is what we want this test to verify. + // Empty packed buffers document the expected schema; this test does not pass + // them to the kernel. Tensor packed_w1 = tfb.zeros({E, 1}); Tensor packed_w3 = tfb.zeros({E, 1}); Tensor packed_w2 = tfb.zeros({E, 1}); @@ -54,9 +52,7 @@ TEST(OpQuantizedMoeFfnTest, RegistrationSmokeTest) { Tensor out = tff.zeros({T, D}); executorch::runtime::KernelRuntimeContext ctx{}; - // We don't actually call the kernel here in the registration smoke test - // because the empty packed buffers would not be valid torchao blobs. - // Just verify the op symbol resolves at link time. + // Verify the op symbol resolves at link time. auto fn = &torch::executor::native::quantized_moe_ffn_out; EXPECT_NE(fn, nullptr); // Silence unused-variable warnings on the input tensors above; they diff --git a/extension/llm/export/BUCK b/extension/llm/export/BUCK index 7ce87cb3e73..d956bca2258 100644 --- a/extension/llm/export/BUCK +++ b/extension/llm/export/BUCK @@ -58,6 +58,24 @@ fbcode_target(_kind = runtime.python_library, ], ) +fbcode_target(_kind = runtime.python_library, + name = "model_metadata", + srcs = [ + "model_metadata.py", + ], + _is_external_target = True, + base_module = "executorch.extension.llm.export", + visibility = [ + "//executorch/backends/...", + "//executorch/examples/...", + "//executorch/extension/llm/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir:scalar_type", + ], +) + fbcode_target(_kind = runtime.python_library, name = "int4", srcs = [ diff --git a/extension/llm/export/config/llm_config.py b/extension/llm/export/config/llm_config.py index acdc9771141..3eb8b8a18f6 100644 --- a/extension/llm/export/config/llm_config.py +++ b/extension/llm/export/config/llm_config.py @@ -50,6 +50,7 @@ class ModelType(str, Enum): qwen3_5_4b = "qwen3_5_4b" phi_4_mini = "phi_4_mini" smollm2 = "smollm2" + smollm2_360m = "smollm2_360m" lfm2_350m = "lfm2_350m" lfm2_700m = "lfm2_700m" lfm2_1_2b = "lfm2_1_2b" diff --git a/extension/llm/export/model_metadata.py b/extension/llm/export/model_metadata.py new file mode 100644 index 00000000000..2057fe910df --- /dev/null +++ b/extension/llm/export/model_metadata.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python3 +# 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. + +"""Per-constant model-metadata writers for exported LLM programs. + +Exports publish their metadata as PTE constant methods; the shared typed C++ +readers in ``extension/llm/runner/model_metadata.h`` consume them. The method +names below are the single Python-side source of truth and must stay in sync +with the constants in ``extension/llm/runner/constants.h``. + +Each ``write_*`` returns the ``{name: value}`` for one constant. Producers +compose the set they publish (the MLX export's ``model_constant_methods`` builds +the full set); the backend-neutral runner test composes them too. This module +deliberately depends only on ``torch`` (and a lazy ``ScalarType`` import), so +those consumers pull in no backend. +""" + +from typing import Optional + +import torch + +# Constant-method names. Keep in sync with extension/llm/runner/constants.h. +MAX_CONTEXT_LEN_METHOD = "get_max_context_len" +MAX_SEQ_LEN_METHOD = "get_max_seq_len" +VOCAB_SIZE_METHOD = "get_vocab_size" +ACTIVATION_DTYPE_METHOD = "get_activation_dtype" +LOGITS_TO_KEEP_MODE_METHOD = "get_logits_to_keep_mode" + +# Serialized logits-to-keep modes. Keep in sync with LogitsToKeepMode in +# extension/llm/runner/model_metadata.h. +LOGITS_TO_KEEP_MODES = {"full": 0, "last": 1, "selected": 2} + + +def _require_positive(name: str, value: int) -> dict[str, int]: + if value <= 0: + raise ValueError(f"Invalid value for {name}: {value}") + return {name: value} + + +def write_max_context_len(max_context_len: int) -> dict[str, int]: + """max_context_len -> get_max_context_len (the KV-cache capacity).""" + return _require_positive(MAX_CONTEXT_LEN_METHOD, max_context_len) + + +def write_vocab_size(vocab_size: int) -> dict[str, int]: + """vocab_size -> get_vocab_size.""" + return _require_positive(VOCAB_SIZE_METHOD, vocab_size) + + +def write_max_seq_len(max_seq_len: Optional[int]) -> dict[str, int]: + """max_seq_len -> get_max_seq_len (largest single forward step); optional.""" + if max_seq_len is None: + return {} + return _require_positive(MAX_SEQ_LEN_METHOD, max_seq_len) + + +def write_activation_dtype(activation_dtype: str) -> dict[str, int]: + """activation_dtype name -> get_activation_dtype (ExecuTorch ScalarType).""" + from executorch.exir.scalar_type import ScalarType + + table = { + "fp16": ScalarType.HALF, + "fp32": ScalarType.FLOAT, + "bf16": ScalarType.BFLOAT16, + } + try: + return {ACTIVATION_DTYPE_METHOD: int(table[activation_dtype])} + except KeyError as error: + raise ValueError(f"Unsupported activation dtype: {activation_dtype}") from error + + +def write_logits_to_keep_mode(logits_to_keep: str) -> dict[str, int]: + """logits_to_keep name -> get_logits_to_keep_mode.""" + try: + return {LOGITS_TO_KEEP_MODE_METHOD: LOGITS_TO_KEEP_MODES[logits_to_keep]} + except KeyError as error: + raise ValueError( + f"Unsupported logits-to-keep mode: {logits_to_keep}" + ) from error + + +def model_vocab_size(model: torch.nn.Module) -> int: + """Return the model's actual output vocabulary width.""" + output_embeddings = model.get_output_embeddings() + if output_embeddings is None or not hasattr(output_embeddings, "weight"): + raise ValueError("Model has no output embedding weight") + vocab_size = int(output_embeddings.weight.shape[0]) + if vocab_size <= 0: + raise ValueError(f"Invalid vocabulary size: {vocab_size}") + return vocab_size diff --git a/extension/llm/export/quantizer_lib.py b/extension/llm/export/quantizer_lib.py index e4564f32360..9c900c13e63 100644 --- a/extension/llm/export/quantizer_lib.py +++ b/extension/llm/export/quantizer_lib.py @@ -109,7 +109,7 @@ def check_embedding_byte_registered(): "Need to specify shared library path to register quantized ops (and their out variants) into EXIR.\n" "Follow the following steps to build the needed lib via cmake.\n" "Then from root executorch dir do the following:\n" - "rm -rf cmake-out && mkdir cmake-out && (cd cmake-out && cmake -DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON ..) && cmake --build . -j16\n" + "rm -rf cmake-out && mkdir cmake-out && (cd cmake-out && cmake -DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON ..) && cmake --build . -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 ))\n" 'To find the location of the lib: find cmake-out -name "libquantized_ops_aot_lib*"\n' "Then specify the said library via -s +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { + +enum class LogitsToKeepMode : std::int64_t { + Full = 0, + Last = 1, + Selected = 2, +}; + +// Readers for the metadata a program publishes as constant methods. Each +// required field's reader returns Error::InvalidProgram (and logs which method) +// if it is absent or malformed (non-positive size, unknown enum). A genuinely +// optional field added later should read through detail::read_int_method +// (nullopt when absent) so old programs stay readable; there are none today. +// vocab_size is returned as published (int64); check_vocab_size() narrows it to +// int32 after cross-checking the forward output. + +namespace detail { + +// Read a named int constant method: nullopt if absent, the value if present, +// Error::InvalidProgram if it does not evaluate to a single int. +inline runtime::Result> read_int_method( + Module& module, + const char* name) { + const auto names = ET_UNWRAP(module.method_names()); + if (names.count(name) == 0) { + return std::optional{}; + } + const auto result = module.execute(name); + if (!result.ok()) { + return result.error(); + } + ET_CHECK_OR_RETURN_ERROR( + result->size() == 1 && result->at(0).isInt(), + InvalidProgram, + "metadata %s must evaluate to a single int", + name); + return std::optional{result->at(0).toInt()}; +} + +// A required int constant that must be present and positive. +inline runtime::Result read_required_positive_int( + Module& module, + const char* name) { + const auto value = ET_UNWRAP(read_int_method(module, name)); + ET_CHECK_OR_RETURN_ERROR( + value.has_value(), InvalidProgram, "metadata %s is required", name); + ET_CHECK_OR_RETURN_ERROR( + *value > 0, + InvalidProgram, + "metadata %s must be positive, got %" PRId64, + name, + *value); + return *value; +} + +} // namespace detail + +// One reader per constant: the name, its encoding, and its validation together. +// Each rejection logs which constant method was at fault. + +inline runtime::Result read_max_context_length(Module& module) { + return detail::read_required_positive_int(module, kMaxContextLen); +} + +inline runtime::Result read_vocab_size(Module& module) { + return detail::read_required_positive_int(module, kVocabSize); +} + +inline runtime::Result read_activation_dtype(Module& module) { + const auto value = + ET_UNWRAP(detail::read_int_method(module, kActivationDtype)); + ET_CHECK_OR_RETURN_ERROR( + value.has_value(), + InvalidProgram, + "metadata %s is required", + kActivationDtype); + switch (*value) { + case static_cast(aten::ScalarType::Half): + return aten::ScalarType::Half; + case static_cast(aten::ScalarType::Float): + return aten::ScalarType::Float; + case static_cast(aten::ScalarType::BFloat16): + return aten::ScalarType::BFloat16; + default: + ET_LOG( + Error, + "metadata %s has unsupported value %" PRId64, + kActivationDtype, + *value); + return runtime::Error::InvalidProgram; + } +} + +inline runtime::Result read_logits_to_keep_mode( + Module& module) { + const auto value = + ET_UNWRAP(detail::read_int_method(module, kLogitsToKeepMode)); + ET_CHECK_OR_RETURN_ERROR( + value.has_value(), + InvalidProgram, + "metadata %s is required", + kLogitsToKeepMode); + switch (*value) { + case static_cast(LogitsToKeepMode::Full): + return LogitsToKeepMode::Full; + case static_cast(LogitsToKeepMode::Last): + return LogitsToKeepMode::Last; + case static_cast(LogitsToKeepMode::Selected): + return LogitsToKeepMode::Selected; + default: + ET_LOG( + Error, + "metadata %s has unsupported value %" PRId64, + kLogitsToKeepMode, + *value); + return runtime::Error::InvalidProgram; + } +} + +inline runtime::Result read_max_seq_len(Module& module) { + const auto value = + ET_UNWRAP(detail::read_required_positive_int(module, kMaxSeqLen)); + // Consumers narrow this to int for chunked prefill, so reject a value that + // would truncate rather than letting the cast overflow. + ET_CHECK_OR_RETURN_ERROR( + value <= std::numeric_limits::max(), + InvalidProgram, + "metadata %s %" PRId64 " exceeds the maximum forward step %d", + kMaxSeqLen, + value, + std::numeric_limits::max()); + return value; +} + +// Check the published vocab size against the model's actual forward output +// width: reject a disagreement or an out-of-int32 width, and hand back the +// (now int32) value the sampler takes. +inline runtime::Result check_vocab_size( + std::int64_t published_vocab_size, + std::int64_t output_vocab_size) { + ET_CHECK_OR_RETURN_ERROR( + output_vocab_size > 0 && + output_vocab_size <= std::numeric_limits::max(), + InvalidProgram, + "forward output vocab width %" PRId64 " is out of range", + output_vocab_size); + ET_CHECK_OR_RETURN_ERROR( + published_vocab_size == output_vocab_size, + InvalidProgram, + "published %s %" PRId64 " disagrees with forward output width %" PRId64, + kVocabSize, + published_vocab_size, + output_vocab_size); + // Equal to output_vocab_size, already checked to fit int32. + return static_cast(published_vocab_size); +} + +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/runner/targets.bzl b/extension/llm/runner/targets.bzl index 9af2597b4f2..c66bec170fe 100644 --- a/extension/llm/runner/targets.bzl +++ b/extension/llm/runner/targets.bzl @@ -35,6 +35,7 @@ def define_common_targets(): runtime.cxx_library( name = "stats" + aten_suffix, exported_headers = [ + "model_metadata.h", "stats.h", "util.h", ], diff --git a/extension/llm/runner/test/CMakeLists.txt b/extension/llm/runner/test/CMakeLists.txt index 81b69c0ab9a..d43b297fa52 100644 --- a/extension/llm/runner/test/CMakeLists.txt +++ b/extension/llm/runner/test/CMakeLists.txt @@ -19,10 +19,12 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) set(_test_srcs test_generation_config.cpp + test_model_metadata.cpp test_text_llm_runner.cpp test_text_prefiller.cpp test_text_decoder_runner.cpp test_multimodal_input.cpp + test_text_stream.cpp test_util.cpp test_wav_loader.cpp ) @@ -32,9 +34,42 @@ if(APPLE) list(APPEND _test_srcs lsan_stub.cpp) endif() +set(_metadata_pte_files + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_full.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_last.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_selected.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_context.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_prefill.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_vocab.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_missing.pte" +) +add_custom_command( + OUTPUT ${_metadata_pte_files} + # Run against the installed executorch (which stages exir/_serialize/*.fbs at + # install time). Do NOT prepend the source tree to PYTHONPATH: it shadows the + # installed package with a source copy that lacks the staged .fbs, breaking + # to_executorch() under a non-editable install. + COMMAND + ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/export_model_metadata.py + --outdir ${CMAKE_CURRENT_BINARY_DIR} + WORKING_DIRECTORY ${EXECUTORCH_ROOT} + DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/export_model_metadata.py + ${EXECUTORCH_ROOT}/extension/llm/export/model_metadata.py +) +add_custom_target( + generated_model_metadata_test_files DEPENDS ${_metadata_pte_files} +) + et_cxx_test( test_runner SOURCES ${_test_srcs} EXTRA_LIBS executorch extension_llm_runner ) +add_dependencies(test_runner generated_model_metadata_test_files) +set_property( + TEST test_runner + PROPERTY + ENVIRONMENT + "ET_MODEL_METADATA_FULL_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_full.pte;ET_MODEL_METADATA_LAST_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_last.pte;ET_MODEL_METADATA_SELECTED_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_selected.pte;ET_MODEL_METADATA_INVALID_CONTEXT_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_context.pte;ET_MODEL_METADATA_INVALID_PREFILL_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_prefill.pte;ET_MODEL_METADATA_INVALID_VOCAB_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_vocab.pte;ET_MODEL_METADATA_MISSING_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_missing.pte" +) # Override sanitizer to this issue: # https://github.com/abseil/abseil-cpp/issues/841 Root issue: diff --git a/extension/llm/runner/test/export_model_metadata.py b/extension/llm/runner/test/export_model_metadata.py new file mode 100644 index 00000000000..71e0255aa6e --- /dev/null +++ b/extension/llm/runner/test/export_model_metadata.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +# 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. + +import argparse +from pathlib import Path + +import torch + +from executorch.exir import to_edge +from executorch.extension.llm.export.model_metadata import ( + write_activation_dtype, + write_logits_to_keep_mode, + write_max_context_len, + write_max_seq_len, + write_vocab_size, +) + + +class Identity(torch.nn.Module): + """Minimal model used to serialize metadata fixtures.""" + + def forward(self, value: torch.Tensor) -> torch.Tensor: + """Return the input unchanged.""" + return value + + +def all_methods(logits_to_keep: str, activation_dtype: str) -> dict[str, int]: + """Compose the full metadata set from the individual per-constant writers.""" + return { + **write_max_context_len(4096), + **write_vocab_size(128256), + **write_activation_dtype(activation_dtype), + **write_logits_to_keep_mode(logits_to_keep), + **write_max_seq_len(512), + } + + +def main() -> None: + """Generate model metadata PTE fixtures.""" + parser = argparse.ArgumentParser() + parser.add_argument("--outdir", required=True) + args = parser.parse_args() + + output_dir = Path(args.outdir) + output_dir.mkdir(parents=True, exist_ok=True) + exported = torch.export.export(Identity(), (torch.ones(1),), strict=True) + for mode, dtype in ( + ("full", "fp32"), + ("last", "fp16"), + ("selected", "bf16"), + ): + program = to_edge( + exported, + constant_methods=all_methods(mode, dtype), + ).to_executorch() + (output_dir / f"ModelMetadata_{mode}.pte").write_bytes(program.buffer) + + for name, invalid_field in ( + ("invalid_context", "get_max_context_len"), + ("invalid_prefill", "get_max_seq_len"), + ("invalid_vocab", "get_vocab_size"), + ): + methods = all_methods("full", "fp32") + methods[invalid_field] = 0 + program = to_edge(exported, constant_methods=methods).to_executorch() + (output_dir / f"ModelMetadata_{name}.pte").write_bytes(program.buffer) + + # No constant methods: exercises rejection of missing required fields. + missing_program = to_edge(exported).to_executorch() + (output_dir / "ModelMetadata_missing.pte").write_bytes(missing_program.buffer) + + +if __name__ == "__main__": + main() diff --git a/extension/llm/runner/test/test_model_metadata.cpp b/extension/llm/runner/test/test_model_metadata.cpp new file mode 100644 index 00000000000..256f915bd14 --- /dev/null +++ b/extension/llm/runner/test/test_model_metadata.cpp @@ -0,0 +1,126 @@ +/* + * 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. + */ + +#include +#include + +#include +#include + +#include + +namespace { + +using ::executorch::extension::Module; +using ::executorch::extension::llm::check_vocab_size; +using ::executorch::extension::llm::LogitsToKeepMode; +using ::executorch::extension::llm::read_activation_dtype; +using ::executorch::extension::llm::read_logits_to_keep_mode; +using ::executorch::extension::llm::read_max_context_length; +using ::executorch::extension::llm::read_max_seq_len; +using ::executorch::extension::llm::read_vocab_size; +using ::executorch::runtime::Error; + +std::unique_ptr load_fixture(const char* environment_variable) { + const char* path = std::getenv(environment_variable); + EXPECT_NE(path, nullptr); + auto module = std::make_unique(path); + EXPECT_EQ(module->load(), Error::Ok); + return module; +} + +struct ModeCase { + const char* environment_variable; + LogitsToKeepMode expected_mode; + ::executorch::aten::ScalarType expected_dtype; +}; + +class ModelMetadataTest : public ::testing::TestWithParam {}; + +TEST_P(ModelMetadataTest, ReadsPythonExportedConstants) { + auto module = load_fixture(GetParam().environment_variable); + + const auto max_context_length = read_max_context_length(*module); + ASSERT_TRUE(max_context_length.ok()); + EXPECT_EQ(*max_context_length, 4096); + + const auto max_seq_len = read_max_seq_len(*module); + ASSERT_TRUE(max_seq_len.ok()); + EXPECT_EQ(*max_seq_len, 512); + + const auto vocab_size = read_vocab_size(*module); + ASSERT_TRUE(vocab_size.ok()); + EXPECT_EQ(*vocab_size, 128256); + + const auto activation_dtype = read_activation_dtype(*module); + ASSERT_TRUE(activation_dtype.ok()); + EXPECT_EQ(*activation_dtype, GetParam().expected_dtype); + + const auto logits_to_keep_mode = read_logits_to_keep_mode(*module); + ASSERT_TRUE(logits_to_keep_mode.ok()); + EXPECT_EQ(*logits_to_keep_mode, GetParam().expected_mode); +} + +TEST(ModelMetadataTest, RejectsNonPositiveSizes) { + { + auto module = load_fixture("ET_MODEL_METADATA_INVALID_CONTEXT_PATH"); + const auto value = read_max_context_length(*module); + ASSERT_FALSE(value.ok()); + EXPECT_EQ(value.error(), Error::InvalidProgram); + } + { + auto module = load_fixture("ET_MODEL_METADATA_INVALID_PREFILL_PATH"); + const auto value = read_max_seq_len(*module); + ASSERT_FALSE(value.ok()); + EXPECT_EQ(value.error(), Error::InvalidProgram); + } + { + auto module = load_fixture("ET_MODEL_METADATA_INVALID_VOCAB_PATH"); + const auto value = read_vocab_size(*module); + ASSERT_FALSE(value.ok()); + EXPECT_EQ(value.error(), Error::InvalidProgram); + } +} + +TEST(ModelMetadataTest, ChecksVocabAgainstForwardOutput) { + auto matched = check_vocab_size(128256, 128256); + ASSERT_TRUE(matched.ok()); + EXPECT_EQ(*matched, 128256); + + auto mismatched = check_vocab_size(128256, 128000); + ASSERT_FALSE(mismatched.ok()); + EXPECT_EQ(mismatched.error(), Error::InvalidProgram); +} + +TEST(ModelMetadataTest, RejectsMissingRequiredFields) { + auto module = load_fixture("ET_MODEL_METADATA_MISSING_PATH"); + EXPECT_FALSE(read_max_context_length(*module).ok()); + EXPECT_FALSE(read_max_seq_len(*module).ok()); + EXPECT_FALSE(read_vocab_size(*module).ok()); + EXPECT_FALSE(read_activation_dtype(*module).ok()); + EXPECT_FALSE(read_logits_to_keep_mode(*module).ok()); +} + +INSTANTIATE_TEST_SUITE_P( + RoundTrip, + ModelMetadataTest, + ::testing::Values( + ModeCase{ + "ET_MODEL_METADATA_FULL_PATH", + LogitsToKeepMode::Full, + ::executorch::aten::ScalarType::Float}, + ModeCase{ + "ET_MODEL_METADATA_LAST_PATH", + LogitsToKeepMode::Last, + ::executorch::aten::ScalarType::Half}, + ModeCase{ + "ET_MODEL_METADATA_SELECTED_PATH", + LogitsToKeepMode::Selected, + ::executorch::aten::ScalarType::BFloat16})); + +} // namespace diff --git a/extension/llm/runner/test/test_text_stream.cpp b/extension/llm/runner/test/test_text_stream.cpp new file mode 100644 index 00000000000..8cddcdb88b5 --- /dev/null +++ b/extension/llm/runner/test/test_text_stream.cpp @@ -0,0 +1,318 @@ +/* + * 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. + */ + +#include + +#include +#include +#include + +#include + +using executorch::extension::llm::TextStream; +using executorch::runtime::Error; + +namespace { + +// Maps each token to the bytes it contributes. `pair_pieces` overrides the +// piece for a (previous, token) pair, so a test can observe that the preceding +// token reaches the tokenizer at all. +class FakeTokenizer : public tokenizers::Tokenizer { + public: + std::map pieces; + std::map, std::string> pair_pieces; + // Tokens the tokenizer refuses to decode. + std::vector rejected; + + tokenizers::Error load(const std::string&) override { + initialized_ = true; + return tokenizers::Error::Ok; + } + + tokenizers::Result decode( + uint64_t previous, + uint64_t token, + bool /*skip_special_tokens*/ = false) const override { + for (uint64_t bad : rejected) { + if (token == bad) { + return tokenizers::Error::Internal; + } + } + auto pair = pair_pieces.find({previous, token}); + if (pair != pair_pieces.end()) { + return pair->second; + } + auto single = pieces.find(token); + if (single != pieces.end()) { + return single->second; + } + return std::string(); + } + + tokenizers::Result> + encode(const std::string&, int8_t, int8_t) const override { + return std::vector{}; + } + tokenizers::Result id_to_piece(uint64_t) const override { + return std::string(); + } + tokenizers::Result piece_to_id(const std::string&) const override { + return uint64_t{0}; + } +}; + +// Collects what the stream emitted, both as separate pieces and joined. +struct Sink { + std::vector pieces; + std::string joined; + + void operator()(const std::string& piece) { + pieces.push_back(piece); + joined += piece; + } +}; + +// A stream writing into `sink`. +TextStream +stream_into(const FakeTokenizer& tokenizer, Sink& sink, uint64_t previous = 0) { + return TextStream( + tokenizer, [&sink](const std::string& piece) { sink(piece); }, previous); +} + +// The three bytes of U+4E16 (CJK), which a byte-level tokenizer can split. +constexpr const char kCjkByte0[] = "\xE4"; +constexpr const char kCjkByte1[] = "\xB8"; +constexpr const char kCjkByte2[] = "\x96"; +constexpr const char kCjk[] = "\xE4\xB8\x96"; + +} // namespace + +TEST(TextStreamTest, EmitsWholeCharactersImmediately) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "Hello"}, {2, " world"}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + EXPECT_EQ(stream.append(std::vector{1, 2}), Error::Ok); + EXPECT_EQ(sink.joined, "Hello world"); + EXPECT_EQ(sink.pieces.size(), 2u); + EXPECT_FALSE(stream.has_pending()); +} + +TEST(TextStreamTest, NeverEmitsAnEmptyPiece) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, ""}, {2, "x"}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(std::vector{1, 2}), Error::Ok); + EXPECT_EQ(sink.pieces, (std::vector{"x"})); +} + +// The reason this class exists: a character split across tokens must not reach +// the sink in pieces, or a consumer treating each piece as a string breaks. +TEST(TextStreamTest, HoldsBackAPartialCharacterUntilItCompletes) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, kCjkByte0}, {2, kCjkByte1}, {3, kCjkByte2}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(1u), Error::Ok); + EXPECT_TRUE(sink.pieces.empty()) << "one third of a character is not text"; + EXPECT_TRUE(stream.has_pending()); + + ASSERT_EQ(stream.append(2u), Error::Ok); + EXPECT_TRUE(sink.pieces.empty()); + + ASSERT_EQ(stream.append(3u), Error::Ok); + EXPECT_EQ(sink.joined, kCjk); + EXPECT_EQ(sink.pieces.size(), 1u) << "the character arrives whole, once"; + EXPECT_FALSE(stream.has_pending()); +} + +TEST(TextStreamTest, EmitsTheCompletePrefixAndKeepsTheRest) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, std::string("ab") + kCjkByte0}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(1u), Error::Ok); + EXPECT_EQ(sink.joined, "ab") << "the finished characters go now"; + EXPECT_TRUE(stream.has_pending()) << "the split character waits"; +} + +TEST(TextStreamTest, FlushReleasesAnUnfinishedCharacter) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, kCjkByte0}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(1u), Error::Ok); + ASSERT_TRUE(sink.pieces.empty()); + + stream.flush(); + EXPECT_EQ(sink.joined, kCjkByte0) + << "a generation that ends mid-character must not swallow the bytes"; + EXPECT_FALSE(stream.has_pending()); +} + +TEST(TextStreamTest, FlushIsIdempotentAndSilentWhenEmpty) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "done"}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(1u), Error::Ok); + stream.flush(); + stream.flush(); + EXPECT_EQ(sink.pieces, (std::vector{"done"})); +} + +// Only SentencePiece reads the preceding token, but the stream must still +// forward it and advance it, or that tokenizer would strip the wrong space. +TEST(TextStreamTest, ForwardsThePrecedingTokenAndAdvancesIt) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{7, "?"}, {8, "!"}}; + tokenizer.pair_pieces = {{{5, 7}, " seeded"}, {{7, 8}, " advanced"}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink, /*previous=*/5); + + ASSERT_EQ(stream.append(7u), Error::Ok); + EXPECT_EQ(sink.joined, " seeded") << "the constructor seed reaches decode"; + + ASSERT_EQ(stream.append(8u), Error::Ok); + EXPECT_EQ(sink.joined, " seeded advanced") + << "the previous token becomes the one just decoded"; +} + +TEST(TextStreamTest, ATokenizerErrorFailsTheStreamForGood) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "a"}, {3, "c"}}; + tokenizer.rejected = {2}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + EXPECT_EQ(stream.append(1u), Error::Ok); + EXPECT_EQ(stream.append(2u), Error::InvalidArgument); + EXPECT_TRUE(stream.failed()); + EXPECT_EQ(stream.append(3u), Error::InvalidState) + << "a failed stream must not resume and emit text out of order"; + EXPECT_EQ(sink.joined, "a"); +} + +// The batch stops where it broke rather than skipping past the bad token. +TEST(TextStreamTest, ABatchStopsAtTheTokenThatFailed) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "a"}, {3, "c"}}; + tokenizer.rejected = {2}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + EXPECT_EQ( + stream.append(std::vector{1, 2, 3}), Error::InvalidArgument); + EXPECT_EQ(sink.joined, "a") << "nothing after the failure is emitted"; + EXPECT_TRUE(stream.failed()); +} + +// A speculative executor hands back several tokens at once, so a batch must +// read the same as the tokens arriving one by one. +TEST(TextStreamTest, ABatchMatchesTokenByTokenDelivery) { + FakeTokenizer tokenizer; + tokenizer.pieces = { + {1, "He"}, {2, kCjkByte0}, {3, kCjkByte1}, {4, kCjkByte2}}; + Sink batched; + Sink one_at_a_time; + + TextStream a = stream_into(tokenizer, batched); + ASSERT_EQ(a.append(std::vector{1, 2, 3, 4}), Error::Ok); + + TextStream b = stream_into(tokenizer, one_at_a_time); + for (uint64_t token : {1u, 2u, 3u, 4u}) { + ASSERT_EQ(b.append(token), Error::Ok); + } + + EXPECT_EQ(batched.joined, one_at_a_time.joined); + EXPECT_EQ(batched.joined, std::string("He") + kCjk); +} + +TEST(TextStreamTest, AnInvalidLeadByteIsEmittedRatherThanStallingOutput) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "\xFF"}, {2, "ok"}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(std::vector{1, 2}), Error::Ok); + EXPECT_EQ( + sink.joined, + "\xFF" + "ok") + << "a byte that can never start a character must not hold up the stream"; + EXPECT_FALSE(stream.has_pending()); +} + +TEST(TextStreamTest, ToleratesAnAbsentSink) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "x"}}; + TextStream stream(tokenizer, nullptr); + + EXPECT_EQ(stream.append(1u), Error::Ok); + stream.flush(); +} + +// A four-byte codepoint takes the len == 4 branch, which the three-byte cases +// above leave untested, and is the widest split a byte-level tokenizer can +// make. +TEST(TextStreamTest, HoldsBackAFourByteCharacterUntilItCompletes) { + FakeTokenizer tokenizer; + tokenizer.pieces = { + {1, "\xF0"}, {2, "\x9F"}, {3, "\x98"}, {4, "\x80"}}; // U+1F600 + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + for (uint64_t id : {1u, 2u, 3u}) { + ASSERT_EQ(stream.append(id), Error::Ok); + EXPECT_TRUE(sink.joined.empty()) << "emitted before the character finished"; + EXPECT_TRUE(stream.has_pending()); + } + ASSERT_EQ(stream.append(4u), Error::Ok); + EXPECT_EQ(sink.joined, "\xF0\x9F\x98\x80"); + EXPECT_FALSE(stream.has_pending()); +} + +// A lead byte promises continuation bytes that never arrive. Holding them would +// stall the stream for good, so they go out as-is: the sink can see invalid +// UTF-8 here, which is the documented trade for never blocking. +TEST(TextStreamTest, AMalformedSequenceIsEmittedRatherThanHeldForever) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "\xE4\x41"}}; // 3-byte lead, then ASCII 'A' + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(1u), Error::Ok); + EXPECT_EQ(sink.joined, "\xE4\x41"); + EXPECT_FALSE(stream.has_pending()) << "a malformed tail must not be held"; +} + +// The bytes held when a decode fails are not lost: the stream is sticky-failed, +// but flush() still surrenders what it was holding. +TEST(TextStreamTest, FlushAfterAFailureStillReleasesTheHeldBytes) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, kCjkByte0}}; + tokenizer.rejected = {2}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(1u), Error::Ok); + ASSERT_TRUE(stream.has_pending()); + EXPECT_NE(stream.append(2u), Error::Ok); + EXPECT_TRUE(stream.has_pending()) << "a failure keeps what was held"; + + stream.flush(); + EXPECT_EQ(sink.joined, kCjkByte0); + EXPECT_FALSE(stream.has_pending()); +} diff --git a/extension/llm/runner/text_stream.h b/extension/llm/runner/text_stream.h new file mode 100644 index 00000000000..8a58b28ea2c --- /dev/null +++ b/extension/llm/runner/text_stream.h @@ -0,0 +1,128 @@ +/* + * 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 + +// Assembles a stream of token ids into text a caller can hand straight to a +// UI, a socket, or a JSON encoder. +// +// A byte-level tokenizer can emit a token that is only part of a character, so +// decoding each token on its own yields pieces that are not valid UTF-8 even +// though their concatenation is. Printing to a terminal survives that; +// anything that treats a piece as a standalone string does not. Holding the +// incomplete tail back until it completes is the state this exists to own, +// because a character split across tokens outlives any one delivery of them. +// +// Stop strings are a separate stage. Assemble here, then filter the text with +// stop_safe_prefix_len before it reaches the sink. + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { + +class ET_EXPERIMENTAL TextStream { + public: + // Never called with an empty string. + using Sink = std::function; + + // The tokenizer is borrowed and must outlive the stream. + // + // `previous` is the token the next one follows, normally the last token of + // the prompt. Only a SentencePiece tokenizer reads it, to drop the leading + // space of the first token after BOS; every BPETokenizerBase discards it. + // The default is for BPE, which cannot observe the value. A SentencePiece + // caller must pass the real previous token: 0 is a valid id there, so + // leaving it defaulted silently mis-handles the space on the first token. + TextStream( + const tokenizers::Tokenizer& tokenizer, + Sink on_text, + uint64_t previous = 0) + : tokenizer_(tokenizer), + on_text_(std::move(on_text)), + previous_(previous) {} + + // Emits every character the token completed. A token that only extends an + // unfinished character emits nothing and is not lost. + // + // On a tokenizer error the stream stops emitting and stays failed, rather + // than leaving the caller to guess which tokens reached the sink. + ::executorch::runtime::Error append(uint64_t token) { + if (failed_) { + return ::executorch::runtime::Error::InvalidState; + } + const tokenizers::Result piece = + tokenizer_.decode(previous_, token); + if (!piece.ok()) { + failed_ = true; + return ::executorch::runtime::Error::InvalidArgument; + } + previous_ = token; + pending_ += *piece; + emit_(utf8_complete_prefix_len(pending_)); + return ::executorch::runtime::Error::Ok; + } + + // Stops at the first token that fails. + ::executorch::runtime::Error append(const std::vector& tokens) { + for (uint64_t token : tokens) { + const ::executorch::runtime::Error error = append(token); + if (error != ::executorch::runtime::Error::Ok) { + return error; + } + } + return ::executorch::runtime::Error::Ok; + } + + // Emits whatever is held back, including a trailing character the tokens + // never finished. Call it once the generation has ended, or those bytes are + // dropped. Idempotent. + void flush() { + emit_(pending_.size()); + } + + // A character whose bytes have not all arrived yet. + bool has_pending() const { + return !pending_.empty(); + } + + bool failed() const { + return failed_; + } + + private: + void emit_(size_t length) { + if (length == 0) { + return; + } + std::string ready = pending_.substr(0, length); + pending_.erase(0, length); + if (on_text_) { + on_text_(ready); + } + } + + const tokenizers::Tokenizer& tokenizer_; + Sink on_text_; + uint64_t previous_; + std::string pending_; + bool failed_ = false; +}; + +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/runner/util.h b/extension/llm/runner/util.h index f5bcb945dbf..b45f1b51f67 100644 --- a/extension/llm/runner/util.h +++ b/extension/llm/runner/util.h @@ -87,8 +87,13 @@ ET_EXPERIMENTAL void inline safe_printf(const char* piece) { // UTF-8 multi-byte sequence. A byte-level tokenizer can emit a token that is // only part of a character (e.g. one byte of a 3-byte CJK codepoint or emoji), // so a caller streaming text must hold the incomplete tail until it completes -// rather than decode the partial bytes. An invalid lead byte counts as length 1 -// (emitted, so the caller can replace it) rather than stalling output. +// rather than decode the partial bytes. +// +// Malformed input is emitted rather than held, so a bad byte never stalls the +// stream: an invalid lead byte counts as length 1, and so does a valid lead +// whose continuation bytes are not 0x80-0xBF. Both reach the caller as single +// bytes it can replace, instead of being run together into a sequence that +// merely looks multi-byte. Only a *well-formed* truncated tail is held back. ET_EXPERIMENTAL size_t inline utf8_complete_prefix_len(const std::string& s) { size_t i = 0; const size_t n = s.size(); @@ -106,6 +111,19 @@ ET_EXPERIMENTAL size_t inline utf8_complete_prefix_len(const std::string& s) { } else { len = 1; // invalid lead byte; emit it and let the caller replace it } + // A lead byte only promises a length; the bytes after it have to agree. + // Checking the ones already present means a malformed sequence degrades to + // the invalid-lead path above rather than being emitted whole, and a + // truncated *valid* sequence is still held for the bytes that finish it. + if (len > 1) { + const size_t seen = (n - i) < len ? (n - i) : len; + for (size_t k = 1; k < seen; ++k) { + if ((static_cast(s[i + k]) & 0xC0) != 0x80) { + len = 1; + break; + } + } + } if (i + len > n) { break; // incomplete trailing sequence: hold it for more bytes } diff --git a/extension/wasm/CMakeLists.txt b/extension/wasm/CMakeLists.txt index 8ffd1801c63..5b52622140c 100644 --- a/extension/wasm/CMakeLists.txt +++ b/extension/wasm/CMakeLists.txt @@ -9,7 +9,7 @@ # cmake-format -i CMakeLists.txt # ~~~ -cmake_minimum_required(VERSION 3.29) +cmake_minimum_required(VERSION 3.26) project(executorch_wasm) diff --git a/extension/wasm/README.md b/extension/wasm/README.md index 54b1168732d..bfcfe6af96d 100644 --- a/extension/wasm/README.md +++ b/extension/wasm/README.md @@ -30,7 +30,7 @@ emcmake cmake . -DEXECUTORCH_BUILD_WASM=ON \ -Bcmake-out-wasm # Build the Wasm extension -cmake --build cmake-out-wasm --target executorch_wasm -j32 +cmake --build cmake-out-wasm --target executorch_wasm -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) ``` To reduce the binary size, you may also use the selective build options found in the [Kernel Library Selective Build guide](../../docs/source/kernel-library-selective-build.md). You may also use optimized kernels with the `EXECUTORCH_BUILD_KERNELS_OPTIMIZED` option. Portable kernels are used by default. diff --git a/extension/wasm/tokenizers/CMakeLists.txt b/extension/wasm/tokenizers/CMakeLists.txt index 03b7ea1ff6b..a1bead4bb80 100644 --- a/extension/wasm/tokenizers/CMakeLists.txt +++ b/extension/wasm/tokenizers/CMakeLists.txt @@ -9,7 +9,7 @@ # cmake-format -i CMakeLists.txt # ~~~ -cmake_minimum_required(VERSION 3.29) +cmake_minimum_required(VERSION 3.26) if(NOT CMAKE_CXX_STANDARD) set(CMAKE_CXX_STANDARD 17) diff --git a/extension/wasm/tokenizers/README.md b/extension/wasm/tokenizers/README.md index e1c48992e94..5e446e18acc 100644 --- a/extension/wasm/tokenizers/README.md +++ b/extension/wasm/tokenizers/README.md @@ -14,7 +14,7 @@ emcmake cmake . -DEXECUTORCH_BUILD_TOKENIZERS_WASM=ON \ -Bcmake-out-wasm # Build the Wasm extension -cmake --build cmake-out-wasm --target tokenizers_wasm -j32 +cmake --build cmake-out-wasm --target tokenizers_wasm -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) ``` Emscripten modules are loaded into the global `Module` object by default. This means you cannot have multiple modules in the same page. If you are also using the ExecuTorch Wasm bindings, it is recommended to use the `MODULARIZE` option to avoid conflicts. diff --git a/install_requirements.py b/install_requirements.py index b7d220179d3..3429f06dd8c 100644 --- a/install_requirements.py +++ b/install_requirements.py @@ -18,6 +18,35 @@ TORCH_URL_BASE = "https://download.pytorch.org/whl/test" TORCHAO_URL_BASE = "https://download.pytorch.org/whl/nightly" TORCHAO_NIGHTLY_VERSION = "0.18.0.dev20260729" +CU134_TORCHAO_NIGHTLY_VERSION = "0.19.0.dev20260811" +# These wheels' metadata pairs August 11 domain libraries with August 10 torch. +CU134_TORCH_PACKAGES = [ + "torch==2.14.0.dev20260810+cu134", + "torchvision==0.29.0.dev20260811+cu134", + "torchaudio==2.11.0.dev20260811+cu134", +] + + +def torchao_from_source(): + return ( + os.environ.get("EXECUTORCH_BUILD_KERNELS_TORCHAO") == "1" + or os.environ.get("TORCHAO_BUILD_EXPERIMENTAL_MPS") == "1" + ) + + +def cu134_requirements(torch_url, include_domains=False): + if not torch_url.endswith("/cu134"): + return [] + packages = list( + CU134_TORCH_PACKAGES if include_domains else CU134_TORCH_PACKAGES[:1] + ) + if not torchao_from_source(): + torchao_variant = ( + "cpu" if platform.machine().lower() in ("aarch64", "arm64") else "cu134" + ) + packages.append(f"torchao=={CU134_TORCHAO_NIGHTLY_VERSION}+{torchao_variant}") + return packages + # Since ExecuTorch often uses main-branch features of pytorch, only the nightly # pip versions will have the required features. @@ -46,6 +75,11 @@ def install_requirements(use_pytorch_nightly): # Determine the appropriate PyTorch URL based on CUDA delegate status torch_url = determine_torch_url(TORCH_URL_BASE) + cu134_packages = cu134_requirements(torch_url) + if cu134_packages: + torch_url = determine_torch_url(TORCHAO_URL_BASE) + if not use_pytorch_nightly: + cu134_packages[0] = "torch" # torchao's CUDA channel publishes x86_64 only, so asking for a CUDA build makes the pin # unsatisfiable on aarch64. Only that case is special-cased: falling back everywhere would # change which torchao a CPU x86_64 install resolves, and the CUDA build is genuinely wanted @@ -62,11 +96,11 @@ def install_requirements(use_pytorch_nightly): torchao_url = determine_torch_url(TORCHAO_URL_BASE) # pip packages needed by exir. - TORCH_PACKAGE = [ + TORCH_PACKAGE = cu134_packages or [ # Setting use_pytorch_nightly to false to test the pinned PyTorch commit. Note # that we don't need to set any version number there because they have already # been installed on CI before this step, so pip won't reinstall them - ("torch==2.13.0" if use_pytorch_nightly else "torch"), + ("torch==2.14.0" if use_pytorch_nightly else "torch"), f"torchao=={TORCHAO_NIGHTLY_VERSION}", ] @@ -91,10 +125,7 @@ def install_requirements(use_pytorch_nightly): ) LOCAL_REQUIREMENTS = [] - if ( - os.environ.get("EXECUTORCH_BUILD_KERNELS_TORCHAO") == "1" - or os.environ.get("TORCHAO_BUILD_EXPERIMENTAL_MPS") == "1" - ): + if torchao_from_source(): LOCAL_REQUIREMENTS.append("third-party/ao") if sys.platform != "win32": # TODO(larryliu0820): Setup a pypi package for this. @@ -122,6 +153,12 @@ def install_requirements(use_pytorch_nightly): # Without --no-build-isolation, setup.py can't find the torch module. "--no-build-isolation", *LOCAL_REQUIREMENTS, + *cu134_packages, + *( + ["--extra-index-url", torch_url, "--extra-index-url", torchao_url] + if cu134_packages + else [] + ), ], env=new_env, check=True, @@ -131,10 +168,22 @@ def install_requirements(use_pytorch_nightly): def install_optional_example_requirements(use_pytorch_nightly): # Determine the appropriate PyTorch URL based on CUDA delegate status torch_url = determine_torch_url(TORCH_URL_BASE) + cu134_packages = ( + cu134_requirements(torch_url, include_domains=True) + if use_pytorch_nightly + else [] + ) + if cu134_packages: + torch_url = determine_torch_url(TORCHAO_URL_BASE) + torchao_index = ( + ["--extra-index-url", f"{TORCHAO_URL_BASE}/cpu"] + if cu134_packages and platform.machine().lower() in ("aarch64", "arm64") + else [] + ) print("Installing torch domain libraries") - DOMAIN_LIBRARIES = [ - ("torchvision==0.28.0" if use_pytorch_nightly else "torchvision"), + DOMAIN_LIBRARIES = cu134_packages or [ + ("torchvision==0.29.0" if use_pytorch_nightly else "torchvision"), ("torchaudio==2.11.0" if use_pytorch_nightly else "torchaudio"), ] # Then install domain libraries @@ -147,6 +196,7 @@ def install_optional_example_requirements(use_pytorch_nightly): *DOMAIN_LIBRARIES, "--extra-index-url", torch_url, + *torchao_index, ], check=True, ) @@ -160,8 +210,10 @@ def install_optional_example_requirements(use_pytorch_nightly): "install", "-r", "requirements-examples.txt", + *cu134_packages, "--extra-index-url", torch_url, + *torchao_index, "--upgrade-strategy", "only-if-needed", ], diff --git a/install_utils.py b/install_utils.py index 276535cf38f..4e1397413b3 100644 --- a/install_utils.py +++ b/install_utils.py @@ -21,6 +21,7 @@ (12, 6), (13, 0), (13, 2), + (13, 4), ) diff --git a/kernels/aten/cpu/util/targets.bzl b/kernels/aten/cpu/util/targets.bzl index 983391d3613..411203c2b9a 100644 --- a/kernels/aten/cpu/util/targets.bzl +++ b/kernels/aten/cpu/util/targets.bzl @@ -15,11 +15,13 @@ def define_common_targets(): "copy_ops_util.h", ], compiler_flags = select({ - "DEFAULT": ["-Wno-missing-prototypes"], + "DEFAULT": select({ + "DEFAULT": ["-Wno-missing-prototypes"], + # GCC's C++ frontend rejects this C-only flag under -Werror. + # Nested under DEFAULT so windows and gcc can't both match. + "ovr_config//compiler:gcc": [], + }), "ovr_config//os:windows": [], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses - # this branch via runtime.is_oss. - "ovr_config//os:zephyr": [], }) if not runtime.is_oss else select({ "DEFAULT": ["-Wno-missing-prototypes"], "ovr_config//os:windows": [], diff --git a/kernels/optimized/cpu/op_bmm.cpp b/kernels/optimized/cpu/op_bmm.cpp index 171f14de399..0dd7bc40b3e 100644 --- a/kernels/optimized/cpu/op_bmm.cpp +++ b/kernels/optimized/cpu/op_bmm.cpp @@ -150,6 +150,11 @@ Tensor& opt_bmm_out( ET_KERNEL_CHECK( ctx, check_bmm_out_args(self, mat2, out), InvalidArgument, out); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dim_order(self, mat2, out), InvalidArgument, out); + + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(self), InvalidArgument, out); + static constexpr auto name = "bmm.out"; auto self_type = self.scalar_type(); diff --git a/kernels/optimized/cpu/op_mm.cpp b/kernels/optimized/cpu/op_mm.cpp index 53385a40dff..eefaf3131b7 100644 --- a/kernels/optimized/cpu/op_mm.cpp +++ b/kernels/optimized/cpu/op_mm.cpp @@ -34,6 +34,11 @@ Tensor& opt_mm_out( InvalidArgument, out); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dim_order(in, mat2, out), InvalidArgument, out); + + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out); + if (out.numel() == 0) { return out; } diff --git a/kernels/optimized/cpu/op_native_layer_norm.cpp b/kernels/optimized/cpu/op_native_layer_norm.cpp index 5fac9faf25e..66e2971087d 100644 --- a/kernels/optimized/cpu/op_native_layer_norm.cpp +++ b/kernels/optimized/cpu/op_native_layer_norm.cpp @@ -149,6 +149,32 @@ std::tuple opt_native_layer_norm_out( InvalidArgument, ret_val); + // Only support default dim order for now. + ET_KERNEL_CHECK( + ctx, tensor_is_default_dim_order(input), InvalidArgument, ret_val); + + ET_KERNEL_CHECK( + ctx, + tensors_have_same_dim_order(input, out, mean_out, rstd_out), + InvalidArgument, + ret_val); + + if (weight.has_value()) { + ET_KERNEL_CHECK( + ctx, + tensors_have_same_dim_order(input, weight.value()), + InvalidArgument, + ret_val); + } + + if (bias.has_value()) { + ET_KERNEL_CHECK( + ctx, + tensors_have_same_dim_order(input, bias.value()), + InvalidArgument, + ret_val); + } + Tensor::SizesType mean_rstd_sizes[kTensorDimensionLimit]; size_t mean_rstd_ndim = 0; get_layer_norm_out_target_size( diff --git a/kernels/optimized/cpu/op_to_copy.cpp b/kernels/optimized/cpu/op_to_copy.cpp new file mode 100644 index 00000000000..2720f46d042 --- /dev/null +++ b/kernels/optimized/cpu/op_to_copy.cpp @@ -0,0 +1,184 @@ +/* + * 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. + */ + +#include +#include +#include + +#if defined(__aarch64__) +#include + +#include +#endif + +#include +#include + +namespace torch { +namespace executor { +namespace native { + +using BFloat16 = executorch::aten::BFloat16; +using MemoryFormat = executorch::aten::MemoryFormat; +using ScalarType = executorch::aten::ScalarType; +using Tensor = executorch::aten::Tensor; + +Tensor& to_copy_out( + KernelRuntimeContext& ctx, + const Tensor& self, + bool non_blocking, + std::optional memory_format, + Tensor& out); + +namespace { + +#if defined(__aarch64__) +static_assert(sizeof(BFloat16) == sizeof(uint16_t)); +static_assert(std::is_trivially_copyable_v); + +void float_to_bfloat16_range( + const float* const input, + BFloat16* const output, + const int64_t begin, + const int64_t end) { + constexpr int64_t kVectorWidth = 8; + // Integer rounding preserves subnormals regardless of FPCR, like BFloat16. + const uint32x4_t magnitude_mask = vdupq_n_u32(0x7FFFFFFF); + const uint32x4_t infinity = vdupq_n_u32(0x7F800000); + const uint32x4_t mantissa_lsb_mask = vdupq_n_u32(1); + const uint32x4_t rounding_bias = vdupq_n_u32(0x7FFF); + const uint32x4_t canonical_nan = vdupq_n_u32(0x7FC00000); + + int64_t i = begin; +#if defined(__clang__) +#pragma unroll 4 +#elif defined(__GNUC__) +#pragma GCC unroll 4 +#endif + for (; i + kVectorWidth <= end; i += kVectorWidth) { + const uint32x4_t low_bits = vreinterpretq_u32_f32(vld1q_f32(input + i)); + const uint32x4_t high_bits = + vreinterpretq_u32_f32(vld1q_f32(input + i + 4)); + + const auto round_and_canonicalize = [&](const uint32x4_t bits) { + const uint32x4_t mantissa_lsb = + vandq_u32(vshrq_n_u32(bits, 16), mantissa_lsb_mask); + const uint32x4_t rounded = + vaddq_u32(bits, vaddq_u32(rounding_bias, mantissa_lsb)); + const uint32x4_t is_nan = + vcgtq_u32(vandq_u32(bits, magnitude_mask), infinity); + return vbslq_u32(is_nan, canonical_nan, rounded); + }; + + const uint16x4_t low = vshrn_n_u32(round_and_canonicalize(low_bits), 16); + const uint16x8_t result = + vshrn_high_n_u32(low, round_and_canonicalize(high_bits), 16); + std::memcpy(output + i, &result, sizeof(result)); + } + + for (; i < end; ++i) { + output[i] = static_cast(input[i]); + } +} + +void bfloat16_to_float_range( + const BFloat16* const input, + float* const output, + const int64_t begin, + const int64_t end) { + constexpr int64_t kVectorWidth = 8; + + int64_t i = begin; +#if defined(__clang__) +#pragma unroll 4 +#elif defined(__GNUC__) +#pragma GCC unroll 4 +#endif + for (; i + kVectorWidth <= end; i += kVectorWidth) { + uint16x8_t input_bits; + // Avoid aliasing BFloat16 storage as a NEON vector. + std::memcpy(&input_bits, input + i, sizeof(input_bits)); + const uint32x4_t low_bits = vshll_n_u16(vget_low_u16(input_bits), 16); + const uint32x4_t high_bits = vshll_high_n_u16(input_bits, 16); + vst1q_f32(output + i, vreinterpretq_f32_u32(low_bits)); + vst1q_f32(output + i + 4, vreinterpretq_f32_u32(high_bits)); + } + + for (; i < end; ++i) { + output[i] = static_cast(input[i]); + } +} + +template +bool convert_contiguous(const Tensor& self, Tensor& out) { + const auto numel = self.numel(); + if (numel == 0) { + return true; + } + + const auto* const input = self.const_data_ptr(); + auto* const output = out.mutable_data_ptr(); + const auto convert_range = [&](const auto begin, const auto end) { + if constexpr (std::is_same_v) { + float_to_bfloat16_range(input, output, begin, end); + } else { + bfloat16_to_float_range(input, output, begin, end); + } + }; + + if (numel > ::executorch::extension::internal::GRAIN_SIZE) { + return ::executorch::extension::parallel_for( + 0, numel, ::executorch::extension::internal::GRAIN_SIZE, convert_range); + } + convert_range(0, numel); + return true; +} +#endif + +} // namespace + +Tensor& opt_to_copy_out( + KernelRuntimeContext& ctx, + const Tensor& self, + bool non_blocking, + std::optional memory_format, + Tensor& out) { +#if defined(__aarch64__) + const bool float_to_bfloat16 = self.scalar_type() == ScalarType::Float && + out.scalar_type() == ScalarType::BFloat16; + const bool bfloat16_to_float = self.scalar_type() == ScalarType::BFloat16 && + out.scalar_type() == ScalarType::Float; + const bool supported_memory_format = !memory_format.has_value() || + memory_format.value() == MemoryFormat::Contiguous; + const bool can_use_optimized_kernel = + (float_to_bfloat16 || bfloat16_to_float) && !non_blocking && + supported_memory_format && tensor_is_default_dim_order(self) && + tensor_is_default_dim_order(out); + if (can_use_optimized_kernel) { + ET_KERNEL_CHECK( + ctx, + resize_tensor(out, self.sizes()) == Error::Ok, + InvalidArgument, + out); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dim_order(self, out), InvalidArgument, out); + + const bool success = float_to_bfloat16 + ? convert_contiguous(self, out) + : convert_contiguous(self, out); + ET_KERNEL_CHECK_MSG(ctx, success, Internal, out, "parallel_for failed"); + return out; + } +#endif + + return to_copy_out(ctx, self, non_blocking, memory_format, out); +} + +} // namespace native +} // namespace executor +} // namespace torch diff --git a/kernels/optimized/lib_defs.bzl b/kernels/optimized/lib_defs.bzl index 2ea329a8baa..29d27d97b28 100644 --- a/kernels/optimized/lib_defs.bzl +++ b/kernels/optimized/lib_defs.bzl @@ -25,16 +25,16 @@ def get_vec_preprocessor_flags(): # various ovr_configs are not available in oss preprocessor_flags = select({ "ovr_config//os:linux-x86_64": [ - "-DET_BUILD_ARM_VEC256_WITH_SLEEF", + "-DAT_BUILD_ARM_VEC256_WITH_SLEEF", ] if not runtime.is_oss else [], "ovr_config//os:iphoneos-arm64": [ - "-DET_BUILD_ARM_VEC256_WITH_SLEEF", + "-DAT_BUILD_ARM_VEC256_WITH_SLEEF", ] if not runtime.is_oss else [], "ovr_config//os:macos-arm64": [ - "-DET_BUILD_ARM_VEC256_WITH_SLEEF", + "-DAT_BUILD_ARM_VEC256_WITH_SLEEF", ] if not runtime.is_oss else [], "ovr_config//os:android-arm64": [ - "-DET_BUILD_ARM_VEC256_WITH_SLEEF", + "-DAT_BUILD_ARM_VEC256_WITH_SLEEF", ] if not runtime.is_oss else [], "DEFAULT": [], }) diff --git a/kernels/optimized/optimized.yaml b/kernels/optimized/optimized.yaml index 5a001afc7a0..1827411ff57 100644 --- a/kernels/optimized/optimized.yaml +++ b/kernels/optimized/optimized.yaml @@ -17,6 +17,11 @@ - arg_meta: null kernel_name: torch::executor::opt_log_softmax_out +- op: _to_copy.out + kernels: + - arg_meta: null + kernel_name: torch::executor::opt_to_copy_out + - op: add.out kernels: - arg_meta: null diff --git a/kernels/portable/cpu/op_cat.cpp b/kernels/portable/cpu/op_cat.cpp index ab15d5249df..5d5d930e29c 100644 --- a/kernels/portable/cpu/op_cat.cpp +++ b/kernels/portable/cpu/op_cat.cpp @@ -28,6 +28,10 @@ Tensor& cat_out( ET_KERNEL_CHECK(ctx, check_cat_args(tensors, dim, out), InvalidArgument, out); + // check_cat_args already requires every input to share out's dim order, so + // checking out is enough to force all of them to the default. + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(out), InvalidArgument, out); + Tensor::SizesType expected_out_size[kTensorDimensionLimit]; size_t expected_out_dim = 0; get_cat_out_target_size(tensors, dim, expected_out_size, &expected_out_dim); diff --git a/kernels/portable/cpu/op_constant_pad_nd.cpp b/kernels/portable/cpu/op_constant_pad_nd.cpp index 0f287a5ac53..0947a14dc2d 100644 --- a/kernels/portable/cpu/op_constant_pad_nd.cpp +++ b/kernels/portable/cpu/op_constant_pad_nd.cpp @@ -252,6 +252,8 @@ Tensor& constant_pad_nd_out( ET_KERNEL_CHECK( ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out); + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out); + // resize out tensor for dynamic shapes ET_KERNEL_CHECK_MSG( ctx, diff --git a/kernels/portable/cpu/op_cumsum.cpp b/kernels/portable/cpu/op_cumsum.cpp index 5023be7b694..3b7abcbed63 100644 --- a/kernels/portable/cpu/op_cumsum.cpp +++ b/kernels/portable/cpu/op_cumsum.cpp @@ -103,6 +103,8 @@ Tensor& cumsum_out( ET_KERNEL_CHECK( ctx, tensors_have_same_dim_order(self, out), InvalidArgument, out); + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(self), InvalidArgument, out); + ET_KERNEL_CHECK( ctx, resize_tensor(out, self.sizes()) == Error::Ok, InvalidArgument, out); diff --git a/kernels/portable/cpu/op_fft_r2c.cpp b/kernels/portable/cpu/op_fft_r2c.cpp new file mode 100644 index 00000000000..ddd5593ef5c --- /dev/null +++ b/kernels/portable/cpu/op_fft_r2c.cpp @@ -0,0 +1,307 @@ +/* + * 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. + */ + +#include +#include + +#include +#include +#include +#include + +namespace torch::executor::native { + +namespace { + +constexpr double kTwoPi = 6.283185307179586476925286766559; + +// A complex-to-complex pass has to read a whole line before it can overwrite +// it. Lines up to this length go through a stack buffer, which keeps the buffer +// at 2 KB for double; longer ones ask the runtime for temporary memory. Only +// multi-dimensional transforms reach this path at all: the single-dimension +// case, which is what torch.fft.rfft lowers to, needs no line buffer. +constexpr size_t kStackLineLimit = 128; + +// Mirrors ATen's fft_norm_mode (ATen/native/SpectralOpsUtils.h), which is how +// the normalization argument is encoded. +enum class fft_norm_mode { + none, // No normalization + by_root_n, // Divide by sqrt(signal_size) + by_n, // Divide by signal_size +}; + +template +std::optional compute_fct( + KernelRuntimeContext& ctx, + const Tensor& t, + IntArrayRef dim, + int64_t normalization) { + constexpr auto one = static_cast(1); + const auto mode = static_cast(normalization); + if (mode == fft_norm_mode::none) { + return one; + } + int64_t n = 1; + for (auto idx : dim) { + n *= t.sizes()[idx]; + } + switch (mode) { + case fft_norm_mode::none: + return one; + case fft_norm_mode::by_n: + return one / static_cast(n); + case fft_norm_mode::by_root_n: + return one / std::sqrt(static_cast(n)); + } + ET_KERNEL_CHECK_MSG( + ctx, + false, + InvalidArgument, + std::nullopt, + "Unsupported normalization type: %" PRId64, + normalization); +} + +// cos and sin of -2*pi*idx/n. +// +// The quarter turns are returned exactly rather than through cos/sin, so that a +// real input's Nyquist bin comes out with a zero imaginary part instead of +// rounding noise on the order of 1e-16. Reducing idx modulo n first also keeps +// the angle inside one period, which matters for accuracy once k * j is large. +void twiddle(size_t idx, size_t n, double& cos_out, double& sin_out) { + idx %= n; + if (idx == 0) { + cos_out = 1.0; + sin_out = 0.0; + } else if (2 * idx == n) { + cos_out = -1.0; + sin_out = 0.0; + } else if (4 * idx == n) { + cos_out = 0.0; + sin_out = -1.0; + } else if (4 * idx == 3 * n) { + cos_out = 0.0; + sin_out = 1.0; + } else { + const double angle = + -kTwoPi * static_cast(idx) / static_cast(n); + cos_out = std::cos(angle); + sin_out = std::sin(angle); + } +} + +// Offset of the start of the line_index'th line along `axis`, for a tensor with +// the given sizes and strides. Lines are enumerated over every dimension except +// `axis`, so the same index names the same logical line in two tensors that +// agree on all dimensions but that one. +size_t line_offset( + size_t line_index, + ArrayRef sizes, + ArrayRef strides, + size_t axis) { + size_t offset = 0; + for (size_t d = sizes.size(); d-- > 0;) { + if (d == axis) { + continue; + } + const size_t size = static_cast(sizes[d]); + offset += (line_index % size) * static_cast(strides[d]); + line_index /= size; + } + return offset; +} + +// Forward real-to-complex DFT along `axis`, writing the onesided output. +// The normalization factor is folded in here: every later pass is linear, so +// applying it once at the front scales the whole transform. +template +void dft_r2c_axis(const Tensor& in, Tensor& out, size_t axis, T fct) { + using C = executorch::runtime::etensor::complex; + const T* const in_data = in.const_data_ptr(); + C* const out_data = out.mutable_data_ptr(); + + const size_t n = static_cast(in.size(axis)); + const size_t n_out = static_cast(out.size(axis)); + const size_t in_stride = static_cast(in.strides()[axis]); + const size_t out_stride = static_cast(out.strides()[axis]); + const size_t num_lines = n == 0 ? 0 : static_cast(in.numel()) / n; + + for (size_t line = 0; line < num_lines; ++line) { + const size_t in_off = line_offset(line, in.sizes(), in.strides(), axis); + const size_t out_off = line_offset(line, out.sizes(), out.strides(), axis); + for (size_t k = 0; k < n_out; ++k) { + double real = 0; + double imag = 0; + for (size_t j = 0; j < n; ++j) { + double c = 0; + double s = 0; + twiddle(k * j, n, c, s); + const double x = static_cast(in_data[in_off + j * in_stride]); + real += x * c; + imag += x * s; + } + out_data[out_off + k * out_stride] = + C{static_cast(real * static_cast(fct)), + static_cast(imag * static_cast(fct))}; + } + } +} + +// In-place forward complex-to-complex DFT along `axis`. `scratch` must hold at +// least out.size(axis) elements. +template +void dft_c2c_axis_(Tensor& out, size_t axis, void* scratch) { + using C = executorch::runtime::etensor::complex; + C* const out_data = out.mutable_data_ptr(); + C* const line_buf = static_cast(scratch); + + const size_t n = static_cast(out.size(axis)); + const size_t stride = static_cast(out.strides()[axis]); + const size_t num_lines = n == 0 ? 0 : static_cast(out.numel()) / n; + + for (size_t line = 0; line < num_lines; ++line) { + const size_t off = line_offset(line, out.sizes(), out.strides(), axis); + for (size_t j = 0; j < n; ++j) { + line_buf[j] = out_data[off + j * stride]; + } + for (size_t k = 0; k < n; ++k) { + double real = 0; + double imag = 0; + for (size_t j = 0; j < n; ++j) { + double c = 0; + double s = 0; + twiddle(k * j, n, c, s); + const double xr = static_cast(line_buf[j].real_); + const double xi = static_cast(line_buf[j].imag_); + real += xr * c - xi * s; + imag += xr * s + xi * c; + } + out_data[off + k * stride] = + C{static_cast(real), static_cast(imag)}; + } + } +} + +} // namespace + +// Reference discrete Fourier transform. +// +// This is a direct O(n^2) evaluation of the transform sum, not a fast Fourier +// transform. kernels/optimized provides a pocketfft-backed _fft_r2c.out that is +// asymptotically faster; this exists so that a graph containing _fft_r2c can be +// run by a build that only has the portable kernels, rather than failing to +// load with OperatorMissing. Audio front-ends that transform a few hundred +// points per frame are the intended case. +Tensor& _fft_r2c_out( + KernelRuntimeContext& ctx, + const Tensor& in, + IntArrayRef dim, + int64_t normalization, + bool onesided, + Tensor& out) { + auto in_sizes = in.sizes(); + ET_KERNEL_CHECK( + ctx, + static_cast(in.dim()) <= kTensorDimensionLimit, + InvalidArgument, + out); + ET_KERNEL_CHECK(ctx, !dim.empty(), InvalidArgument, out); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out); + + ET_KERNEL_CHECK_MSG( + ctx, + onesided, + InvalidArgument, + out, + "onesided=False is not supported yet in _fft_r2c"); + + ET_KERNEL_CHECK_MSG( + ctx, + out.scalar_type() == executorch::runtime::toComplexType(in.scalar_type()), + InvalidArgument, + out, + "the output type for _fft_r2c must be the Complex type corresponding to the input type"); + + for (auto d : dim) { + ET_KERNEL_CHECK_MSG( + ctx, + d >= 0 && d < in.dim(), + InvalidArgument, + out, + "dims must be in bounds (got %" PRId64 ")", + d); + } + + std::array out_sizes_storage; + executorch::runtime::Span out_sizes( + out_sizes_storage.data(), in_sizes.size()); + std::copy(in_sizes.begin(), in_sizes.end(), out_sizes.begin()); + out_sizes[dim.back()] = out_sizes[dim.back()] / 2 + 1; + + ET_KERNEL_CHECK_MSG( + ctx, + resize_tensor( + out, + executorch::runtime::ArrayRef( + out_sizes.data(), out_sizes.size())) == Error::Ok, + InvalidArgument, + out, + "Failed to resize output tensor (last dim %d).", + out_sizes[dim.back()]); + + // NOTE: as of this writing, upstream PyTorch only supports float/double, so + // we follow suit. + ET_SWITCH_FLOAT_TYPES(in.scalar_type(), ctx, "_fft_r2c.out", CTYPE_IN, [&] { + auto fct = compute_fct(ctx, in, dim, normalization); + if (!fct) { + // Check failed, just bail out of the lambda. + return; + } + + // The real transform runs along the last requested dimension, which is the + // one that is halved; the remaining dimensions are complex transforms of + // the result, matching pocketfft's multi-axis r2c. + const size_t real_axis = static_cast(dim.back()); + dft_r2c_axis(in, out, real_axis, *fct); + + if (dim.size() == 1) { + return; + } + + using Complex = executorch::runtime::etensor::complex; + size_t max_line = 0; + for (size_t i = 0; i + 1 < dim.size(); ++i) { + max_line = std::max(max_line, static_cast(out.size(dim[i]))); + } + + std::array stack_buf; + void* line_buf = stack_buf.data(); + if (max_line > kStackLineLimit) { + Result scratch = ctx.allocate_temp(max_line * sizeof(Complex)); + ET_KERNEL_CHECK_MSG( + ctx, + scratch.ok(), + MemoryAllocationFailed, + , + "_fft_r2c needs %zu bytes of temporary memory to transform a " + "dimension of length %zu, but no temp allocator is available", + max_line * sizeof(Complex), + max_line); + line_buf = scratch.get(); + } + + for (size_t i = 0; i + 1 < dim.size(); ++i) { + dft_c2c_axis_(out, static_cast(dim[i]), line_buf); + } + }); + + return out; +} + +} // namespace torch::executor::native diff --git a/kernels/portable/cpu/op_pixel_unshuffle.cpp b/kernels/portable/cpu/op_pixel_unshuffle.cpp index 68d7bbbc27a..c6b69a56a04 100644 --- a/kernels/portable/cpu/op_pixel_unshuffle.cpp +++ b/kernels/portable/cpu/op_pixel_unshuffle.cpp @@ -81,6 +81,11 @@ Tensor& pixel_unshuffle_out( InvalidArgument, out); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out); + + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out); + // @lint-ignore CLANGTIDY facebook-hte-CArray Tensor::SizesType expected_out_size[kTensorDimensionLimit]; size_t expected_out_dim = 0; diff --git a/kernels/portable/cpu/op_slice_scatter.cpp b/kernels/portable/cpu/op_slice_scatter.cpp index 29c4ff7ab90..5a59df8d25b 100644 --- a/kernels/portable/cpu/op_slice_scatter.cpp +++ b/kernels/portable/cpu/op_slice_scatter.cpp @@ -42,7 +42,10 @@ Tensor& slice_scatter_out( out); ET_KERNEL_CHECK( - ctx, tensors_have_same_dim_order(input, out), InvalidArgument, out); + ctx, tensors_have_same_dim_order(input, src, out), InvalidArgument, out); + + ET_KERNEL_CHECK( + ctx, tensor_is_default_dim_order(input), InvalidArgument, out); if (input.numel() == 0) { return out; diff --git a/kernels/portable/cpu/op_split_copy.cpp b/kernels/portable/cpu/op_split_copy.cpp index fdc89727897..0f97dc76345 100644 --- a/kernels/portable/cpu/op_split_copy.cpp +++ b/kernels/portable/cpu/op_split_copy.cpp @@ -49,6 +49,8 @@ void split_copy_Tensor_out( for (size_t i = 0; i < out.size(); ++i) { ET_KERNEL_CHECK( ctx, tensors_have_same_dim_order(input, out[i]), InvalidArgument, ); + ET_KERNEL_CHECK( + ctx, tensor_is_default_dim_order(out[i]), InvalidArgument, ); } const size_t leading_dims = getLeadingDims(input, dim); diff --git a/kernels/portable/cpu/op_split_with_sizes_copy.cpp b/kernels/portable/cpu/op_split_with_sizes_copy.cpp index c99a7fb6815..0353e048b9e 100644 --- a/kernels/portable/cpu/op_split_with_sizes_copy.cpp +++ b/kernels/portable/cpu/op_split_with_sizes_copy.cpp @@ -43,6 +43,8 @@ void split_with_sizes_copy_out( for (const auto i : c10::irange(out.size())) { ET_KERNEL_CHECK( ctx, tensors_have_same_dim_order(in, out[i]), InvalidArgument, ); + ET_KERNEL_CHECK( + ctx, tensor_is_default_dim_order(out[i]), InvalidArgument, ); } // If out is empty, then nothing needs to be done after checking the args. diff --git a/kernels/portable/cpu/op_topk.cpp b/kernels/portable/cpu/op_topk.cpp index 3082bc94662..7bda44fccd6 100644 --- a/kernels/portable/cpu/op_topk.cpp +++ b/kernels/portable/cpu/op_topk.cpp @@ -170,6 +170,14 @@ std::tuple topk_values( ET_KERNEL_CHECK( ctx, check_topk_args(in, k, dim, values, indices), InvalidArgument, out); + ET_KERNEL_CHECK( + ctx, + tensors_have_same_dim_order(in, values, indices), + InvalidArgument, + out); + + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out); + if (dim < 0) { dim += nonzero_dim(in); } diff --git a/kernels/portable/cpu/pattern/targets.bzl b/kernels/portable/cpu/pattern/targets.bzl index 10159c7b540..47cd80ddf1b 100644 --- a/kernels/portable/cpu/pattern/targets.bzl +++ b/kernels/portable/cpu/pattern/targets.bzl @@ -56,8 +56,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], exported_deps = [ "//executorch/kernels/portable/cpu/util:broadcast_util", diff --git a/kernels/portable/cpu/util/targets.bzl b/kernels/portable/cpu/util/targets.bzl index 99f1f3adce3..fe30188e3a8 100644 --- a/kernels/portable/cpu/util/targets.bzl +++ b/kernels/portable/cpu/util/targets.bzl @@ -49,8 +49,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ "//executorch/runtime/core/exec_aten/util:tensor_shape_to_c_string", @@ -103,8 +103,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ "//executorch/runtime/kernel:kernel_includes", @@ -119,8 +119,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], exported_deps = [ ":broadcast_indexes_range", @@ -155,8 +155,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ ":broadcast_util", @@ -174,8 +174,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], exported_deps = [ ":broadcast_util", @@ -194,8 +194,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ "//executorch/runtime/kernel:kernel_includes", @@ -214,8 +214,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ "//executorch/runtime/kernel:kernel_includes", @@ -231,8 +231,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ ":broadcast_util", @@ -252,8 +252,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ "//executorch/runtime/kernel:kernel_includes", @@ -269,8 +269,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ "//executorch/runtime/kernel:kernel_includes", diff --git a/kernels/portable/functions.yaml b/kernels/portable/functions.yaml index ecf62ee3606..61f32677b99 100644 --- a/kernels/portable/functions.yaml +++ b/kernels/portable/functions.yaml @@ -32,6 +32,11 @@ - arg_meta: null kernel_name: torch::executor::_conj_physical_out +- op: _fft_r2c.out + kernels: + - arg_meta: null + kernel_name: torch::executor::_fft_r2c_out + - op: _log_softmax.out kernels: - arg_meta: null diff --git a/kernels/test/CMakeLists.txt b/kernels/test/CMakeLists.txt index da9fbc2f55b..fe70f09f00d 100644 --- a/kernels/test/CMakeLists.txt +++ b/kernels/test/CMakeLists.txt @@ -207,6 +207,7 @@ set(all_test_sources "op_exp_test.cpp" "op_expand_copy_test.cpp" "op_expm1_test.cpp" + "op_fft_r2c_test.cpp" "op_fill_test.cpp" "op_flip_test.cpp" "op_floor_divide_test.cpp" @@ -261,6 +262,7 @@ set(all_test_sources "op_pdist_forward_test.cpp" "op_permute_copy_test.cpp" "op_pixel_shuffle_test.cpp" + "op_pixel_unshuffle_test.cpp" "op_prod_test.cpp" "op_rand_test.cpp" "op_randn_test.cpp" diff --git a/kernels/test/ScalarOverflowTestMacros.h b/kernels/test/ScalarOverflowTestMacros.h index 46a2425b0fa..6567ad71564 100644 --- a/kernels/test/ScalarOverflowTestMacros.h +++ b/kernels/test/ScalarOverflowTestMacros.h @@ -11,28 +11,36 @@ // Macro to generate scalar overflow test cases for a given test suite. // The test suite must have a method called expect_bad_scalar_value_dies // that takes a template parameter for ScalarType and a Scalar value. -#define GENERATE_SCALAR_OVERFLOW_TESTS(TEST_SUITE_NAME) \ - TEST_F(TEST_SUITE_NAME, ByteTensorTooLargeScalarDies) { \ - /* Cannot be represented by a uint8_t. */ \ - expect_bad_scalar_value_dies(256); \ - } \ - \ - TEST_F(TEST_SUITE_NAME, CharTensorTooSmallScalarDies) { \ - /* Cannot be represented by a int8_t. */ \ - expect_bad_scalar_value_dies(-129); \ - } \ - \ - TEST_F(TEST_SUITE_NAME, ShortTensorTooLargeScalarDies) { \ - /* Cannot be represented by a int16_t. */ \ - expect_bad_scalar_value_dies(32768); \ - } \ - \ - TEST_F(TEST_SUITE_NAME, FloatTensorTooSmallScalarDies) { \ - /* Cannot be represented by a float. */ \ - expect_bad_scalar_value_dies(-3.41e+38); \ - } \ - \ - TEST_F(TEST_SUITE_NAME, FloatTensorTooLargeScalarDies) { \ - /* Cannot be represented by a float. */ \ - expect_bad_scalar_value_dies(3.41e+38); \ +#define GENERATE_SCALAR_OVERFLOW_TESTS(TEST_SUITE_NAME) \ + TEST_F(TEST_SUITE_NAME, ByteTensorTooLargeScalarDies) { \ + /* Cannot be represented by a uint8_t. */ \ + expect_bad_scalar_value_dies(256); \ + } \ + \ + TEST_F(TEST_SUITE_NAME, CharTensorTooSmallScalarDies) { \ + /* Cannot be represented by a int8_t. */ \ + expect_bad_scalar_value_dies(-129); \ + } \ + \ + TEST_F(TEST_SUITE_NAME, ShortTensorTooLargeScalarDies) { \ + /* Cannot be represented by a int16_t. */ \ + expect_bad_scalar_value_dies(32768); \ + } \ + \ + TEST_F(TEST_SUITE_NAME, FloatTensorTooSmallScalarDies) { \ + /* Cannot be represented by a float. */ \ + expect_bad_scalar_value_dies(-3.41e+38); \ + } \ + \ + TEST_F(TEST_SUITE_NAME, FloatTensorTooLargeScalarDies) { \ + /* Cannot be represented by a float. */ \ + expect_bad_scalar_value_dies(3.41e+38); \ + } \ + \ + TEST_F(TEST_SUITE_NAME, LongTensorTooLargeScalarDies) { \ + /* 2^63 is one past the largest int64_t, so converting it is undefined \ + * unless the range check rejects it first. The add suites reject it \ + * earlier, on the alpha type, so there this case only repeats their \ + * existing floating-point alpha coverage. */ \ + expect_bad_scalar_value_dies(9223372036854775808.0); \ } diff --git a/kernels/test/op_bmm_test.cpp b/kernels/test/op_bmm_test.cpp index afc4be856cf..944d8a2d48e 100644 --- a/kernels/test/op_bmm_test.cpp +++ b/kernels/test/op_bmm_test.cpp @@ -471,3 +471,22 @@ TEST_F(OpBmmOutTest, DISABLED_DynamicShapeUnbound) { Tensor ret = op_bmm_out(x, y, out); EXPECT_TENSOR_CLOSE(out, expected_result); } + +TEST_F(OpBmmOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + // All three tensors share the same non-default dim order, so the kernel's + // same dim order check passes and only the default dim order check rejects. + Tensor x = + tf.make_with_dimorder({2, 3, 4}, std::vector(24, 2), {0, 2, 1}); + Tensor y = + tf.make_with_dimorder({2, 4, 5}, std::vector(40, 3), {0, 2, 1}); + Tensor out = + tf.make_with_dimorder({2, 3, 5}, std::vector(30), {0, 2, 1}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE(context_, op_bmm_out(x, y, out)); +} diff --git a/kernels/test/op_cat_test.cpp b/kernels/test/op_cat_test.cpp index d3bda1e8abd..cfdd6e7a426 100644 --- a/kernels/test/op_cat_test.cpp +++ b/kernels/test/op_cat_test.cpp @@ -464,3 +464,21 @@ TEST_F(OpCatOutTest, DynamicShapeUnbound) { op_cat_out(x, 0, out); EXPECT_TENSOR_EQ(out, expected); } + +TEST_F(OpCatOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + Tensor x = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor y = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor out = tf.zeros_channels_last({1, 6, 2, 2}); + std::vector inputs = {x, y}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, op_cat_out(TensorList(inputs.data(), inputs.size()), 1, out)); +} diff --git a/kernels/test/op_constant_pad_nd_test.cpp b/kernels/test/op_constant_pad_nd_test.cpp index 7bd908e0ecb..00a3b5bbef9 100644 --- a/kernels/test/op_constant_pad_nd_test.cpp +++ b/kernels/test/op_constant_pad_nd_test.cpp @@ -484,3 +484,21 @@ TEST_F(OpConstantPadNDOutTest, IncorrectOutputShapeFail) { } GENERATE_SCALAR_OVERFLOW_TESTS(OpConstantPadNDOutTest) + +TEST_F(OpConstantPadNDOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + Tensor self = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor out = tf.zeros_channels_last({1, 3, 2, 4}); + const std::vector padding = {1, 1}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_constant_pad_nd_out( + self, IntArrayRef(padding.data(), padding.size()), 0.0, out)); +} diff --git a/kernels/test/op_cumsum_test.cpp b/kernels/test/op_cumsum_test.cpp index 8ddc197217b..bbf333db4c6 100644 --- a/kernels/test/op_cumsum_test.cpp +++ b/kernels/test/op_cumsum_test.cpp @@ -286,3 +286,18 @@ TEST_F(OpCumSumOutTest, DISABLED_DynamicShapeUnbound) { Tensor ret = op_cumsum_out(x, 1, ScalarType::Float, out); EXPECT_TENSOR_CLOSE(out, expected_result); } + +TEST_F(OpCumSumOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + Tensor in = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor out = tf.zeros_channels_last({1, 3, 2, 2}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, op_cumsum_out(in, 1, ScalarType::Float, out)); +} diff --git a/kernels/test/op_full_test.cpp b/kernels/test/op_full_test.cpp index 752c4e710f4..e412b67dfc4 100644 --- a/kernels/test/op_full_test.cpp +++ b/kernels/test/op_full_test.cpp @@ -84,6 +84,19 @@ ET_FORALL_REALHBF16_TYPES(GENERATE_TEST) GENERATE_SCALAR_OVERFLOW_TESTS(OpFullOutTest) +// The other half of the boundary change: 127.5 used to be refused for an int8 +// tensor and now truncates to 127. +TEST_F(OpFullOutTest, CharTensorFractionalScalarTruncates) { + TensorFactory tf; + std::vector sizes = {2, 2}; + std::vector sizes_int64_t(sizes.begin(), sizes.end()); + auto aref = IntArrayRef(sizes_int64_t.data(), sizes_int64_t.size()); + Tensor out = tf.zeros(sizes); + + op_full_out(aref, 127.5, out); + EXPECT_TENSOR_EQ(out, tf.full(sizes, 127)); +} + TEST_F(OpFullOutTest, HalfSupport) { TensorFactory tf; diff --git a/kernels/test/op_mm_test.cpp b/kernels/test/op_mm_test.cpp index 05d6a7b8d7e..a2044f383f7 100644 --- a/kernels/test/op_mm_test.cpp +++ b/kernels/test/op_mm_test.cpp @@ -294,3 +294,19 @@ TEST_F(OpMmOutTest, DISABLED_DynamicShapeUnbound) { Tensor ret = op_mm_out(x, y, out); EXPECT_TENSOR_CLOSE(out, expected_result); } + +TEST_F(OpMmOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + // All three tensors share the same non-default dim order, so the kernel's + // same dim order check passes and only the default dim order check rejects. + Tensor x = tf.make_with_dimorder({3, 4}, std::vector(12, 2), {1, 0}); + Tensor y = tf.make_with_dimorder({4, 5}, std::vector(20, 3), {1, 0}); + Tensor out = tf.make_with_dimorder({3, 5}, std::vector(15), {1, 0}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE(context_, op_mm_out(x, y, out)); +} diff --git a/kernels/test/op_native_layer_norm_test.cpp b/kernels/test/op_native_layer_norm_test.cpp index e1345a10354..0504716a83a 100644 --- a/kernels/test/op_native_layer_norm_test.cpp +++ b/kernels/test/op_native_layer_norm_test.cpp @@ -452,3 +452,95 @@ TEST_F(OpNativeLayerNormTest, DynamicShapeUnbound) { test_dynamic_shape( {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND); } + +TEST_F(OpNativeLayerNormTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + // mean and rstd share the input's rank with the normalized dims set to 1. + // All four are channels-last so the same-dim-order check passes and only the + // default dim order check can reject. + Tensor input = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor out0 = tf.zeros_channels_last({1, 3, 2, 2}); + Tensor out1 = tf.zeros_channels_last({1, 3, 2, 1}); + Tensor out2 = tf.zeros_channels_last({1, 3, 2, 1}); + const std::vector normalized_shape = {2}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_native_layer_norm_out( + input, + IntArrayRef(normalized_shape.data(), normalized_shape.size()), + exec_aten::optional(), + exec_aten::optional(), + 1e-5, + out0, + out1, + out2)); +} + +TEST_F(OpNativeLayerNormTest, NonDefaultDimOrderWeightDies) { + TensorFactory tf; + + // weight takes the shape of normalized_shape, so it has to be rank 4 to carry + // a channels-last dim order, which makes the input rank 5. Everything else is + // default, so the same dim order check on weight shall be what rejects. + Tensor input = tf.make({2, 2, 2, 2, 2}, std::vector(32, 1)); + Tensor weight = tf.make_with_dimorder( + {2, 2, 2, 2}, std::vector(16, 1), {0, 2, 3, 1}); + Tensor out0 = tf.zeros({2, 2, 2, 2, 2}); + Tensor out1 = tf.zeros({2, 1, 1, 1, 1}); + Tensor out2 = tf.zeros({2, 1, 1, 1, 1}); + const std::vector normalized_shape = {2, 2, 2, 2}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_native_layer_norm_out( + input, + IntArrayRef(normalized_shape.data(), normalized_shape.size()), + exec_aten::optional(weight), + exec_aten::optional(), + 1e-5, + out0, + out1, + out2)); +} + +TEST_F(OpNativeLayerNormTest, NonDefaultDimOrderBiasDies) { + TensorFactory tf; + + // bias takes the shape of normalized_shape, so it has to be rank 4 to carry + // a channels-last dim order, which makes the input rank 5. Everything else is + // default, so the same dim order check on bias shall be what rejects. + Tensor input = tf.make({2, 2, 2, 2, 2}, std::vector(32, 1)); + Tensor bias = tf.make_with_dimorder( + {2, 2, 2, 2}, std::vector(16, 1), {0, 2, 3, 1}); + Tensor out0 = tf.zeros({2, 2, 2, 2, 2}); + Tensor out1 = tf.zeros({2, 1, 1, 1, 1}); + Tensor out2 = tf.zeros({2, 1, 1, 1, 1}); + const std::vector normalized_shape = {2, 2, 2, 2}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_native_layer_norm_out( + input, + IntArrayRef(normalized_shape.data(), normalized_shape.size()), + exec_aten::optional(), + exec_aten::optional(bias), + 1e-5, + out0, + out1, + out2)); +} diff --git a/kernels/test/op_pixel_unshuffle_test.cpp b/kernels/test/op_pixel_unshuffle_test.cpp index 21bed318b9c..9cb4498aabf 100644 --- a/kernels/test/op_pixel_unshuffle_test.cpp +++ b/kernels/test/op_pixel_unshuffle_test.cpp @@ -9,6 +9,7 @@ #include // Declares the operator #include #include +#include #include #include #include @@ -126,3 +127,33 @@ TEST_F(OpPixelUnshuffleOutTest, NegativeUpscaleFactorDies) { // Using a negative upscale factor should exit with an error code. ET_EXPECT_KERNEL_FAILURE(context_, op_pixel_unshuffle_out(a, -3, out)); } + +TEST_F(OpPixelUnshuffleOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + Tensor a = tf.channels_last_like(tf.make( + {1, 1, 4, 4}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})); + Tensor out = tf.zeros_channels_last({1, 4, 2, 2}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE(context_, op_pixel_unshuffle_out(a, 2, out)); +} + +TEST_F(OpPixelUnshuffleOutTest, MixedDimOrderDies) { + TensorFactory tf; + + // Only out has a non-default dim order, so the same dim order check shall be + // what rejects it. + Tensor a = tf.make( + {1, 1, 4, 4}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); + Tensor out = tf.zeros_channels_last({1, 4, 2, 2}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE(context_, op_pixel_unshuffle_out(a, 2, out)); +} diff --git a/kernels/test/op_slice_scatter_test.cpp b/kernels/test/op_slice_scatter_test.cpp index 309f2b8b5f7..256e6e261f5 100644 --- a/kernels/test/op_slice_scatter_test.cpp +++ b/kernels/test/op_slice_scatter_test.cpp @@ -884,3 +884,41 @@ TEST_F(OpSliceScatterTensorOutTest, LargeEndValue) { EXPECT_TENSOR_EQ(ret, out); EXPECT_TENSOR_EQ(ret, expected); } + +TEST_F(OpSliceScatterTensorOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + Tensor input = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor src = tf.zeros_channels_last({1, 1, 2, 2}); + Tensor out = tf.zeros_channels_last({1, 3, 2, 2}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, op_slice_scatter_out(input, src, 1, 0, 1, 1, out)); +} + +TEST_F(OpSliceScatterTensorOutTest, MixedDimOrderDies) { + TensorFactory tf; + + // Only src has a non-default dim order, so the same dim order check shall be + // what rejects it. + Tensor input = + tf.make({1, 3, 2, 4}, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, + 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23}); + Tensor src = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor out = tf.zeros({1, 3, 2, 4}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_slice_scatter_out( + input, src, /*dim=*/3, /*start=*/0, /*end=*/2, /*step=*/1, out)); +} diff --git a/kernels/test/op_split_copy_test.cpp b/kernels/test/op_split_copy_test.cpp index 34df2c749ff..7714770d383 100644 --- a/kernels/test/op_split_copy_test.cpp +++ b/kernels/test/op_split_copy_test.cpp @@ -576,3 +576,22 @@ TEST_F(OpSplitCopyTensorOutTest, DISABLED_DynamicShapeUnbound) { test_dynamic_shape( {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND); } + +TEST_F(OpSplitCopyTensorOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + Tensor input = tf.channels_last_like(tf.make( + {1, 4, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})); + std::vector outs = { + tf.zeros_channels_last({1, 2, 2, 2}), + tf.zeros_channels_last({1, 2, 2, 2})}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_split_copy_tensor_out( + input, 2, 1, TensorList(outs.data(), outs.size()))); +} diff --git a/kernels/test/op_split_with_sizes_copy_test.cpp b/kernels/test/op_split_with_sizes_copy_test.cpp index cc81ffff19d..f33b622490f 100644 --- a/kernels/test/op_split_with_sizes_copy_test.cpp +++ b/kernels/test/op_split_with_sizes_copy_test.cpp @@ -9,6 +9,7 @@ #include // Declares the operator #include #include +#include #include #include #include @@ -115,3 +116,28 @@ TEST_F(OpSplitWithSizesCopyOutTest, DynamicShape) { test_tensor_shape_dynamism( executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND); } + +TEST_F(OpSplitWithSizesCopyOutTest, NonDefaultDimOrderDies) { + torch::executor::testing::TensorFactory + tf; + + executorch::aten::Tensor self = tf.channels_last_like(tf.make( + {1, 4, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})); + std::vector outs = { + tf.zeros_channels_last({1, 2, 2, 2}), + tf.zeros_channels_last({1, 2, 2, 2})}; + const std::vector split_sizes = {2, 2}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_split_with_sizes_copy_out( + self, + executorch::aten::ArrayRef( + split_sizes.data(), split_sizes.size()), + 1, + executorch::aten::TensorList(outs.data(), outs.size()))); +} diff --git a/kernels/test/op_to_copy_test.cpp b/kernels/test/op_to_copy_test.cpp index 45b2b2f6020..a6fb6239390 100644 --- a/kernels/test/op_to_copy_test.cpp +++ b/kernels/test/op_to_copy_test.cpp @@ -7,7 +7,9 @@ */ #include +#include #include +#include #include #include @@ -22,6 +24,7 @@ #include using namespace ::testing; +using executorch::aten::BFloat16; using executorch::aten::MemoryFormat; using executorch::aten::ScalarType; using executorch::aten::Tensor; @@ -81,6 +84,22 @@ class OpToTest : public OperatorTest { const std::vector data_out; }; + static float float_from_bits(uint32_t bits) { + float value; + std::memcpy(&value, &bits, sizeof(value)); + return value; + } + + static uint32_t float_bits(float value) { + uint32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; + } + + static bool is_bfloat16_nan(BFloat16 value) { + return (value.x & 0x7FFF) > 0x7F80; + } + // Each test has different combination of input and output types. Therefore it // is a little bit mess if create template test case and custom data types for // both input data and output data. @@ -121,6 +140,81 @@ class OpToTest : public OperatorTest { } } + template < + typename INPUT_CTYPE, + ScalarType INPUT_DTYPE, + typename OUTPUT_CTYPE, + ScalarType OUTPUT_DTYPE> + void test_conversion_at_sizes( + const std::vector& sizes, + const std::vector& input_pattern, + const std::vector& expected_pattern) { + static_assert( + (std::is_same_v && + std::is_same_v) || + (std::is_same_v && + std::is_same_v), + "Only float/BFloat16 conversion pairs are supported"); + ASSERT_EQ(input_pattern.size(), expected_pattern.size()); + + TensorFactory tf_in; + TensorFactory tf_out; + for (const int32_t numel : sizes) { + SCOPED_TRACE(::testing::Message() << "numel=" << numel); + std::vector input_data; + std::vector expected_data; + input_data.reserve(numel); + expected_data.reserve(numel); + for (int32_t i = 0; i < numel; ++i) { + input_data.push_back(input_pattern[i % input_pattern.size()]); + expected_data.push_back(expected_pattern[i % expected_pattern.size()]); + } + + Tensor input = tf_in.make({numel}, input_data); + Tensor output = tf_out.zeros({numel}); + + Tensor& ret = op_to_copy_out( + input, + /*non_blocking=*/false, + executorch::aten::MemoryFormat::Contiguous, + output); + + EXPECT_EQ(&ret, &output); + const auto* const actual_data = ret.const_data_ptr(); + if (actual_data == nullptr) { + ADD_FAILURE() << "conversion returned a null data pointer"; + continue; + } + if constexpr (std::is_same_v) { + const bool is_aten = + torch::executor::testing::SupportedFeatures::get()->is_aten; + std::vector actual_bits; + std::vector expected_bits; + actual_bits.reserve(numel); + expected_bits.reserve(numel); + for (int32_t i = 0; i < numel; ++i) { + if (is_aten && is_bfloat16_nan(expected_data[i])) { + EXPECT_TRUE(is_bfloat16_nan(actual_data[i])) << "index=" << i; + continue; + } + actual_bits.push_back(actual_data[i].x); + expected_bits.push_back(expected_data[i].x); + } + EXPECT_EQ(actual_bits, expected_bits); + } else { + std::vector actual_bits; + std::vector expected_bits; + actual_bits.reserve(numel); + expected_bits.reserve(numel); + for (int32_t i = 0; i < numel; ++i) { + actual_bits.push_back(float_bits(actual_data[i])); + expected_bits.push_back(float_bits(expected_data[i])); + } + EXPECT_EQ(actual_bits, expected_bits); + } + } + } + template void test_runner_to_bool( std::vector test_case, @@ -360,6 +454,140 @@ TEST_F(OpToTest, NanInfSupported) { #undef TEST_KERNEL } +TEST_F(OpToTest, FloatToBFloat16RawBitsAtVectorAndGrainBoundaries) { + std::vector sizes; + for (int32_t size = 1; size <= 17; ++size) { + sizes.push_back(size); + } + sizes.insert(sizes.end(), {32767, 32768, 32769}); + + const std::vector input_pattern = { + float_from_bits(0x00000000), float_from_bits(0x80000000), + float_from_bits(0x00007FFF), float_from_bits(0x00008000), + float_from_bits(0x00008001), float_from_bits(0x00017FFF), + float_from_bits(0x00018000), float_from_bits(0x00018001), + float_from_bits(0x3F807FFF), float_from_bits(0x3F808000), + float_from_bits(0x3F808001), float_from_bits(0x3F817FFF), + float_from_bits(0x3F818000), float_from_bits(0x3F818001), + float_from_bits(0xBF807FFF), float_from_bits(0xBF808000), + float_from_bits(0xBF808001), float_from_bits(0xBF817FFF), + float_from_bits(0xBF818000), float_from_bits(0xBF818001), + float_from_bits(0x7F800000), float_from_bits(0xFF800000), + float_from_bits(0x7FC12345), float_from_bits(0xFFC12345), + float_from_bits(0x7FA12345), float_from_bits(0xFFA12345), + }; + const std::vector expected_pattern = { + BFloat16(0x0000, BFloat16::from_bits()), + BFloat16(0x8000, BFloat16::from_bits()), + BFloat16(0x0000, BFloat16::from_bits()), + BFloat16(0x0000, BFloat16::from_bits()), + BFloat16(0x0001, BFloat16::from_bits()), + BFloat16(0x0001, BFloat16::from_bits()), + BFloat16(0x0002, BFloat16::from_bits()), + BFloat16(0x0002, BFloat16::from_bits()), + BFloat16(0x3F80, BFloat16::from_bits()), + BFloat16(0x3F80, BFloat16::from_bits()), + BFloat16(0x3F81, BFloat16::from_bits()), + BFloat16(0x3F81, BFloat16::from_bits()), + BFloat16(0x3F82, BFloat16::from_bits()), + BFloat16(0x3F82, BFloat16::from_bits()), + BFloat16(0xBF80, BFloat16::from_bits()), + BFloat16(0xBF80, BFloat16::from_bits()), + BFloat16(0xBF81, BFloat16::from_bits()), + BFloat16(0xBF81, BFloat16::from_bits()), + BFloat16(0xBF82, BFloat16::from_bits()), + BFloat16(0xBF82, BFloat16::from_bits()), + BFloat16(0x7F80, BFloat16::from_bits()), + BFloat16(0xFF80, BFloat16::from_bits()), + BFloat16(0x7FC0, BFloat16::from_bits()), + BFloat16(0x7FC0, BFloat16::from_bits()), + BFloat16(0x7FC0, BFloat16::from_bits()), + BFloat16(0x7FC0, BFloat16::from_bits()), + }; + + test_conversion_at_sizes< + float, + ScalarType::Float, + BFloat16, + ScalarType::BFloat16>(sizes, input_pattern, expected_pattern); +} + +#if defined(__aarch64__) +TEST_F(OpToTest, FloatToBFloat16SubnormalsIgnoreFlushToZero) { + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen conversion may flush subnormals to zero"); + + struct RestoreFpcr { + uint64_t value{}; + RestoreFpcr() { + asm volatile("mrs %0, fpcr" : "=r"(value)); + } + ~RestoreFpcr() { + asm volatile("msr fpcr, %0" : : "r"(value) : "memory"); + } + } original_fpcr; + + constexpr uint64_t kFlushToZero = uint64_t{1} << 24; + for (const bool flush_to_zero : {false, true}) { + SCOPED_TRACE(::testing::Message() << "flush_to_zero=" << flush_to_zero); + const uint64_t fpcr = flush_to_zero ? original_fpcr.value | kFlushToZero + : original_fpcr.value & ~kFlushToZero; + asm volatile("msr fpcr, %0" : : "r"(fpcr) : "memory"); + test_conversion_at_sizes< + float, + ScalarType::Float, + BFloat16, + ScalarType::BFloat16>( + {1, 7, 8, 9, 15, 16, 17}, + {float_from_bits(0x00018000), float_from_bits(0x80018000)}, + {BFloat16(0x0002, BFloat16::from_bits()), + BFloat16(0x8002, BFloat16::from_bits())}); + } +} +#endif + +TEST_F(OpToTest, BFloat16ToFloatRawBitsAtVectorAndGrainBoundaries) { + std::vector sizes; + for (int32_t size = 1; size <= 17; ++size) { + sizes.push_back(size); + } + sizes.insert(sizes.end(), {32767, 32768, 32769}); + + const std::vector input_pattern = { + BFloat16(0x0000, BFloat16::from_bits()), + BFloat16(0x8000, BFloat16::from_bits()), + BFloat16(0x3F80, BFloat16::from_bits()), + BFloat16(0x3F81, BFloat16::from_bits()), + BFloat16(0x3F82, BFloat16::from_bits()), + BFloat16(0x7F80, BFloat16::from_bits()), + BFloat16(0xFF80, BFloat16::from_bits()), + BFloat16(0x7FC1, BFloat16::from_bits()), + BFloat16(0xFFC1, BFloat16::from_bits()), + BFloat16(0x7FA1, BFloat16::from_bits()), + BFloat16(0xFFA1, BFloat16::from_bits()), + }; + const std::vector expected_pattern = { + float_from_bits(0x00000000), + float_from_bits(0x80000000), + float_from_bits(0x3F800000), + float_from_bits(0x3F810000), + float_from_bits(0x3F820000), + float_from_bits(0x7F800000), + float_from_bits(0xFF800000), + float_from_bits(0x7FC10000), + float_from_bits(0xFFC10000), + float_from_bits(0x7FA10000), + float_from_bits(0xFFA10000), + }; + + test_conversion_at_sizes< + BFloat16, + ScalarType::BFloat16, + float, + ScalarType::Float>(sizes, input_pattern, expected_pattern); +} + TEST_F(OpToTest, HardcodeFloatConvertInt) { // Hardcode input and output generated from core PyTorch // clang-format off diff --git a/kernels/test/op_topk_test.cpp b/kernels/test/op_topk_test.cpp index 17c7141d12d..331d3c276eb 100644 --- a/kernels/test/op_topk_test.cpp +++ b/kernels/test/op_topk_test.cpp @@ -8,6 +8,8 @@ #include // Declares the operator #include +#include +#include #include #include #include @@ -173,3 +175,49 @@ TEST_F(OpTopkValuesTest, NonPartialSort) { EXPECT_TENSOR_EQ(indices, indices_expected); } } + +TEST_F(OpTopkValuesTest, NonDefaultDimOrderDies) { + TensorFactory tf; + TensorFactory tf_long; + + Tensor in = tf.channels_last_like(tf.make( + {1, 4, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})); + Tensor values = tf.zeros_channels_last({1, 2, 2, 2}); + Tensor indices = tf_long.zeros_channels_last({1, 2, 2, 2}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + TempMemoryAllocator allocator = TempMemoryAllocator(); + executorch::ET_RUNTIME_NAMESPACE::KernelRuntimeContext context( + nullptr, &allocator); + torch::executor::aten::topk_outf( + context, in, 2, 1, true, true, values, indices); + + EXPECT_NE(context.failure_state(), torch::executor::Error::Ok); +} + +TEST_F(OpTopkValuesTest, MixedDimOrderDies) { + TensorFactory tf; + TensorFactory tf_long; + + // Only values has a non-default dim order, so the same dim order check shall + // be what rejects it. + Tensor in = tf.make( + {1, 4, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); + Tensor values = tf.zeros_channels_last({1, 2, 2, 2}); + Tensor indices = tf_long.zeros({1, 2, 2, 2}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + TempMemoryAllocator allocator = TempMemoryAllocator(); + executorch::ET_RUNTIME_NAMESPACE::KernelRuntimeContext context( + nullptr, &allocator); + torch::executor::aten::topk_outf( + context, in, 2, 1, true, true, values, indices); + + EXPECT_NE(context.failure_state(), torch::executor::Error::Ok); +} diff --git a/kernels/test/targets.bzl b/kernels/test/targets.bzl index 837c7327c4f..3d36071260b 100644 --- a/kernels/test/targets.bzl +++ b/kernels/test/targets.bzl @@ -252,7 +252,7 @@ def define_common_targets(): _common_op_test("op_expand_copy_test", ["aten", "portable"]) _common_op_test("op_expm1_test", ["aten", "portable"]) _common_op_test("op_fft_c2r_test", ["aten", "optimized"]) - _common_op_test("op_fft_r2c_test", ["aten", "optimized"]) + _common_op_test("op_fft_r2c_test", ["aten", "portable", "optimized"]) _common_op_test("op_fill_test", ["aten", "portable"]) _common_op_test("op_flip_test", ["aten", "portable"]) _common_op_test("op_floor_divide_test", ["aten", "portable"]) @@ -352,7 +352,7 @@ def define_common_targets(): _common_op_test("op_t_copy_test", ["aten", "portable"]) _common_op_test("op_tan_test", ["aten", "portable"]) _common_op_test("op_tanh_test", ["aten", "portable"]) - _common_op_test("op_to_copy_test", ["aten", "portable"]) + _common_op_test("op_to_copy_test", ["aten", "portable", "optimized"]) _common_op_test("op_topk_test", ["aten", "portable"]) _common_op_test("op_transpose_copy_test", ["aten", "portable"]) _common_op_test("op_tril_test", ["aten", "portable"]) diff --git a/pyproject.toml b/pyproject.toml index 8f630a71e42..f7cf98679fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [build-system] requires = [ - "cmake>=3.24,<4.0.0", # For building binary targets in the wheel. 4.0.0 breaks third-party CMake build so temporarily pin the version. + "cmake>=3.26,<4.0.0", # For building binary targets in the wheel. 4.0.0 breaks third-party CMake build so temporarily pin the version. "packaging>=24.2", # Lower bound required by setuptools "patchelf; sys_platform == 'linux'", # Writes the runtime search paths that let the shipped libraries find each other. "pip>=23", # For building the pip package. @@ -100,9 +100,6 @@ Changelog = "https://github.com/pytorch/executorch/releases" [project.scripts] flatc = "executorch.data.bin:flatc" -# TODO(dbort): Could use py_modules to restrict the set of modules we -# package, and package_data to restrict the set up non-python files we -# include. See also setuptools/discovery.py for custom finders. [tool.setuptools] license-files = ["LICENSE"] @@ -118,10 +115,9 @@ license-files = ["LICENSE"] "executorch" = "src/executorch" [tool.setuptools.package-data] -# TODO(dbort): Prune /test[s]/ dirs, /third-party/ dirs, yaml files that we -# don't need. -# TODO(RobertKalmar): When test[s] dirs pruned the PROJECT_DIR resolution in backends.nxp.tests_models.config.py can -# avoid exporting and reading env variable. +# TODO(RobertKalmar): the PROJECT_DIR env variable is still needed. Test directories still install, +# since the suites import shared helpers from them, and the artifacts that config.py resolves are +# not in the wheel, so a path derived from __file__ would point at files that are not there. "*" = [ # Some backends like XNNPACK need their .fbs files. "*.fbs", diff --git a/requirements-dev.txt b/requirements-dev.txt index c916ccfa472..1220547aeb5 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,6 +1,6 @@ # Pip packages needed to build from source. Mainly for development of ExecuTorch. -cmake>=3.24, <4.0.0 # For building binary targets in the wheel. +cmake>=3.26, <4.0.0 # For building binary targets in the wheel. packaging>=24.2 # Lower bound required by setuptools patchelf; sys_platform == 'linux' # Writes the runtime search paths that let the shipped libraries find each other. pip>=23 # For building the pip package. diff --git a/runtime/core/portable_type/c10/c10/targets.bzl b/runtime/core/portable_type/c10/c10/targets.bzl index 675c04f97a1..a39b2b9f566 100644 --- a/runtime/core/portable_type/c10/c10/targets.bzl +++ b/runtime/core/portable_type/c10/c10/targets.bzl @@ -114,7 +114,6 @@ def define_common_targets(): "util/bit_cast.h", "util/complex.h", "util/complex_math.h", - "util/complex_utils.h", "util/floating_point_utils.h", "util/irange.h", "util/llvmMathExtras.h", diff --git a/runtime/core/portable_type/c10/c10/util/complex.h b/runtime/core/portable_type/c10/c10/util/complex.h index 4e699684bc3..f9849a94ced 100644 --- a/runtime/core/portable_type/c10/c10/util/complex.h +++ b/runtime/core/portable_type/c10/c10/util/complex.h @@ -31,19 +31,11 @@ C10_HOST_DEVICE T abs(const c10::complex& z) { #endif } -#if defined(USE_ROCM) -#define ROCm_Bug(x) -#else -#define ROCm_Bug(x) x -#endif - template C10_HOST_DEVICE T arg(const c10::complex& z) { - return ROCm_Bug(std)::atan2(std::imag(z), std::real(z)); + return std::atan2(std::imag(z), std::real(z)); } -#undef ROCm_Bug - template constexpr T norm(const c10::complex& z) { return z.real() * z.real() + z.imag() * z.imag(); @@ -73,6 +65,9 @@ constexpr c10::complex conj(const c10::complex& z) { #define C10_INTERNAL_INCLUDE_COMPLEX_REMAINING_H // math functions are included in a separate file #include // IWYU pragma: keep -// utilities for complex types -#include // IWYU pragma: keep #undef C10_INTERNAL_INCLUDE_COMPLEX_REMAINING_H + +namespace c10 { +using torch::headeronly::is_complex; +using torch::headeronly::scalar_value_type; +} // namespace c10 diff --git a/runtime/core/portable_type/c10/c10/util/llvmMathExtras.h b/runtime/core/portable_type/c10/c10/util/llvmMathExtras.h index da297241449..8ae5cde4f02 100644 --- a/runtime/core/portable_type/c10/c10/util/llvmMathExtras.h +++ b/runtime/core/portable_type/c10/c10/util/llvmMathExtras.h @@ -400,6 +400,7 @@ constexpr inline bool isShiftedUInt(uint64_t x) { N + S <= 64, "isShiftedUInt with N + S > 64 is too wide."); // Per the two static_asserts above, S must be strictly less than 64. So // 1 << S is not undefined behavior. + // NOLINTNEXTLINE(bugprone-chained-comparison) return isUInt(x) && (x % (UINT64_C(1) << S) == 0); } diff --git a/runtime/core/portable_type/c10/c10/util/overflows.h b/runtime/core/portable_type/c10/c10/util/overflows.h index 183a2f62a32..348ee5ffa40 100644 --- a/runtime/core/portable_type/c10/c10/util/overflows.h +++ b/runtime/core/portable_type/c10/c10/util/overflows.h @@ -61,13 +61,28 @@ template std::enable_if_t, bool> overflows( From f, bool strict_unsigned [[maybe_unused]] = false) { - using limit = std::numeric_limits::type>; + using ToScalar = typename scalar_value_type::type; + using limit = std::numeric_limits; if (limit::has_infinity && std::isinf(static_cast(f))) { return false; } if (!limit::has_quiet_NaN && (f != f)) { return true; } + if constexpr (std::is_integral_v) { + // limit::max() for wide integer types is NOT exactly representable in + // floating point (e.g. int64 max = 2^63-1 rounds up to 2^63), so `f > + // limit::max()` lets a just-out-of-range value like 2^63 slip through and + // then become INT64_MIN via static_cast. Compare against the + // exactly-representable upper bound max()+1 == 2^digits instead. lowest() + // is 0 or a negated power of two, so it stays exact. (digits-1 keeps the + // shift < 64 for the uint64 case; the *2 recovers 2^digits without a 1<<64 + // overflow.) + constexpr int digits = limit::digits; + constexpr From upper = + static_cast(uint64_t{1} << (digits - 1)) * From{2}; + return f < static_cast(limit::lowest()) || f >= upper; + } return f < limit::lowest() || f > limit::max(); } diff --git a/runtime/core/portable_type/c10/torch/headeronly/macros/Macros.h b/runtime/core/portable_type/c10/torch/headeronly/macros/Macros.h index cef99df3f56..08c4e9f1f84 100644 --- a/runtime/core/portable_type/c10/torch/headeronly/macros/Macros.h +++ b/runtime/core/portable_type/c10/torch/headeronly/macros/Macros.h @@ -123,6 +123,15 @@ #define C10_HAS_CPP_ATTRIBUTE(x) (0) #endif +/// Bind a returned reference/pointer's lifetime to a parameter (or *this) so +/// Clang can warn when it would dangle. Expands to nothing on compilers that +/// lack the attribute (e.g. non-clang, older nvcc). +#if C10_HAS_CPP_ATTRIBUTE(clang::lifetimebound) +#define C10_LIFETIMEBOUND [[clang::lifetimebound]] +#else +#define C10_LIFETIMEBOUND +#endif + #ifndef FBCODE_CAFFE2 /// DEPRECATED: Warn if a type or return value is discarded. #define C10_NODISCARD [[nodiscard]] diff --git a/runtime/core/portable_type/c10/torch/headeronly/util/Half.h b/runtime/core/portable_type/c10/torch/headeronly/util/Half.h index e5aa622656c..401472357ec 100644 --- a/runtime/core/portable_type/c10/torch/headeronly/util/Half.h +++ b/runtime/core/portable_type/c10/torch/headeronly/util/Half.h @@ -213,7 +213,7 @@ C10_HOST_DEVICE inline float fp16_ieee_to_fp32_value(uint16_t h) { * Now, remember that denormalized half-precision numbers are represented as: * FP16 = mantissa * 2**(-24). * The trick is to construct a normalized single-precision number with the - * same mantissa and thehalf-precision input and with an exponent which would + * same mantissa and the half-precision input and with an exponent which would * scale the corresponding mantissa bits to 2**(-24). A normalized * single-precision floating-point number is represented as: FP32 = (1 + * mantissa * 2**(-23)) * 2**(exponent - 127) Therefore, when the biased diff --git a/runtime/core/portable_type/c10/torch/headeronly/util/TypeSafeSignMath.h b/runtime/core/portable_type/c10/torch/headeronly/util/TypeSafeSignMath.h index c33a286bc5b..8e897957fee 100644 --- a/runtime/core/portable_type/c10/torch/headeronly/util/TypeSafeSignMath.h +++ b/runtime/core/portable_type/c10/torch/headeronly/util/TypeSafeSignMath.h @@ -14,20 +14,6 @@ C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-int-float-conversion") namespace c10 { -/// Returns false since we cannot have x < 0 if x is unsigned. -template -inline constexpr bool is_negative( - const T& /*x*/, - std::true_type /*is_unsigned*/) { - return false; -} - -/// Returns true if a signed variable x < 0 -template -inline constexpr bool is_negative(const T& x, std::false_type /*is_unsigned*/) { - return x < T(0); -} - /// Returns true if x < 0 /// NOTE: Will fail on an unsigned custom type /// For the most part it's possible to fix this if @@ -35,19 +21,12 @@ inline constexpr bool is_negative(const T& x, std::false_type /*is_unsigned*/) { /// However, notably, c10::Half does not :-( template inline constexpr bool is_negative(const T& x) { - return is_negative(x, std::is_unsigned()); -} - -/// Returns the sign of an unsigned variable x as 0, 1 -template -inline constexpr int signum(const T& x, std::true_type /*is_unsigned*/) { - return T(0) < x; -} - -/// Returns the sign of a signed variable x as -1, 0, 1 -template -inline constexpr int signum(const T& x, std::false_type /*is_unsigned*/) { - return (T(0) < x) - (x < T(0)); + if constexpr (std::is_unsigned_v) { + // An unsigned value can never be less than zero. + return false; + } else { + return x < T(0); + } } /// Returns the sign of x as -1, 0, 1 @@ -57,7 +36,11 @@ inline constexpr int signum(const T& x, std::false_type /*is_unsigned*/) { /// However, notably, c10::Half does not :-( template inline constexpr int signum(const T& x) { - return signum(x, std::is_unsigned()); + if constexpr (std::is_unsigned_v) { + return T(0) < x; + } else { + return (T(0) < x) - (x < T(0)); + } } /// Returns true if a and b are not both negative @@ -86,53 +69,22 @@ inline constexpr bool greater_than_max(const T& x) { #pragma GCC diagnostic pop #endif -/// Returns true if x < lowest(Limit). Standard comparison -template -inline constexpr bool less_than_lowest( - const T& x, - std::false_type /*limit_is_unsigned*/, - std::false_type /*x_is_unsigned*/) { - return x < std::numeric_limits::lowest(); -} - -/// Returns false since all the limit is signed and therefore includes -/// negative values but x cannot be negative because it is unsigned -template -inline constexpr bool less_than_lowest( - const T& /*x*/, - std::false_type /*limit_is_unsigned*/, - std::true_type /*x_is_unsigned*/) { - return false; -} - -/// Returns true if x < 0, where 0 is constructed from T. -/// Limit is not signed, so its lower value is zero -template -inline constexpr bool less_than_lowest( - const T& x, - std::true_type /*limit_is_unsigned*/, - std::false_type /*x_is_unsigned*/) { - return x < T(0); -} - -/// Returns false sign both types are unsigned -template -inline constexpr bool less_than_lowest( - const T& /*x*/, - std::true_type /*limit_is_unsigned*/, - std::true_type /*x_is_unsigned*/) { - return false; -} - -/// Returns true if x is less than the lowest value of type T +/// Returns true if x is less than the lowest value of type Limit /// NOTE: Will fail on an unsigned custom type /// For the most part it's possible to fix this if /// the custom type has a constexpr constructor. /// However, notably, c10::Half does not : template inline constexpr bool less_than_lowest(const T& x) { - return less_than_lowest( - x, std::is_unsigned(), std::is_unsigned()); + if constexpr (std::is_unsigned_v) { + // x is unsigned, so it can never be below the lowest value of any type. + return false; + } else if constexpr (std::is_unsigned_v) { + // Limit is unsigned, so its lowest value is zero. + return x < T(0); + } else { + return x < std::numeric_limits::lowest(); + } } } // namespace c10 diff --git a/runtime/core/portable_type/c10/torch/headeronly/util/complex.h b/runtime/core/portable_type/c10/torch/headeronly/util/complex.h index 733a22d5dbb..c349602dcf0 100644 --- a/runtime/core/portable_type/c10/torch/headeronly/util/complex.h +++ b/runtime/core/portable_type/c10/torch/headeronly/util/complex.h @@ -3,6 +3,7 @@ #include #include +#include #include #if defined(__CUDACC__) || defined(__HIPCC__) @@ -588,6 +589,60 @@ struct alignas(4) complex { } }; +template <> +struct alignas(4) complex { + BFloat16 real_; + BFloat16 imag_; + + // Constructors + complex() = default; + // BFloat16 constructor is not constexpr so the following constructor can't + // be constexpr + C10_HOST_DEVICE explicit inline complex( + const BFloat16& real, + const BFloat16& imag) + : real_(real), imag_(imag) {} + C10_HOST_DEVICE inline complex(const c10::complex& value) + : real_(value.real()), imag_(value.imag()) {} + + // Conversion operator + inline C10_HOST_DEVICE operator c10::complex() const { + return {real_, imag_}; + } + + constexpr C10_HOST_DEVICE BFloat16 real() const { + return real_; + } + constexpr C10_HOST_DEVICE BFloat16 imag() const { + return imag_; + } + + C10_HOST_DEVICE complex& operator+=( + const complex& other) { + real_ = static_cast(real_) + static_cast(other.real_); + imag_ = static_cast(imag_) + static_cast(other.imag_); + return *this; + } + + C10_HOST_DEVICE complex& operator-=( + const complex& other) { + real_ = static_cast(real_) - static_cast(other.real_); + imag_ = static_cast(imag_) - static_cast(other.imag_); + return *this; + } + + C10_HOST_DEVICE complex& operator*=( + const complex& other) { + auto a = static_cast(real_); + auto b = static_cast(imag_); + auto c = static_cast(other.real()); + auto d = static_cast(other.imag()); + real_ = a * c - b * d; + imag_ = a * d + b * c; + return *this; + } +}; + } // namespace c10 HIDDEN_NAMESPACE_BEGIN(torch, headeronly) @@ -614,3 +669,8 @@ using c10::complex_literals::operator""_id; HIDDEN_NAMESPACE_END(torch, headeronly) C10_CLANG_DIAGNOSTIC_POP() + +#define C10_INTERNAL_INCLUDE_COMPLEX_REMAINING_H +// utilities for complex types +#include // IWYU pragma: keep +#undef C10_INTERNAL_INCLUDE_COMPLEX_REMAINING_H diff --git a/runtime/core/portable_type/c10/c10/util/complex_utils.h b/runtime/core/portable_type/c10/torch/headeronly/util/complex_utils.h similarity index 80% rename from runtime/core/portable_type/c10/c10/util/complex_utils.h rename to runtime/core/portable_type/c10/torch/headeronly/util/complex_utils.h index 1ca105f1d0a..ddc66ffe776 100644 --- a/runtime/core/portable_type/c10/c10/util/complex_utils.h +++ b/runtime/core/portable_type/c10/torch/headeronly/util/complex_utils.h @@ -1,11 +1,13 @@ +#pragma once + #if !defined(C10_INTERNAL_INCLUDE_COMPLEX_REMAINING_H) #error \ - "c10/util/complex_utils.h is not meant to be individually included. Include c10/util/complex.h instead." + "torch/headeronly/util/complex_utils.h is not meant to be individually included. Include torch/headeronly/util/complex.h instead." #endif #include -namespace c10 { +HIDDEN_NAMESPACE_BEGIN(torch, headeronly) template struct is_complex : public std::false_type {}; @@ -31,7 +33,7 @@ struct scalar_value_type> { using type = T; }; -} // namespace c10 +HIDDEN_NAMESPACE_END(torch, headeronly) namespace std { diff --git a/scripts/build_apple_frameworks.sh b/scripts/build_apple_frameworks.sh index bc201858b30..0fe7b27e7c1 100755 --- a/scripts/build_apple_frameworks.sh +++ b/scripts/build_apple_frameworks.sh @@ -8,7 +8,7 @@ set -euxo pipefail MODES=() -PRESETS=("ios" "ios-simulator" "macos") +PRESETS=("apple-framework-ios" "apple-framework-ios-simulator" "apple-framework-macos") # To support backwards compatibility, we want to retain the same output directory. PRESETS_RELATIVE_OUT_DIR=("ios" "simulator" "macos") diff --git a/setup.py b/setup.py index c80131657f5..a00dfd69909 100644 --- a/setup.py +++ b/setup.py @@ -47,7 +47,9 @@ # other computer software, distribute, and sublicense such enhancements or # derivative works thereof, in binary and source code form. +import ast import contextlib +import functools # Import this before distutils so that setuptools can intercept the distuils # imports. @@ -64,7 +66,7 @@ from distutils import log # type: ignore[import-not-found] from distutils.sysconfig import get_python_lib # type: ignore[import-not-found] from pathlib import Path, PurePosixPath -from typing import List, Optional +from typing import Dict, FrozenSet, List, Optional, Set, Tuple # Clean dynamic import using importlib _install_utils_path = Path(__file__).parent / "install_utils.py" @@ -177,10 +179,117 @@ def _minimal_cmake_flags() -> List[str]: ] +_VENDORED_DIR_NAMES = frozenset({"third-party", "third_party"}) + +# Used only when .gitmodules cannot be read, as in a source distribution. A test keeps it in step. +_VENDORED_SUBMODULE_FALLBACK = ( + "backends/cadence/utils/FACTO", + "extension/llm/tokenizers", +) + + +@functools.lru_cache(maxsize=None) +def _vendored_prefixes() -> Tuple[str, ...]: + """Source-tree prefixes holding code from another repository. + + Two shapes reach the wheel. Most vendored code sits in a directory named third-party, + which the name above covers wherever it appears. The rest are git submodules checked out + under an ordinary name, so they can only be recognized by asking git what they are. + + None of them are importable from where they sit. FACTO is pure Python but its nested copy + cannot satisfy backends/cadence/utils/facto_util.py, which imports the top level facto.specdb, + and the tokenizers ship separately as pytorch-tokenizers in the dependency list. The rest, + XNNPACK and the Vulkan headers among them, are C++ sources that the wheel has no use for once + the libraries are built. + + Read through git rather than by scanning the file, so only real submodule entries count. + A hand-rolled reader accepts a `path` line from any section, and one stray line elsewhere + in the file would drop a first-party package from the wheel with nothing to warn about. + + Submodules at the repository root are skipped. Those are build tooling, never copied into + the package, and carrying a bare single-word name here would make the match below drop any + directory that happened to share it. + """ + root = Path(__file__).parent + if not (root / ".gitmodules").is_file(): + # A source distribution carries no .gitmodules, so nothing can be read there. Fall + # back to the directories the vendored trees occupy, or the exclusion would quietly + # do half its job and those files would ship again. + return _VENDORED_SUBMODULE_FALLBACK + try: + listed = subprocess.run( + [ + "git", + "config", + "-z", + "-f", + ".gitmodules", + "--get-regexp", + r"^submodule\..*\.path$", + ], + cwd=root, + capture_output=True, + text=True, + check=False, + ) + except OSError: + # No git on PATH, so fall back for the same reason as above. + return _VENDORED_SUBMODULE_FALLBACK + + if listed.returncode or not listed.stdout.strip(): + # git ran and told us nothing useful, which happens when the file has a bad section + # header or conflict markers in it. Reading that as "no submodules" would turn the + # exclusion off without a word, so fall back rather than trust an empty answer. + return _VENDORED_SUBMODULE_FALLBACK + + prefixes = [] + # -z separates each record with NUL and its key from its value with a newline, so neither a + # name nor a path containing a space can be misread. Splitting the default space-separated + # output cannot do that: "submodule.a b.path c d/e" is ambiguous either way round. + for record in listed.stdout.split("\0"): + if not record: + continue + _, separator, value = record.partition("\n") + if not separator: + continue + # Normalize, because git accepts a trailing slash, a ./ prefix and doubled + # separators as the same path, and the raw text would stop matching the real + # directory. + parts = Path(value).parts + if len(parts) < 2 or any(part in _VENDORED_DIR_NAMES for part in parts): + continue + prefixes.append("/".join(parts)) + return tuple(sorted(prefixes)) + + +def _is_vendored_path(path: str) -> bool: + """Whether a source-tree path holds code from another repository.""" + parts = Path(path).parts + if any(part in _VENDORED_DIR_NAMES for part in parts): + return True + # A submodule path is relative to the repository root, while a path here may be relative + # to src/executorch or carry a src/executorch prefix, so match on any suffix boundary. + # Whole-component match: the prefix must be the entire path, or sit at its start, end, or + # middle bounded by separators. Substring matching would let a directory whose name merely + # begins with a prefix be dropped. + posix = "/".join(parts) + return any( + posix == prefix + or posix.startswith(f"{prefix}/") + or posix.endswith(f"/{prefix}") + or f"/{prefix}/" in posix + for prefix in _vendored_prefixes() + ) + + def _minimal_packages() -> List[str]: return sorted( find_namespace_packages( - where="src", + # Anchored on this file, not the working directory, so the list does not change + # with where the build was started from. A cwd-relative path returns nothing when + # the build runs from anywhere but the repository root, and a wheel with no packages + # in it ships no Python at all. + where=str(Path(__file__).parent / "src"), include=[ "executorch", "executorch.data", @@ -204,6 +313,334 @@ def _minimal_packages() -> List[str]: ) +_WALK_SKIP_DIRS = frozenset( + {".git", "pip-out", "cmake-out", "third-party", "third_party", "__pycache__"} +) + +_TEST_DIR_NAMES = frozenset({"test", "tests"}) + + +@functools.lru_cache(maxsize=None) +def _top_level_package_dirs() -> FrozenSet[str]: + """The first path component of every package the wheel ships. + + Derived from the tree rather than listed, so a new top-level directory is covered without an + edit here. Used to recognize the unprefixed spelling of a first-party import. + """ + root = Path(__file__).parent / "src" / "executorch" + if not root.is_dir(): + return frozenset() + return frozenset(entry.name for entry in root.iterdir() if entry.is_dir()) + + +# Named only by a workflow or by a documented command, so no import reaches them. Listed here +# rather than scanned from .github, which a source distribution does not carry; a test re-derives +# the list so it cannot drift. +_CI_ENTRY_POINTS = ( + "executorch.backends.mlx.test.run_all_tests", + "executorch.backends.mlx.test.test_sample", + "executorch.backends.mlx.test.test_slot_recycling", + "executorch.backends.samsung.test.utils.run_tests", + "executorch.backends.test.suite.generate_markdown_summary_json", + "executorch.examples.models.muse_glimmer.tests.gen_prompt_golden", + "executorch.examples.models.muse_glimmer.tests.test_mlx_pipeline", + "executorch.examples.models.muse_glimmer.tests.test_prompt_tokens", + "executorch.extension.pybindings.test.test_pybindings", +) + +# Directories whose test modules are reached without any import statement naming them, so no scan +# of the source can find them: mlx.yml runs each file it discovers under custom_kernel_ops, the +# webgpu scripts import one module per operator, runner.py resolves a suite root out of a dict and +# then walks it, and the llava README documents a `python -m` command. Directories rather than file +# names, so a new test is covered when it is added. +_CI_ENTRY_POINT_DIRS = ( + "executorch.backends.mlx.custom_kernel_ops", + "executorch.backends.webgpu.test", + "executorch.backends.test.suite", + "executorch.examples.models.llava.test", +) + + +def _is_test_module(dotted: str) -> bool: + return any(part in _TEST_DIR_NAMES for part in dotted.split(".")) + + +def _module_name(root: Path, path: Path) -> str: + parts = list(path.relative_to(root).parts) + if parts[-1] == "__init__.py": + parts = parts[:-1] + else: + parts[-1] = parts[-1][: -len(".py")] + return ".".join(["executorch"] + parts) + + +def _first_party_module(name: str) -> Optional[str]: + """The `executorch.`-prefixed spelling of an import target, or None if it is not ours. + + This repository imports itself two ways. Most code says `executorch.backends.x`, but some + says `backends.x`, which resolves because pytest puts the repository root on sys.path. Both + name the same file, so both have to count as a reference or a helper reached only by the + second spelling is dropped from the wheel while its importers still expect it. + """ + if name.startswith("executorch."): + return name + if name.split(".", 1)[0] in _top_level_package_dirs(): + return f"executorch.{name}" + return None + + +def _scan_imports(path: Path, package: str, out: Set[str], dynamic: Set[str]) -> None: + """Collect into out the executorch modules one file refers to.""" + try: + tree = ast.parse(path.read_text(encoding="utf-8", errors="replace")) + except SyntaxError: + return + for node in ast.walk(tree): + out.update(_import_targets(node, package, dynamic)) + + +_GENERATED_DIR_NAMES = frozenset( + { + ".venv", + "venv", + "build", + "dist", + "buck-out", + ".cache", + ".hypothesis", + ".mypy_cache", + ".pytest_cache", + ".tox", + "test-build", + "arm_test", + "riscv_test", + } +) + + +def _unshipped_directories(root: Path) -> List[Path]: + """Checkout directories the wheel does not carry, whose imports still have to be followed. + + src/executorch is a subset of the repository, so a file under test/ or tools/ is never + packaged, yet a module it imports still has to ship. + + Generated directories are left out, because a build tree or an in-tree virtualenv holds an + INSTALLED copy of this package, and reading it would let the last wheel vote on what the next + one ships. Listed by name rather than asked of git, because `git check-ignore` needs a working + repository and answers differently for a pattern with a trailing slash depending on whether the + directory exists yet, which made the same build behave differently on two platforms. + """ + repository = Path(__file__).parent + if not repository.is_dir() or not root.is_dir(): + return [] + shipped = {entry.name for entry in root.iterdir()} + return [ + entry + for entry in sorted(repository.iterdir()) + if entry.is_dir() + and entry.name not in shipped + and entry.name not in _WALK_SKIP_DIRS + and entry.name not in _GENERATED_DIR_NAMES + and entry.name not in ("src", ".github") + ] + + +def _import_graph(root: Path) -> Tuple[Set[str], Dict[str, Set[str]], Set[str]]: + """Every module under root, what each imports, and literal importlib targets. + + The walk covers root, but the SEED covers more: a file elsewhere in the checkout can import + a module that ships, so its imports are collected too and attributed to a synthetic name. + Without that, a helper whose only importer lives outside the shipped tree looks unreachable. + """ + modules: Set[str] = set() + edges: Dict[str, Set[str]] = {} + dynamic: Set[str] = set() + + # followlinks, because src/executorch is a tree of symlinks into the repository root. + for dirpath, dirnames, filenames in os.walk(root, followlinks=True): + # Vendored trees are skipped by the same test that excludes them from the package list, + # not only by directory name. A submodule checked out under an ordinary name, FACTO and + # the tokenizers among them, is otherwise read as first-party, and its imports would keep + # test modules the wheel has no reason to carry. + dirnames[:] = [ + d + for d in dirnames + if d not in _WALK_SKIP_DIRS + and not _is_vendored_path(os.path.relpath(os.path.join(dirpath, d), root)) + ] + for filename in filenames: + if not filename.endswith(".py"): + continue + path = Path(dirpath) / filename + me = _module_name(root, path) + modules.add(me) + package = me if filename == "__init__.py" else me.rsplit(".", 1)[0] + _scan_imports(path, package, edges.setdefault(me, set()), dynamic) + + # Directories of the checkout that the wheel does not ship, test/ among them. Their files are + # never packaged, so they are not modules, but what they import must still ship: for example + # test/end2end/test_end2end.py imports two model helpers out of exir/tests. + for entry in _unshipped_directories(root): + for dirpath, dirnames, filenames in os.walk(entry, followlinks=False): + dirnames[:] = [d for d in dirnames if d not in _WALK_SKIP_DIRS] + for filename in filenames: + if not filename.endswith(".py"): + continue + outside = f"{dirpath}/{filename}" + _scan_imports( + Path(dirpath) / filename, + "", + edges.setdefault(outside, set()), + dynamic, + ) + + return modules, edges, dynamic + + +def _import_targets(node: ast.AST, package: str, dynamic: Set[str]) -> Set[str]: + """The executorch modules one AST node refers to.""" + found: Set[str] = set() + if isinstance(node, ast.Import): + found.update( + name for name in (_first_party_module(a.name) for a in node.names) if name + ) + elif isinstance(node, ast.ImportFrom): + if node.level: + if not package: + # A file outside the shipped tree, so a relative import stays inside that tree + # and cannot name anything the wheel carries. + return found + # A relative import names a real module too, and inside a kept package its target + # has to ship: stages/__init__.py does `from .export import Export`, so dropping + # stages.export would break every importer of that package. + parts = package.split(".") + if node.level > 1: + parts = parts[: len(parts) - (node.level - 1)] + base = ".".join(parts + (node.module.split(".") if node.module else [])) + elif node.module and (prefixed := _first_party_module(node.module)): + base = prefixed + else: + return found + if base.startswith("executorch"): + found.add(base) + # `from pkg import name` may name a submodule rather than an attribute, and there + # is no way to tell without importing, so both readings are kept. + found.update(f"{base}.{a.name}" for a in node.names) + elif ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "import_module" + and node.args + and isinstance(node.args[0], ast.Constant) + and isinstance(node.args[0].value, str) + ): + target = node.args[0].value + if target.startswith("."): + if not package: + return found + target = package + target + resolved = _first_party_module(target) + if resolved: + dynamic.add(resolved) + return found + + +@functools.lru_cache(maxsize=None) +def _reachable_test_modules() -> FrozenSet[str]: + """Test modules something can still reach once the wheel is installed. + + A test case that nothing imports is dead weight in the wheel: pytest loads it from a path in + the checkout, never through the installed name. A shared helper is the opposite, because the + suites import each other by installed name, so it has to ship or collection breaks. + + Reachable means named by something, anywhere in the checkout, including by a test module + itself. That looks circular and is not: a test collected from the checkout still resolves + `from executorch.x.test import helper` through the INSTALLED package, so the helper must be + in the wheel even though the file importing it is not. + """ + root = Path(__file__).parent / "src" / "executorch" + modules, edges, dynamic = _import_graph(root) + + # Every name anything refers to. No transitive walk is needed: this is already the union of + # every edge target, so following an edge could only rediscover a name that is in here. + referenced = set(dynamic) + referenced.update(_CI_ENTRY_POINTS) + for targets in edges.values(): + referenced.update(targets) + + keep = {name for name in referenced if _is_test_module(name)} & modules + # Everything under a directory whose tests are run one file at a time by a discovery loop. + keep |= { + name + for name in modules + if _is_test_module(name) + and any( + name == prefix or name.startswith(f"{prefix}.") + for prefix in _CI_ENTRY_POINT_DIRS + ) + } + # Parent packages of anything kept, or the dotted path cannot resolve. + for name in list(keep): + parts = name.split(".") + for end in range(2, len(parts)): + parent = ".".join(parts[:end]) + if _is_test_module(parent): + keep.add(parent) + return frozenset(keep) + + +_SHADER_TEMPLATE_MARKERS = ( + "parameter_names_with_default_values", + "shader_variants", + "generate_variant_forall", +) + + +@functools.lru_cache(maxsize=None) +def _is_shader_template(path: str) -> bool: + """Whether a yaml file is a shader codegen input rather than data the wheel needs. + + gen_vulkan_spv.py and gen_wgsl_headers.py expand these into SPIR-V and WGSL headers during + the cmake build, so the wheel already carries the compiled result. Matched on content rather + than on a directory list, because the same shape appears under vulkan and webgpu and a path + list goes stale as soon as a backend adds one. The op and kernel definitions that ARE read + at run time, edge.yaml among them, carry none of these keys. + """ + if not path.endswith(".yaml"): + return False + full = Path(__file__).parent / path + try: + head = full.read_text(encoding="utf-8", errors="replace") + except OSError: + return False + return any(marker in head for marker in _SHADER_TEMPLATE_MARKERS) + + +def _full_packages() -> List[str]: + """Every package the full wheel ships. + + Without an explicit list setuptools discovers all of src/executorch, which pulls in the + Python files and codegen scripts of the vendored third-party checkouts. Those exist to + build the C++ targets, so once the libraries are built no shipped module imports them. + + Test packages deliberately stay. The suites in this repository import each other through + the installed name, for example `from executorch.backends.arm.test import common`, so + dropping them from the wheel stops the suites collecting under a non-editable install. + """ + return sorted( + package + # Anchored on this file rather than the working directory, so the list does not + # change with where the build or a test was started from. + for package in find_namespace_packages( + where=str(Path(__file__).parent / "src"), + include=["executorch", "executorch.*"], + ) + # The include patterns above DO match these, since they are ordinary dotted names, + # which is exactly why they have to be removed here instead. + if not _is_vendored_path(package.replace(".", "/")) + ) + + # The published project names for the CUDA runtime components a CUDA wheel links but # does not bundle, keyed by CUDA major version. Not derivable from a suffix rule: the # CUDA 12 wheels carry a "-cu12" suffix while the CUDA 13 ones are published under @@ -238,6 +675,9 @@ def _minimal_packages() -> List[str]: "13": ("nvidia/cu13/lib",), } +# Arch subdirectory names MKL's exported link interface appends to its prefix. +_MKL_ARCH_DIRECTORIES = ("intel64", "intel64_win", "win-x64") + def _cmake_args() -> List[str]: """CMAKE_ARGS split into arguments, tolerating an unbalanced quote. @@ -600,6 +1040,36 @@ def cuda_named(part: str) -> bool: return len(parts) >= 4 and parts[-3] == "targets" and cuda_named(parts[-4]) +def _is_unresolved_math_library_directory(entry: str) -> bool: + """Whether a runtime search path entry is a maths library directory whose prefix resolved empty. + + Torch's exported CMake package creates a caffe2::mkl imported target, and linking torch brings it + in even though this project never asks for MKL. Its link directories are a hardcoded list of four, + spelled below MKL_ROOT, which resolves to nothing here, so what the linker records is left + anchored at the filesystem root: /lib, /lib/intel64, /lib/intel64_win, /lib/win-x64. The bare + /lib does not survive, because CMake filters its own implicit link directories out of the link + line. + + An empty prefix is the whole signature, so the entry must be exactly /lib/. A real + installation spells the same arch directory below a prefix, as /opt/intel/mkl/lib/intel64, and + that one is a directory the environment genuinely provides. + + Only the three arch directories are handled. The bare prefix/lib is deliberately left, because + with a resolving prefix it is an ordinary library directory this project has no business + dropping, and with an empty one CMake filters it out as an implicit link directory before the + link line is built, so it never reaches a shipped library. + + Dropping these loses nothing: no shipped library names an MKL or OpenMP runtime in DT_NEEDED, so + nothing resolves through them, and they sit ahead of the relative hops appended below. + """ + parts = PurePosixPath(entry).parts + return ( + len(parts) == 3 + and parts[:2] == ("/", "lib") + and parts[2] in _MKL_ARCH_DIRECTORIES + ) + + def _package_relative_depth(library: Path) -> int: """How many directories separate a shipped library from the installed package root. @@ -636,6 +1106,11 @@ def _torchao_requirement() -> str: spec.loader.exec_module(module) version = module.TORCHAO_NIGHTLY_VERSION + if ( + install_utils.determine_torch_url(module.TORCH_URL_BASE).endswith("/cu134") + and not module.torchao_from_source() + ): + version = module.CU134_TORCHAO_NIGHTLY_VERSION major, minor = (int(part) for part in version.split(".")[:2]) return f"torchao>={version},<{major}.{minor + 1}" @@ -1446,6 +1921,9 @@ def _is_usable_runtime_path( # directory from the machine that built it. It is kept when no relative route exists, because # then it is the only way this library finds torch. # + # A maths library directory whose prefix resolved empty is dropped for the same reason: nothing in + # the wheel resolves through it, and it sits ahead of the relative hops appended below. + # # Anything else absolute stays, because it is a dependency the environment provides and the wheel # has no relative answer for. # @@ -1465,6 +1943,8 @@ def _is_usable_runtime_path( # CUDA version, which is the one absolute path that has to survive. if safe_to_drop_toolkit_paths and _is_cuda_toolkit_directory(entry): return False + if _is_unresolved_math_library_directory(entry): + return False if entry.rstrip("/").endswith("/torch/lib") and has_relative_torch_route: return False return True @@ -1590,6 +2070,68 @@ class CustomBuildPy(build_py): a file to a different relative location under the output package directory. """ + def _prune_unstaged_files(self) -> None: + """Delete .py and .yaml this command staged in an earlier build and no longer wants. + + Restricted to files that exist in the source tree, because build_py is the first of + build's sub commands and everything a later one stages is still sitting in the build + directory when this runs. build_ext generates executorch/data/bin/__init__.py, the + target of the flatc console script, from a template that lives elsewhere, so a walk + that removed anything absent from build_py's own file list would delete it. + """ + if self.editable_mode: + return + if not self.packages: + # Nothing to compare against, so every staged file would look unwanted. Refuse + # rather than empty the build directory. + return + + wanted = set() + for package, module, _ in self.find_all_modules(): + parts = package.split(".") if package else [] + wanted.add(os.path.join(self.build_lib, *parts, f"{module}.py")) + for package in self.packages or (): + src_dir = self.get_package_dir(package) + build_dir = os.path.join(*([self.build_lib] + package.split("."))) + for filename in self.find_data_files(package, src_dir): + wanted.add(os.path.join(build_dir, os.path.relpath(filename, src_dir))) + + source_root = Path(__file__).parent / "src" + for dirpath, _, filenames in os.walk(self.build_lib): + for filename in filenames: + if not filename.endswith((".py", ".yaml")): + continue + staged = os.path.join(dirpath, filename) + if staged in wanted: + continue + relative = os.path.relpath(staged, self.build_lib) + if not (source_root / relative).is_file(): + # Generated by another command, so build_py must not remove it. + continue + os.remove(staged) + + def find_package_modules(self, package, package_dir): + modules = super().find_package_modules(package, package_dir) + if self.editable_mode or not _is_test_module(package): + # An editable install exposes the whole source tree whatever is listed here, and a + # package outside a test directory has nothing to drop. + return modules + keep = _reachable_test_modules() + return [ + entry + for entry in modules + if entry[1] == "__init__" or f"{package}.{entry[1]}" in keep + ] + + def find_data_files(self, package, src_dir): + files = super().find_data_files(package, src_dir) + if self.editable_mode: + return files + root = os.path.dirname(os.path.abspath(__file__)) + return [ + _f for _f in files if not _is_shader_template(os.path.relpath(_f, root)) + ] + def analyze_manifest(self): super().analyze_manifest() # Recent versions of setuptools may include bare directory symlinks from version @@ -1604,6 +2146,13 @@ def analyze_manifest(self): _f for _f in self.manifest_files[_pkg] if os.path.isfile(os.path.join(_root, _f)) + # A directory left out of `packages` is not simply skipped. setuptools + # walks up to the nearest listed package and records the file as that + # package's data, so a vendored *.yaml still arrives under its parent. + # Filter with the same list so the two agree. + and not _is_vendored_path(_f) + # Shader templates are consumed by the cmake build, not at run time. + and not _is_shader_template(_f) ] def _copy_extra_files(self, src_to_dst, dst_root: str) -> None: @@ -1646,6 +2195,12 @@ def run(self): # defined by the py_module list and package_data patterns. build_py.run(self) + # A rebuild over a staging directory left by an earlier build keeps whatever that build + # put there, because build_py only ever copies and never deletes. So a file this build + # deliberately leaves out is still present from last time, and the wheel packages it. + # Remove what is no longer wanted rather than only skipping the copy. + self._prune_unstaged_files() + # dst_root is the root of the `executorch` module in the output package # directory. build_lib is the platform-independent root of the output # package, and will look like `pip-out/lib`. It can contain multiple @@ -2287,6 +2842,7 @@ def iter_distribution_names(self): setup_kwargs["packages"] = _minimal_packages() setup_kwargs["install_requires"] = _minimal_dependencies() else: + setup_kwargs["packages"] = _full_packages() # A CUDA wheel links the CUDA runtime but does not bundle it, so the wheels that # carry it are declared here. A CPU wheel adds nothing. setup_kwargs["install_requires"] = _base_dependencies() + _cuda_dependencies() diff --git a/shim_et/BUCK b/shim_et/BUCK index a1a9bdaf65d..2020ed0a73e 100644 --- a/shim_et/BUCK +++ b/shim_et/BUCK @@ -5,6 +5,8 @@ load("@prelude//toolchains:genrule.bzl", "system_genrule_toolchain") load("@prelude//toolchains:python.bzl", "system_python_bootstrap_toolchain", "system_python_toolchain") load("@prelude//toolchains:remote_test_execution.bzl", "remote_test_execution_toolchain") +oncall("executorch") + # Although the non-Android toolchains below are present in shim/BUCK, it appears that we # have to duplicate them here or builds won't work. system_cxx_toolchain( diff --git a/shim_et/third-party/nlohmann-json/BUCK b/shim_et/third-party/nlohmann-json/BUCK index c0b4f27eb52..0a42614a385 100644 --- a/shim_et/third-party/nlohmann-json/BUCK +++ b/shim_et/third-party/nlohmann-json/BUCK @@ -1,3 +1,5 @@ load(":targets.bzl", "define_common_targets") +oncall("executorch") + define_common_targets() diff --git a/shim_et/third-party/re2/BUCK b/shim_et/third-party/re2/BUCK index c0b4f27eb52..0a42614a385 100644 --- a/shim_et/third-party/re2/BUCK +++ b/shim_et/third-party/re2/BUCK @@ -1,3 +1,5 @@ load(":targets.bzl", "define_common_targets") +oncall("executorch") + define_common_targets() diff --git a/shim_et/xplat/executorch/build/build_variables.bzl b/shim_et/xplat/executorch/build/build_variables.bzl index c75436d3dc7..8dc8944b052 100644 --- a/shim_et/xplat/executorch/build/build_variables.bzl +++ b/shim_et/xplat/executorch/build/build_variables.bzl @@ -278,6 +278,7 @@ OPTIMIZED_KERNELS_SRCS = [ "kernels/optimized/cpu/op_native_layer_norm.cpp", "kernels/optimized/cpu/op_sub.cpp", "kernels/optimized/cpu/op_sum.cpp", + "kernels/optimized/cpu/op_to_copy.cpp", "kernels/optimized/cpu/op_where.cpp", ] @@ -320,6 +321,7 @@ OPTIMIZED_NATIVE_CPU_OPS_SRCS = [ "kernels/optimized/cpu/op_mul.cpp", "kernels/optimized/cpu/op_native_layer_norm.cpp", "kernels/optimized/cpu/op_sub.cpp", + "kernels/optimized/cpu/op_to_copy.cpp", "kernels/optimized/cpu/op_where.cpp", ] diff --git a/shim_et/xplat/executorch/build/runtime_wrapper.bzl b/shim_et/xplat/executorch/build/runtime_wrapper.bzl index e84af76f3d9..5d525b73c00 100644 --- a/shim_et/xplat/executorch/build/runtime_wrapper.bzl +++ b/shim_et/xplat/executorch/build/runtime_wrapper.bzl @@ -145,34 +145,72 @@ def _has_pytorch_dep(dep_list): return True return False +def _is_aten_target(kwargs): + """Whether a target compiles against ATen. + + Keyed on exact dep names, not a substring: every label contains "torch". + """ + aten_external_deps = [ + "c10", + "gmock_aten", + "gtest_aten", + "libtorch", + "libtorch_python", + "torch-core-cpp", + ] + aten_resolved_external_deps = [ + "c10", + "libtorch", + "libtorch_python", + "torch-core-cpp", + ] + # The ATen-flavored gtest and gmock names resolve to the same internal + # labels as their ordinary variants, so only their short names are unique. + for key in ["external_deps", "exported_external_deps"]: + for dep in kwargs.get(key) or []: + if dep in aten_external_deps: + return True + + # A target can also name one of those through external_dep_location, which + # hands back the resolved label and puts it in an ordinary dep list. + aten_targets = [] + + def _note_aten_targets(targets): + for target in targets: + if target not in aten_targets: + aten_targets.append(target) + return targets + + for name in aten_resolved_external_deps: + resolved = env.resolve_external_dep(name) + if resolved != env.EXTERNAL_DEP_FALLTHROUGH: + selects.apply(obj = resolved, function = _note_aten_targets) + + # A dep list can be a select(), so collect through selects.apply rather than + # walking it. The lists it holds are the same shape either way. + found = [] + + def _note_aten_deps(targets): + for dep in targets: + if dep in aten_targets: + found.append(dep) + return targets + + for key in ["deps", "exported_deps"]: + if kwargs.get(key): + selects.apply(obj = kwargs.get(key), function = _note_aten_deps) + if found: + return True + + for key in ["xplat_deps", "fbcode_deps"]: + if _has_pytorch_dep(kwargs.get(key)): + return True + return False + def _patch_test_compiler_flags(kwargs): if "compiler_flags" not in kwargs: kwargs["compiler_flags"] = [] - # Determine C++ standard based on whether this is an aten test. - # Aten tests require at least C++20 to compile against PyTorch, while - # non-aten tests are pinned to C++17 for embedded. - name = kwargs.get("name", "") - external_deps = kwargs.get("external_deps", []) - deps = kwargs.get("deps", []) - xplat_deps = kwargs.get("xplat_deps", []) - fbcode_deps = kwargs.get("fbcode_deps", []) - is_aten_test = ( - "_aten" in name or - "aten_" in name or - "libtorch" in external_deps or - "gtest_aten" in external_deps or - "gmock_aten" in external_deps or - _has_pytorch_dep(deps) or - _has_pytorch_dep(xplat_deps) or - _has_pytorch_dep(fbcode_deps) - ) - - if not is_aten_test: - kwargs["compiler_flags"] += [ - "-std=c++17", - ] - # Relaxing some constraints for tests kwargs["compiler_flags"] += [ "-Wno-missing-prototypes", @@ -267,9 +305,21 @@ def _patch_kwargs_cxx(kwargs): env.remove_platform_specific_args(kwargs) return _patch_kwargs_common(kwargs) +def _patch_aten_mode_std(kwargs, aten_mode): + """Raises an ATen-mode target to C++20, which PyTorch's headers require. + + A plain compiler flag, which the prelude places after the toolchain's. + """ + if aten_mode: + kwargs["compiler_flags"] = kwargs.get("compiler_flags", []) + ["-std=c++20"] + return kwargs + def _cxx_library_common(*args, **kwargs): + # Before _patch_kwargs_cxx, which consumes external_deps. + aten_mode = _is_aten_target(kwargs) _patch_kwargs_cxx(kwargs) _patch_build_mode_flags(kwargs) + _patch_aten_mode_std(kwargs, aten_mode) env.patch_platform_build_mode_flags(kwargs) env.patch_headers(kwargs) @@ -294,8 +344,11 @@ def _cxx_library(*args, **kwargs): _cxx_library_common(*args, **kwargs) def _cxx_binary_helper(*args, **kwargs): + # Before _patch_kwargs_cxx, which consumes external_deps. + aten_mode = _is_aten_target(kwargs) _patch_kwargs_cxx(kwargs) _patch_build_mode_flags(kwargs) + _patch_aten_mode_std(kwargs, aten_mode) env.patch_platform_build_mode_flags(kwargs) env.patch_cxx_compiler_flags(kwargs) @@ -324,11 +377,14 @@ def _cxx_test(*args, **kwargs): env.cxx_test(*args, **kwargs) def _cxx_python_extension(*args, **kwargs): + # Before _patch_kwargs_common, which consumes external_deps. + aten_mode = _is_aten_target(kwargs) _patch_kwargs_common(kwargs) _remove_caffe2_deps(kwargs) kwargs["srcs"] = _patch_executorch_references(kwargs["srcs"]) if "types" in kwargs: kwargs["types"] = _patch_executorch_references(kwargs["types"]) + _patch_aten_mode_std(kwargs, aten_mode) env.cxx_python_extension(*args, **kwargs) def _export_file(*args, **kwargs): diff --git a/shim_et/xplat/executorch/codegen/codegen.bzl b/shim_et/xplat/executorch/codegen/codegen.bzl index 318996784a1..9eba82e09e0 100644 --- a/shim_et/xplat/executorch/codegen/codegen.bzl +++ b/shim_et/xplat/executorch/codegen/codegen.bzl @@ -633,13 +633,10 @@ def build_portable_lib( # Currently fbcode links all dependent libraries through shared # library, and it blocks users like unit tests to use kernel # implementation directly. So we enable this for xplat only. - # -Wno-missing-prototypes is Clang-only for C++; GCC (used by Zephyr ARM - # cross-compilation) rejects it with -Werror, so exclude it for Zephyr. - # OSS bypasses the select since ovr_config//os:zephyr is not in the OSS - # buck2 prelude. + # GCC's C++ frontend rejects this C-only flag under -Werror. compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - "ovr_config//os:zephyr": [], + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"] if not expose_operator_symbols and is_xplat(): # Removing '-fvisibility=hidden' exposes operator symbols. @@ -686,13 +683,11 @@ def build_optimized_lib(name, oplist_header_name, portable_header_lib, feature = # Currently fbcode links all dependent libraries through shared # library, and it blocks users like unit tests to use kernel # implementation directly. So we enable this for xplat only. - # -Wno-missing-prototypes and -Wno-global-constructors are Clang-only for - # C++; GCC (used by Zephyr ARM cross-compilation) rejects them with - # -Werror, so exclude them for Zephyr. OSS bypasses the select since - # ovr_config//os:zephyr is not in the OSS buck2 prelude. + # Drop the Clang-only flags for GCC: its C++ frontend rejects + # -Wno-missing-prototypes under -Werror. compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes", "-Wno-pass-failed", "-Wno-global-constructors", "-Wno-shadow"], - "ovr_config//os:zephyr": ["-Wno-pass-failed", "-Wno-shadow"], + "ovr_config//compiler:gcc": ["-Wno-pass-failed", "-Wno-shadow"], }) if not runtime.is_oss else ["-Wno-missing-prototypes", "-Wno-pass-failed", "-Wno-global-constructors", "-Wno-shadow"] if not expose_operator_symbols and is_xplat(): # Removing '-fvisibility=hidden' exposes operator symbols. diff --git a/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl index 7e32f5b7473..25041897059 100644 --- a/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl @@ -96,19 +96,19 @@ def define_op_library(name, compiler_flags, deps): compiler_flags = (select({ # kernels often have helpers with no prototypes just disabling the warning here as the headers # are codegend and linked in later - # -Wno-missing-prototypes is Clang-only for C++; GCC (used by - # Zephyr ARM cross-compilation) rejects it with -Werror, so - # exclude it for Zephyr. OSS bypasses the select since - # ovr_config//os:zephyr is not in the OSS buck2 prelude. - "DEFAULT": [ - "-Wno-missing-prototypes", - # pragma unroll fails with -Os, don't need to warn us and - # fail Werror builds; see https://godbolt.org/z/zvf85vTsr - "-Wno-pass-failed", - ], - "ovr_config//os:zephyr": [ - "-Wno-pass-failed", - ], + # GCC's C++ frontend rejects this C-only flag under -Werror. Nested + # under DEFAULT so the windows (OS) and gcc keys can't both match. + "DEFAULT": select({ + "DEFAULT": [ + "-Wno-missing-prototypes", + # pragma unroll fails with -Os, don't need to warn us and + # fail Werror builds; see https://godbolt.org/z/zvf85vTsr + "-Wno-pass-failed", + ], + "ovr_config//compiler:gcc": [ + "-Wno-pass-failed", + ], + }), # The vendored ATen vec headers trip several -Werror warnings on # the Windows (clang) host, so disable warnings-as-errors there. "ovr_config//os:windows": select({ @@ -327,6 +327,14 @@ OPTIMIZED_ATEN_OPS = ( "//executorch/kernels/portable/cpu/util:reduce_util", ], ), + op_target( + name = "op_to_copy", + deps = [ + "//executorch/extension/threadpool:threadpool", + "//executorch/kernels/portable/cpu:op_to_copy", + "//executorch/kernels/portable/cpu/util:copy_ops_util", + ], + ), op_target( name = "op_where", deps = [ diff --git a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl index f1a46616295..c76c1af67ba 100644 --- a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl @@ -119,21 +119,20 @@ def define_op_library(name, deps, android_deps, aten_target, _allow_third_party_ visibility = ["PUBLIC"], # kernels often have helpers with no prototypes just disabling the warning here as the headers # are codegend and linked in later - # -Wno-missing-prototypes is Clang-only for C++; GCC (used by Zephyr - # ARM cross-compilation) rejects it with -Werror, so exclude it for - # Zephyr and Windows builds. OSS bypasses the zephyr branch via - # runtime.is_oss since ovr_config//os:zephyr is not in the OSS - # buck2 prelude. + # GCC's C++ frontend rejects this C-only flag under -Werror. Nested under + # DEFAULT so the windows (OS) and gcc (compiler) keys can't both match. # The vendored ATen vec headers pulled in on the Windows host trip # several -Werror warnings (e.g. -Wundef on __GNUC__), so disable # warnings-as-errors for the Windows (clang) kernel compiles. compiler_flags = (select({ - "DEFAULT": ["-Wno-missing-prototypes"], + "DEFAULT": select({ + "DEFAULT": ["-Wno-missing-prototypes"], + "ovr_config//compiler:gcc": [], + }), "ovr_config//os:windows": select({ "DEFAULT": ["-Wno-error"], "ovr_config//compiler:msvc": [], }), - "ovr_config//os:zephyr": [], }) if not runtime.is_oss else select({ "DEFAULT": ["-Wno-missing-prototypes"], # OSS buck2 has no compiler constraint (ovr_config//compiler:msvc @@ -574,6 +573,13 @@ ATEN_OPS = ( "//executorch/kernels/portable/cpu/pattern:pattern", ], ), + op_target( + name = "op_fft_r2c", + deps = [ + "//executorch/runtime/core/exec_aten/util:scalar_type_util", + "//executorch/runtime/core/exec_aten/util:tensor_util", + ], + ), op_target( name = "op_fill", deps = [ diff --git a/tools/cmake/preset/README.md b/tools/cmake/preset/README.md index 2eaad11ef37..a1dbbbe0f3b 100644 --- a/tools/cmake/preset/README.md +++ b/tools/cmake/preset/README.md @@ -12,7 +12,7 @@ See: https://github.com/pytorch/executorch/discussions/10661. tl;dr instead of t ```bash $ cmake --preset macos -$ cmake --build cmake-out -j100 --target executor_runner +$ cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target executor_runner ``` ## Working with Presets diff --git a/tools/cmake/preset/default.cmake b/tools/cmake/preset/default.cmake index f0b1285209e..89e679aec7c 100644 --- a/tools/cmake/preset/default.cmake +++ b/tools/cmake/preset/default.cmake @@ -128,6 +128,10 @@ define_overridable_option( EXECUTORCH_BUILD_EXTENSION_APPLE "Build the Apple extension" BOOL OFF ) define_overridable_option(EXECUTORCH_BUILD_MLX "Build the MLX backend" BOOL OFF) +define_overridable_option( + EXECUTORCH_MLX_SWIFTPM_RESOURCES + "Load the MLX metallib from the ExecuTorch SwiftPM resource bundle" BOOL OFF +) define_overridable_option( EXECUTORCH_BUILD_NEURON "Build the backends/mediatek directory" BOOL OFF ) diff --git a/torch_pin.py b/torch_pin.py index ca593b1ef05..f46d5b67ec0 100644 --- a/torch_pin.py +++ b/torch_pin.py @@ -1,2 +1,2 @@ -TORCH_VERSION = "2.13.0" +TORCH_VERSION = "2.14.0" # NIGHTLY_VERSION = "dev20260318" Temporarily pinning to stable release candidate. Revert https://github.com/pytorch/executorch/pull/18287 diff --git a/version.txt b/version.txt index bc80560fad6..dc1e644a101 100644 --- a/version.txt +++ b/version.txt @@ -1 +1 @@ -1.5.0 +1.6.0